"""LLM client — AWS Bedrock."""

from functools import lru_cache
from typing import Any, NamedTuple, Protocol

import boto3

from ai_eval_runner import config


class LLMResult(NamedTuple):
    tools_called: list[str]
    answer_text: str
    # True only when the model still wanted to call a tool on the final turn —
    # i.e. it was cut off, not that it happened to finish exactly on the last turn.
    max_turns_exceeded: bool = False


class JudgeVerdict(NamedTuple):
    passed: bool
    rationale: str


class ToolExecutor(Protocol):
    def __call__(self, name: str, arguments: dict[str, Any]) -> str: ...


_JUDGE_TOOL_NAME = "submit_verdict"
_JUDGE_TOOL_DESCRIPTION = (
    "Submit your evaluation verdict for whether the answer meets the given criteria."
)
_JUDGE_TOOL_INPUT_SCHEMA = {
    "type": "object",
    "properties": {
        "passed": {
            "type": "boolean",
            "description": "Whether the answer meets the criteria",
        },
        "rationale": {
            "type": "string",
            "description": "Brief explanation for the verdict",
        },
    },
    "required": ["passed", "rationale"],
}


@lru_cache(maxsize=1)
def _bedrock_client() -> Any:
    return boto3.client("bedrock-runtime", region_name=config.BEDROCK_REGION)


def _tool_config(tools: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "tools": [
            {
                "toolSpec": {
                    "name": t["name"],
                    "description": t.get("description", ""),
                    "inputSchema": {"json": t["inputSchema"]},
                }
            }
            for t in tools
        ]
    }


def run_prompt(
    prompt: str,
    tools: list[dict[str, Any]],
    tool_executor: ToolExecutor | None = None,
    system_prompt: str | None = None,
    max_turns: int | None = None,
) -> LLMResult:
    """Run a prompt, letting the model call tools (via tool_executor) and iterate
    until it produces a final answer with no further tool calls — the same loop
    a real MCP client (Claude Code, etc.) runs on behalf of a user.
    """
    max_turns = max_turns if max_turns is not None else config.MCP_MAX_TOOL_TURNS
    system = [{"text": system_prompt}] if system_prompt else []
    tool_config = _tool_config(tools) if tools else None
    messages: list[dict[str, Any]] = [{"role": "user", "content": [{"text": prompt}]}]
    tools_called: list[str] = []
    answer_text = ""
    tool_use: dict[str, Any] | None = None

    for _ in range(max_turns):
        request: dict[str, Any] = {
            "modelId": config.BEDROCK_MODEL,
            "system": system,
            "messages": messages,
        }
        if tool_config:
            request["toolConfig"] = tool_config

        response = _bedrock_client().converse(**request)
        message = response["output"]["message"]
        content = message["content"]
        tool_use = next(
            (block["toolUse"] for block in content if "toolUse" in block), None
        )
        answer_text = "".join(block["text"] for block in content if "text" in block)

        if not tool_use or not tool_executor:
            break

        tools_called.append(tool_use["name"])
        messages.append(message)
        try:
            result_text = tool_executor(tool_use["name"], tool_use.get("input", {}))
            status = "success"
        except Exception as exc:  # noqa: BLE001 — fed back to the model as a tool failure
            result_text = str(exc)
            status = "error"
        messages.append(
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": tool_use["toolUseId"],
                            "content": [{"text": result_text}],
                            "status": status,
                        }
                    }
                ],
            }
        )
    else:
        # Loop ran max_turns times without breaking — the model still wanted to call
        # a tool on the final turn, so it was cut off rather than finishing naturally.
        return LLMResult(
            tools_called=tools_called,
            answer_text=answer_text,
            max_turns_exceeded=bool(tool_use and tool_executor),
        )

    return LLMResult(tools_called=tools_called, answer_text=answer_text)


def _judge_prompt(instructions: str, answer: str) -> str:
    return f"{instructions}\n\nAnswer to evaluate:\n{answer}"


def run_judge(instructions: str, answer: str) -> JudgeVerdict:
    response = _bedrock_client().converse(
        modelId=config.BEDROCK_MODEL,
        messages=[
            {"role": "user", "content": [{"text": _judge_prompt(instructions, answer)}]}
        ],
        toolConfig={
            "tools": [
                {
                    "toolSpec": {
                        "name": _JUDGE_TOOL_NAME,
                        "description": _JUDGE_TOOL_DESCRIPTION,
                        "inputSchema": {"json": _JUDGE_TOOL_INPUT_SCHEMA},
                    }
                }
            ],
            "toolChoice": {"tool": {"name": _JUDGE_TOOL_NAME}},
        },
    )
    content = response["output"]["message"]["content"]
    tool_use = next(block["toolUse"] for block in content if "toolUse" in block)
    verdict_input = tool_use["input"]
    return JudgeVerdict(
        passed=bool(verdict_input["passed"]),
        rationale=str(verdict_input["rationale"]),
    )