"""Agent-based seed grader: analyzes why a probe worked and whether it's novel."""

from __future__ import annotations

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

from probegpt.data import load_grader_config

from .models import GradedSeed, SeedAnalysis

_MIN_PROMPT_LEN = 20
_MAX_PROMPT_LEN = 2000
_MAX_RESPONSE_LEN = 3000
_DUPLICATE_NOVELTY_THRESHOLD = 0.25

# Load vocabulary and system prompt template from grader.yaml at import time
_GRADER_CONFIG = load_grader_config()
TECHNIQUE_VOCABULARY: list[str] = _GRADER_CONFIG["technique_vocabulary"]
TECHNIQUE_VOCABULARY_STR: str = "\n".join(f"  {t}" for t in TECHNIQUE_VOCABULARY)
_SYSTEM_TEMPLATE: str = _GRADER_CONFIG["system_prompt"]


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


def _format_existing(seeds: list[GradedSeed]) -> str:
    if not seeds:
        return "  (none yet)"
    lines = []
    for s in seeds:
        tags = ", ".join(s.analysis.techniques) if s.analysis.techniques else "untagged"
        lines.append(f"  - [{tags}] {s.analysis.summary}")
    return "\n".join(lines)


class SeedGrader:
    """LLM agent that analyzes a successful probe and produces a SeedAnalysis."""

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

    def analyze(
        self,
        prompt: str,
        response: str,
        existing_seeds: list[GradedSeed],
    ) -> SeedAnalysis | None:
        """Analyze a probe and return a SeedAnalysis, or None on failure."""
        if len(prompt.strip()) < _MIN_PROMPT_LEN:
            return None

        system = (
            _SYSTEM_TEMPLATE.replace("{vocabulary}", TECHNIQUE_VOCABULARY_STR)
            .replace("{objective}", self._objective)
            .replace("{existing_summaries}", _format_existing(existing_seeds))
        )
        agent = Agent(
            self._model,
            output_type=SeedAnalysis,
            system_prompt=system,
            output_retries=3,
        )
        user_msg = (
            f"Probe prompt:\n\n{_truncate(prompt, _MAX_PROMPT_LEN)}"
            f"\n\nModel response:\n\n{_truncate(response, _MAX_RESPONSE_LEN)}"
        )
        try:
            result = agent.run_sync(user_msg)
            return result.output
        except UnexpectedModelBehavior:
            return None
        except Exception:
            return None

    def is_duplicate(
        self,
        analysis: SeedAnalysis,
        existing_seeds: list[GradedSeed],
    ) -> bool:
        if not existing_seeds:
            return False
        return analysis.novelty < _DUPLICATE_NOVELTY_THRESHOLD
