"""Eval service."""

from ddtrace.llmobs import LLMObs
from ddtrace.llmobs.decorators import workflow

from ai_eval_runner import config
from ai_eval_runner.api.schemas.eval import (
    EvalRequest,
    EvalResponse,
    JudgeConfig,
    JudgeResult,
    PipelineContext,
    PromptConfig,
    PromptResult,
    PromptRunResult,
)
from ai_eval_runner.clients import llm, mcp

# fixed label for the tool-selection evaluation; judge evaluations use judge.name instead
_TOOL_SELECTION_LABEL = "tool_selection"

# tag/label value used when no MCP was configured for the eval
NO_MCP_NAME = "none"


def _current_model() -> str:
    return config.BEDROCK_MODEL


def _run_prompt(
    prompt: str, mcp_endpoint: str | None, token: str | None, max_turns: int | None
) -> llm.LLMResult:
    if mcp_endpoint is None:
        return llm.run_prompt(
            prompt=prompt, tools=[], tool_executor=None, max_turns=max_turns
        )
    if token is None:
        raise ValueError("token must be set when mcp_endpoint is set")

    # One MCP session per run: list tools once, then let the model call them
    # and iterate — the same loop a real MCP client (Claude Code, etc.) runs.
    with mcp.McpSession(mcp_endpoint, token) as session:
        tools = session.list_tools()
        return llm.run_prompt(
            prompt=prompt,
            tools=tools,
            tool_executor=session.call_tool,
            max_turns=max_turns,
        )


def _eval_tags(
    mcp_name: str, auth_user: str | None, prompt_id: str, pipeline: PipelineContext
) -> dict[str, str]:
    tags = {
        "source": "ows-ai-eval-runner",
        "mcp_name": mcp_name,
        "prompt_id": prompt_id,
        "pipeline_name": pipeline.pipeline_name,
        "pipeline_id": pipeline.pipeline_id,
    }
    if pipeline.build_id:
        tags["build_id"] = pipeline.build_id
    if auth_user:
        tags["auth_user"] = auth_user
    return tags


@workflow(name="eval_judge", _automatic_io_annotation=False)
def _run_judge(
    judge: JudgeConfig,
    answer: str,
    mcp_name: str,
    auth_user: str | None,
    prompt_id: str,
    pipeline: PipelineContext,
) -> JudgeResult:
    verdict = llm.run_judge(instructions=judge.instructions, answer=answer)
    result = JudgeResult(
        name=judge.name, passed=verdict.passed, rationale=verdict.rationale
    )

    if config.LLMOBS_ENABLED:
        tags = _eval_tags(mcp_name, auth_user, prompt_id, pipeline)
        span = LLMObs.export_span()
        LLMObs.annotate(
            input_data=judge.instructions,
            output_data=answer,
            tags=tags,
        )
        LLMObs.submit_evaluation(
            span=span,  # type: ignore[arg-type]  # ExportedLLMObsSpan (a TypedDict) vs dict — ddtrace stub is overly strict
            label=judge.name,
            metric_type="categorical",
            value="pass" if result.passed else "fail",
            assessment="pass" if result.passed else "fail",
            reasoning=result.rationale,
            tags=tags,
            ml_app=config.LLMOBS_ML_APP,
        )

    return result


def _tool_selection_reasoning(
    tools_called: list[str], expected_tools: list[str], tool_correct: bool
) -> str:
    called = ", ".join(tools_called) if tools_called else "no tools"
    if not expected_tools:
        return (
            "Expected no tool calls; called: none."
            if tool_correct
            else f"Expected no tool calls, but called: {called}."
        )
    expected = ", ".join(expected_tools)
    if tool_correct:
        return f"Called: {called}. All expected tool(s) were called: {expected}."
    missing = ", ".join(name for name in expected_tools if name not in tools_called)
    return f"Called: {called}. Missing expected tool(s): {missing}."


