"""Eval service."""

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

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

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


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


def _get_auth_token(auth_user: str) -> str:
    # TODO: fetch credentials from Secrets Manager, exchange for Auth0 token via ROPG
    return "stub-token"


def _resolve_auth_token(prompt_config: PromptConfig, request: EvalRequest) -> str:
    token = prompt_config.auth_token or request.auth_token
    if token:
        return token

    auth_user = prompt_config.auth_user or request.auth_user
    if not auth_user:
        raise ValueError("prompt requires either auth_token or auth_user to be set")
    return _get_auth_token(auth_user)


def _run_prompt(prompt: str, mcp_endpoint: str, token: str) -> llm.LLMResult:
    # 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
        )


def _eval_tags(
    mcp_name: str, prompt_id: str, pipeline: PipelineContext
) -> dict[str, str]:
    tags = {
        "source": "ows-ai-evals",
        "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
    return tags


@workflow(name="eval_judge", _automatic_io_annotation=False)
def _run_judge(
    judge: JudgeConfig,
    answer: str,
    mcp_name: str,
    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, prompt_id, pipeline)
        span = LLMObs.export_span()
        LLMObs.annotate(
            input_data=judge.instructions,
            output_data=answer,
            tags=tags,
        )
        LLMObs.submit_evaluation(
            span=span,
            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,
    prompt_id: str,
    pipeline: PipelineContext,
) -> None:
    tags = _eval_tags(mcp_name, prompt_id, pipeline)
    span = LLMObs.export_span()

    LLMObs.submit_evaluation(
        span=span,
        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,
    token: str,
    mcp_name: str,
    prompt_id: str,
    pipeline: PipelineContext,
) -> PromptRunResult:
    result = _run_prompt(
        prompt=prompt_config.prompt, mcp_endpoint=mcp_endpoint, token=token
    )
    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,
            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,
    )

    if config.LLMOBS_ENABLED:
        LLMObs.annotate(
            input_data=prompt_config.prompt,
            output_data=result.answer_text,
            tags=_eval_tags(mcp_name, prompt_id, pipeline),
        )
        _submit_tool_selection_evaluation(
            result.tools_called,
            prompt_config.expected_tools,
            tool_correct,
            mcp_name,
            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] = {}

    for prompt_id, prompt_config in request.prompts.items():
        token = _resolve_auth_token(prompt_config, request)
        runs = prompt_config.runs or request.runs

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

        prompt_results[prompt_id] = _aggregate_prompt_results(run_results, runs)

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