from app.llm.client import call_llm
from app.models.schemas import ChatResult, SecurityPipelineStep
from app.security.layer import security_layer


async def handle_secure(
    message: str,
    *,
    scenario_id: str,
    task: str,
) -> ChatResult:
    pre = security_layer.process(message, scenario_id=scenario_id, task=task)
    pipeline = list(pre.get("pipeline", []))

    if pre["blocked"]:
        return ChatResult(
            output=pre["output"],
            attack_success=False,
            consequences=[],
            blocked_by=pre["blocked_by"],
            defense_explanation=pre["defense_explanation"],
            risk_level=pre["risk_level"],  # type: ignore[arg-type]
            audit_entries=pre["audit_entries"],
            security_pipeline=[SecurityPipelineStep(**s) for s in pipeline],
        )

    raw = await call_llm(
        pre["prompt_for_llm"] or message,
        user_input=message,
        task=task,
        mode="secure",
        injection_detected=pre["injection_detected"],
    )

    pipeline.append(
        security_layer.build_llm_step(
            called=True,
            prompt_preview=pre["prompt_for_llm"],
        )
    )

    post = security_layer.post_process_output(
        raw,
        scenario_id=scenario_id,
        user_input=message,
        injection_detected=pre["injection_detected"],
        input_guard_note=pre.get("input_guard_note"),
    )

    pipeline.extend(post.get("pipeline", []))
    audit_entries = pre["audit_entries"] + post["audit_entries"]

    return ChatResult(
        output=post["output"],
        attack_success=post["attack_success"],
        consequences=post["consequences"],
        blocked_by=post["blocked_by"],
        defense_explanation=post["defense_explanation"],
        risk_level=post["risk_level"],  # type: ignore[arg-type]
        audit_entries=audit_entries,
        security_pipeline=[SecurityPipelineStep(**s) for s in pipeline],
    )