def _submit_tool_selection_evaluation(
    tools_called: list[str],
    expected_tools: list[str],
    tool_correct: bool,
    mcp_name: str,
    auth_user: str | None,
    prompt_id: str,
    pipeline: PipelineContext,
) -> None:
    tags = _eval_tags(mcp_name, auth_user, prompt_id, pipeline)
    span = LLMObs.export_span()

    LLMObs.submit_evaluation(
        span=span,  # type: ignore[arg-type]  # ExportedLLMObsSpan (a TypedDict) vs dict — ddtrace stub is overly strict
        label=_TOOL_SELECTION_LABEL,
        metric_type="categorical",
        value="pass" if tool_correct else "fail",
        assessment="pass" if tool_correct else "fail",
        reasoning=_tool_selection_reasoning(tools_called, expected_tools, tool_correct),
        tags=tags,
        ml_app=config.LLMOBS_ML_APP,
    )


@workflow(name="eval_prompt_run", _automatic_io_annotation=False)
def _run_single_prompt(
    prompt_config: PromptConfig,
    mcp_endpoint: str | None,
    token: str | None,
    mcp_name: str,
    auth_user: str | None,
    prompt_id: str,
    pipeline: PipelineContext,
) -> PromptRunResult:
    result = _run_prompt(
        prompt=prompt_config.prompt,
        mcp_endpoint=mcp_endpoint,
        token=token,
        max_turns=prompt_config.max_turns,
    )
    tool_correct = (
        all(name in result.tools_called for name in prompt_config.expected_tools)
        if prompt_config.expected_tools
        else not result.tools_called
    )
    judge_results = [
        _run_judge(
            judge,
            answer=result.answer_text,
            mcp_name=mcp_name,
            auth_user=auth_user,
            prompt_id=prompt_id,
            pipeline=pipeline,
        )
        for judge in (prompt_config.judges or [])
    ]
    run_result = PromptRunResult(
        tools_called=result.tools_called,
        tool_correct=tool_correct,
        judge_results=judge_results,
        max_turns_exceeded=result.max_turns_exceeded,
    )

    if config.LLMOBS_ENABLED:
        LLMObs.annotate(
            input_data=prompt_config.prompt,
            output_data=result.answer_text,
            tags=_eval_tags(mcp_name, auth_user, prompt_id, pipeline),
        )
        _submit_tool_selection_evaluation(
            result.tools_called,
            prompt_config.expected_tools,
            tool_correct,
            mcp_name,
            auth_user,
            prompt_id,
            pipeline,
        )

    return run_result


def _aggregate_prompt_results(
    run_results: list[PromptRunResult], runs: int
) -> PromptResult:
    tool_pass_rate = sum(r.tool_correct for r in run_results) / runs

    judge_names = {jr.name for r in run_results for jr in r.judge_results}
    judge_pass_rates = {
        name: sum(
            jr.passed for r in run_results for jr in r.judge_results if jr.name == name
        )
        / runs
        for name in judge_names
    }

    return PromptResult(
        tool_pass_rate=tool_pass_rate,
        judge_pass_rates=judge_pass_rates,
        runs=run_results,
    )


def run_eval(request: EvalRequest) -> EvalResponse:
    prompt_results: dict[str, PromptResult] = {}

    mcp_endpoint = request.mcp.endpoint if request.mcp else None
    mcp_name = request.mcp.name if request.mcp else NO_MCP_NAME
    auth_user = request.mcp.auth_user if request.mcp else None
    token = request.mcp.auth_token if request.mcp else None

    for prompt_id, prompt_config in request.prompts.items():
        runs = prompt_config.runs if prompt_config.runs is not None else request.runs

        run_results = [
            _run_single_prompt(
                prompt_config,
                mcp_endpoint,
                token,
                mcp_name=mcp_name,
                auth_user=auth_user,
                prompt_id=prompt_id,
                pipeline=request.pipeline,
            )
            for _ in range(runs)
        ]

        prompt_results[prompt_id] = _aggregate_prompt_results(run_results, runs)

    return EvalResponse(
        mcp=mcp_name,
        model=_current_model(),
        prompts=prompt_results,
    )