"""LLM grader: score each assertion against the agent's output.

Grading is a second, structured-output model call. Each assertion is graded
independently as pass/fail with quoted evidence, under a strict prompt that gives
no benefit of the doubt. This is why assertions must be atomic — one claim each.
"""

import json
from dataclasses import dataclass
from typing import Any, cast

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import HumanMessage
from pydantic import BaseModel, model_validator

from skill_eval_runner.usage import message_usage

GRADER_PROMPT = """Grade the AI output below against each assertion.

Rules:
- PASS requires concrete evidence quoted or referenced from the output.
- Do not give benefit of the doubt — if not clearly present, it FAILS.
- Evidence must be specific (quote a phrase, cite a section heading, etc.).

## Output
{response}

## Assertions
{assertions}

For each assertion, set passed=true only if concrete evidence is present in the output."""


class AssertionResult(BaseModel):
    """One assertion's verdict with supporting evidence."""

    text: str
    passed: bool
    evidence: str


class GradingResponse(BaseModel):
    """The grader's structured response: one result per assertion."""

    assertion_results: list[AssertionResult]

    @model_validator(mode="before")
    @classmethod
    def _coerce_stringified_list(cls, value: Any) -> Any:
        """Parse ``assertion_results`` if a model returned it as a JSON string."""
        if isinstance(value, dict) and isinstance(value.get("assertion_results"), str):
            value["assertion_results"] = json.loads(value["assertion_results"])
        return value


@dataclass
class GradingResult:
    """A graded suite case: the gradeable result plus the grader call's token usage.

    ``grading`` is the dict written verbatim to ``grading.json`` (``assertion_results``
    + ``summary``). ``input_tokens`` / ``output_tokens`` are the grader model's own
    usage, kept separate from the agent's so cost accounting can price each — the
    grader model may differ from the agent model.
    """

    grading: dict[str, Any]
    input_tokens: int
    output_tokens: int


def grade(response: str, assertions: list[str], model: BaseChatModel) -> GradingResult:
    """Grade ``response`` against ``assertions``; return the result and grader usage."""
    if not assertions:
        return GradingResult(
            grading={
                "assertion_results": [],
                "summary": {"passed": 0, "failed": 0, "total": 0, "pass_rate": 0.0},
            },
            input_tokens=0,
            output_tokens=0,
        )

    numbered = "\n".join(f"{i + 1}. {a}" for i, a in enumerate(assertions))
    prompt = GRADER_PROMPT.format(response=response, assertions=numbered)

    # include_raw=True so the underlying message (and its token usage) survives the
    # structured-output parsing — a bare .invoke() would return only the parsed model.
    grader = model.with_structured_output(
        GradingResponse, method="json_schema", include_raw=True
    )
    raw = cast(dict[str, Any], grader.invoke([HumanMessage(content=prompt)]))
    parsed = raw.get("parsed")
    if parsed is None:
        raise ValueError(
            f"Grader returned no parseable result: {raw.get('parsing_error')}"
        )
    grading = cast(GradingResponse, parsed)
    input_tokens, output_tokens = message_usage(raw.get("raw"))

    passed = sum(1 for r in grading.assertion_results if r.passed)
    total = len(grading.assertion_results)
    return GradingResult(
        grading={
            "assertion_results": [r.model_dump() for r in grading.assertion_results],
            "summary": {
                "passed": passed,
                "failed": total - passed,
                "total": total,
                "pass_rate": round(passed / total, 3) if total else 0.0,
            },
        },
        input_tokens=input_tokens,
        output_tokens=output_tokens,
    )