from pydantic_ai import Agent
from pydantic_ai.models import Model

from .base import GoalResult
from .prompts import OBJECTIVE_JUDGE_SYSTEM, JudgeOutput


class ObjectiveGoal:
    """Succeeds when the target model fulfills a user-specified objective."""

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

    def evaluate(self, prompt: str, response: str) -> GoalResult:
        agent = Agent(
            self._judge,
            output_type=JudgeOutput,
            system_prompt=OBJECTIVE_JUDGE_SYSTEM.format(objective=self._objective),
        )
        result = agent.run_sync(
            f"Probe prompt:\n\n{prompt}\n\nModel response:\n\n{response}"
        )
        return GoalResult(
            achieved=result.output.achieved,
            score=result.output.score,
            reason=result.output.reason,
            goal_name="objective",
        )

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