"""A fake LLM that answers per stage: the guard and the planner get the same user text, so MockProvider's
`contains` rules can't tell them apart."""

from typing import Any

from app.llm.provider import ChatMessage, JsonSchema, LLMError, LLMResult


class ScriptedLLM:
    def __init__(self, guard: dict[str, Any] | str | None = None, plan: dict[str, Any] | str | None = None) -> None:
        self.guard: dict[str, Any] | str | None = guard if guard is not None else {"verdict": "safe", "reason": "ok"}
        self.plan = plan
        self.down = False
        self.calls: list[tuple[str, list[ChatMessage]]] = []

    async def complete(self, messages: list[ChatMessage], schema: JsonSchema | None = None) -> LLMResult:
        stage = schema.name if schema else "text"
        self.calls.append((stage, messages))
        if self.down:
            raise LLMError("connection refused")
        reply = self.guard if stage == "guard_verdict" else self.plan
        if isinstance(reply, dict):
            return LLMResult(text=str(reply), data=reply, prompt_tokens=10, completion_tokens=5)
        return LLMResult(text=reply or "", data=None, prompt_tokens=10, completion_tokens=5)

    async def aclose(self) -> None:
        pass

    def stages(self) -> list[str]:
        return [s for s, _ in self.calls]


def plan(**overrides: Any) -> dict[str, Any]:
    base: dict[str, Any] = {
        "container": "crash",
        "since_minutes": 60,
        "pattern": "connection|refused",
        "ignore_case": True,
        "invert": False,
        "count_only": False,
        "context": 0,
        "limit": 100,
    }
    return {**base, **overrides}