"""LLM-powered pairwise evaluation judge."""

from __future__ import annotations

import json
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.cancellation_token import CancellationToken
from agentbyte.llm.base import BaseChatCompletionClient
from agentbyte.messages import SystemMessage, UserMessage

from agentbyte.eval.base import BaseEvalJudge
from agentbyte.eval.types import AnswerStrategy, EvalScore, EvalTrajectory, PairwiseResult


class _PairwiseResponse(BaseModel):
    model_config = ConfigDict(extra="allow")

    winner: str = Field(default="tie")
    margin: float = Field(default=0.0)
    reasoning: str = Field(default="No reasoning provided")


class PairwiseJudge(BaseEvalJudge):
    """Compare two trajectories for the same task and decide which is better.

    Uses an LLM to make a relative judgment: "which response is better?"
    This is more reliable than absolute pointwise scoring for subjective
    qualities like helpfulness and clarity.

    The primary interface is ``compare(trajectory_a, trajectory_b)``.
    The inherited ``score()`` method is not applicable and raises ``TypeError``.

    ``answer_strategy`` defaults to ``"last_assistant"`` because pairwise judges
    compare agent outputs — semantically an ``AssistantMessage`` concern.
    Pass ``source_filter`` to target a specific agent in orchestrator trajectories.
    """

    def __init__(
        self,
        client: BaseChatCompletionClient,
        *,
        name: str | None = None,
        criteria: list[str] | None = None,
        answer_strategy: AnswerStrategy = "last_assistant",
        source_filter: str | None = None,
        custom_instructions: str | None = None,
    ) -> None:
        super().__init__(
            name=name or f"Pairwise-{client.model}",
            answer_strategy=answer_strategy,
            source_filter=source_filter,
        )
        self.client = client
        self.criteria = criteria or ["accuracy", "helpfulness", "clarity"]
        self.custom_instructions = custom_instructions

    async def compare(
        self,
        trajectory_a: EvalTrajectory,
        trajectory_b: EvalTrajectory,
        cancellation_token: CancellationToken | None = None,
    ) -> PairwiseResult:
        """Compare two trajectories and return a pairwise decision.

        Returns a tie with ``margin=0.0`` on cancellation or LLM parse failure.
        The failure reason is preserved in ``metadata["fallback_reason"]``.
        """
        if cancellation_token and cancellation_token.is_cancelled():
            return self._tie(trajectory_a, trajectory_b, fallback_reason="Cancelled")

        answer_a = self.extract_answer(trajectory_a)
        answer_b = self.extract_answer(trajectory_b)

        system_prompt = (
            "You are a pairwise evaluation judge. "
            "Compare Response A and Response B for the same task and decide which is better. "
            "Return strict JSON with keys:\n"
            '  "winner": "a" if A is better, "b" if B is better, "tie" if equivalent\n'
            '  "margin": float 0.0 (barely any difference) to 1.0 (decisive clear winner)\n'
            '  "reasoning": brief explanation of your decision\n'
            "Ties must have margin 0.0."
        )
        if self.custom_instructions:
            system_prompt += f"\n\nAdditional instructions:\n{self.custom_instructions.strip()}"

        user_payload: dict[str, Any] = {
            "task_name": trajectory_a.task.name,
            "task_input": trajectory_a.task.input,
            "expected_output": trajectory_a.task.expected_output,
            "criteria": self.criteria,
            "response_a": answer_a,
            "response_b": answer_b,
        }

        try:
            result = await self.client.create(
                messages=[
                    SystemMessage(content=system_prompt, source="system"),
                    UserMessage(
                        content=json.dumps(user_payload, ensure_ascii=True, indent=2),
                        source="user",
                    ),
                ],
                output_format=_PairwiseResponse,
            )
            parsed = self._parse_response(result.structured_output, result.message.content)
            winner = parsed.winner if parsed.winner in ("a", "b", "tie") else "tie"
            margin = max(0.0, min(1.0, float(parsed.margin)))
            if winner == "tie":
                margin = 0.0
            return PairwiseResult(
                task=trajectory_a.task,
                winner=winner,  # type: ignore[arg-type]
                margin=margin,
                reasoning=parsed.reasoning,
                trajectory_a=trajectory_a,
                trajectory_b=trajectory_b,
                metadata={"judge": self.name, "model": result.model},
            )
        except Exception as exc:
            return self._tie(trajectory_a, trajectory_b, fallback_reason=str(exc))

    async def score(self, *args: Any, **kwargs: Any) -> EvalScore:
        """Not supported on PairwiseJudge — use compare() instead."""
        raise TypeError(
            "PairwiseJudge does not support score(). Use compare(trajectory_a, trajectory_b)."
        )

    def _parse_response(
        self,
        structured_output: Any,
        raw_content: str,
    ) -> _PairwiseResponse:
        if isinstance(structured_output, _PairwiseResponse):
            return structured_output
        if isinstance(structured_output, BaseModel):
            return _PairwiseResponse.model_validate(structured_output.model_dump())
        if structured_output is not None:
            return _PairwiseResponse.model_validate(structured_output)
        return _PairwiseResponse.model_validate_json(raw_content)

    def _tie(
        self,
        trajectory_a: EvalTrajectory,
        trajectory_b: EvalTrajectory,
        *,
        fallback_reason: str,
    ) -> PairwiseResult:
        return PairwiseResult(
            task=trajectory_a.task,
            winner="tie",
            margin=0.0,
            reasoning=fallback_reason,
            trajectory_a=trajectory_a,
            trajectory_b=trajectory_b,
            metadata={"judge": self.name, "fallback": True, "fallback_reason": fallback_reason},
        )


__all__ = ["PairwiseJudge"]
