"""LangSmith evaluators for the harmonic analysis agent."""

from langchain_anthropic import ChatAnthropic
from langsmith import Client
from langsmith.evaluation import EvaluationResult
from langsmith.schemas import Example, Run
from pydantic import BaseModel

_JUDGE_MODEL = "claude-sonnet-4-6"


class _MentionedSongs(BaseModel):
    titles: list[str]


def _extract_mentioned_songs(response: str) -> list[str]:
    judge = ChatAnthropic(model=_JUDGE_MODEL).with_structured_output(_MentionedSongs)
    result = judge.invoke(
        "Extract all song titles explicitly mentioned in the following text. "
        "Return only the titles as a list — no artist names, no explanation.\n\n"
        + response
    )
    return result.titles

_FORMATS_TOOL = "get_supported_chord_formats"

_EXPECTED_NOTABLE = frozenset({
    "Let It Be",
    "Ho Hey",
    "With or Without You",
    "Stand By Me",
    "Someone Like You",
    "Wonderwall",
    "Can't Help Falling in Love",
    "Riptide",
    "Hallelujah",
    "Piano Man",
    "Hey Soul Sister",
})


def _count_tool_use_in_llm_runs(llm_runs: list[Run], tool_name: str) -> int:
    """Count calls to a named tool across LLM run outputs.

    Handles two serialisation formats:
    - wrap_anthropic (OpenAI-compatible): outputs["tool_calls"][].function.name
    - LangChain ChatAnthropic: outputs["generations"][][]["message"]["tool_calls"][].name
    """
    count = 0
    for r in llm_runs:
        outputs = r.outputs or {}
        # wrap_anthropic / OpenAI-compatible format
        for tc in outputs.get("tool_calls") or []:
            if tc.get("function", {}).get("name") == tool_name:
                count += 1
        # LangChain ChatAnthropic generations format
        for gen_list in outputs.get("generations") or []:
            for gen in gen_list:
                msg = gen.get("message") or {}
                msg_data = msg.get("kwargs", msg)
                for tc in msg_data.get("tool_calls") or []:
                    if tc.get("name") == tool_name:
                        count += 1
    return count


def evaluate_formats_tool_called_once(run: Run, example: Example) -> EvaluationResult:
    """Check that the supported-formats tool was called exactly once.

    :param run: LangSmith run — LLM child runs are fetched via the client.
    :param example: Dataset example (unused).
    """
    client = Client()
    llm_runs = list(client.list_runs(trace_id=run.id, run_type="llm"))
    count = _count_tool_use_in_llm_runs(llm_runs, _FORMATS_TOOL)
    return EvaluationResult(
        key="formats_tool_called_once",
        score=int(count == 1),
        comment=f"{_FORMATS_TOOL} called {count} time(s)",
    )


def _title_in_text(title: str, text: str) -> bool:
    return title.lower() in text.lower()


def evaluate_notable_songs(run: Run, example: Example) -> list[dict]:
    """Precision, recall, and F1 over the expected notable songs.

    Uses an LLM judge to extract song titles from the prose response, then
    matches deterministically against ``_EXPECTED_NOTABLE``. All three metrics
    are derived from a single extraction call.

    :param run: LangSmith run — expects ``outputs["response"]`` as a string.
    :param example: Dataset example (unused).
    """
    response = (run.outputs or {}).get("response", "")
    if not response:
        return [
            {"key": "notable_songs_precision", "score": None, "comment": "no response in run outputs"},
            {"key": "notable_songs_recall", "score": None, "comment": "no response in run outputs"},
            {"key": "notable_songs_f1", "score": None, "comment": "no response in run outputs"},
        ]

    mentioned = _extract_mentioned_songs(response)
    mentioned_lower = {t.lower() for t in mentioned}
    expected_lower = {t.lower() for t in _EXPECTED_NOTABLE}

    tp = len(mentioned_lower & expected_lower)
    precision = tp / len(mentioned_lower) if mentioned_lower else 0.0
    recall = tp / len(expected_lower)
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

    return [
        {"key": "notable_songs_precision", "score": precision, "comment": f"{tp}/{len(mentioned_lower)}"},
        {"key": "notable_songs_recall", "score": recall, "comment": f"{tp}/{len(expected_lower)}"},
        {"key": "notable_songs_f1", "score": f1, "comment": f"p={precision:.2f} r={recall:.2f}"},
    ]


def evaluate_result_count(run: Run, example: Example) -> EvaluationResult:
    """Check the number of returned songs does not exceed the requested limit.

    :param run: LangSmith run — expects ``outputs["neighbours"]`` as a list.
    :param example: Dataset example — expects ``inputs["limit"]`` as an int.
    """
    neighbours = (run.outputs or {}).get("neighbours", [])
    limit = (example.inputs or {}).get("limit")

    if limit is None:
        return EvaluationResult(
            key="result_count_within_limit",
            score=None,
            comment="limit not found in dataset example inputs",
        )

    count = len(neighbours)
    return EvaluationResult(
        key="result_count_within_limit",
        score=int(count <= limit),
        comment=f"{count} results returned, limit was {limit}",
    )