from pydantic import ValidationError
from pydantic_ai import Agent
from pydantic_ai.exceptions import UnexpectedModelBehavior
from pydantic_ai.models import Model

from probegpt.data import load_judge_prompt

from .models import GoalResult, JudgeOutput

_MAX_PROMPT_LEN = 2000
_MAX_RESPONSE_LEN = 4000


def _truncate(text: str, max_len: int) -> str:
    text = text.strip()
    if len(text) <= max_len:
        return text
    return text[:max_len] + "\n\n[truncated]"


class ObjectiveJudge:
    """Scores target model responses against a user-defined objective."""

    def __init__(self, objective: str, judge_model: Model) -> None:
        self._objective = objective
        self._judge = judge_model
        self._system_prompt = load_judge_prompt(objective)

    def evaluate(self, prompt: str, response: str) -> GoalResult:
        truncated_prompt = _truncate(prompt, _MAX_PROMPT_LEN)
        truncated_response = _truncate(response, _MAX_RESPONSE_LEN)

        agent = Agent(
            self._judge,
            output_type=JudgeOutput,
            system_prompt=self._system_prompt,
            output_retries=3,
        )
        try:
            result = agent.run_sync(
                f"Probe prompt:\n\n{truncated_prompt}\n\nModel response:\n\n{truncated_response}"
            )
            return GoalResult(
                achieved=result.output.achieved,
                score=result.output.score,
                reason=result.output.reason,
                goal_name="objective",
            )
        except (UnexpectedModelBehavior, ValidationError) as exc:
            return GoalResult(
                achieved=False,
                score=0.0,
                reason=f"[EVAL ERROR: {exc}]",
                goal_name="objective",
            )

    def get_name(self) -> str:
        return "objective"
