"""LLM-powered 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


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

    overall: float | None = None
    dimensions: dict[str, float] = Field(default_factory=dict)
    reasoning: dict[str, str] = Field(default_factory=dict)


class LLMEvalJudge(BaseEvalJudge):
    """Use an LLM to judge a trajectory on one or more criteria."""

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

    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalScore:
        # resolve: call-level arg > task metadata > judge default
        if criteria is None:
            raw = trajectory.task.metadata.get("_criteria")
            if isinstance(raw, list):
                criteria = raw
            elif isinstance(raw, str):
                criteria = [raw]
        eval_criteria = criteria or self.default_criteria
        if cancellation_token and cancellation_token.is_cancelled():
            return self._neutral_score(
                trajectory,
                eval_criteria,
                reason_prefix="Evaluation cancelled",
            )

        answer = self.extract_answer(trajectory)
        expected_output = trajectory.task.expected_output
        system_prompt = (
            "You are an evaluation judge. "
            "Return strict JSON with keys: overall, dimensions, reasoning. "
            "Each dimension score must be between 0 and 10."
        )
        if self.custom_instructions:
            system_prompt += f"\n\nAdditional instructions:\n{self.custom_instructions.strip()}"

        user_payload = {
            "task_name": trajectory.task.name,
            "task_input": trajectory.task.input,
            "expected_output": expected_output,
            "actual_output": answer,
            "success": trajectory.success,
            "error": trajectory.error,
            "criteria": eval_criteria,
            "usage": trajectory.usage.model_dump(mode="json")
            if trajectory.usage is not None
            else None,
            "metadata": trajectory.metadata,
        }

        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=_JudgeResponse,
            )
            parsed = self._parse_response(result.structured_output, result.message.content)
            dimensions = self._normalize_dimensions(parsed.dimensions, eval_criteria)
            reasoning = {
                criterion: parsed.reasoning.get(criterion, "No reasoning provided")
                for criterion in eval_criteria
            }
            overall = parsed.overall
            if overall is None:
                overall = sum(dimensions.values()) / max(len(dimensions), 1)
            overall = self._clamp_score(overall)
            return EvalScore(
                overall=overall,
                dimensions=dimensions,
                reasoning=reasoning,
                trajectory=trajectory,
                metadata={
                    "judge": self.name,
                    "model": result.model,
                    "finish_reason": result.finish_reason,
                },
            )
        except Exception as exc:
            return self._neutral_score(
                trajectory,
                eval_criteria,
                reason_prefix=f"Judge fallback: {exc}",
            )

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

    def _normalize_dimensions(
        self,
        dimensions: dict[str, float],
        criteria: list[str],
    ) -> dict[str, float]:
        return {
            criterion: self._clamp_score(dimensions.get(criterion, 5.0))
            for criterion in criteria
        }

    def _neutral_score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str],
        *,
        reason_prefix: str,
    ) -> EvalScore:
        return EvalScore(
            overall=5.0,
            dimensions={criterion: 5.0 for criterion in criteria},
            reasoning={criterion: reason_prefix for criterion in criteria},
            trajectory=trajectory,
            metadata={"judge": self.name, "fallback": True},
        )

    @staticmethod
    def _clamp_score(value: float) -> float:
        return min(max(float(value), 0.0), 10.0)


__all__ = ["LLMEvalJudge"]
