"""Data models for graded seeds."""

from __future__ import annotations

from pydantic import BaseModel, Field, field_validator

from probegpt.core.models import Probe


class SeedAnalysis(BaseModel):
    """LLM-generated analysis of why a probe worked and how transferable it is."""

    techniques: list[str] = Field(
        description="Up to 5 short technique tags, e.g. 'academic-framing', 'persona-injection'",
        max_length=5,
    )
    novelty: float = Field(
        ge=0.0,
        le=1.0,
        description="How different this probe's approach is from existing seeds (0=duplicate, 1=entirely new)",
    )
    transferability: float = Field(
        ge=0.0,
        le=1.0,
        description="How likely this technique is to work on variations of the objective (0=very specific, 1=broadly applicable)",
    )
    summary: str = Field(
        max_length=200,
        description="One-sentence explanation of why this probe succeeded",
    )

    @field_validator("novelty", "transferability", mode="before")
    @classmethod
    def coerce_float(cls, v: object) -> float:
        try:
            return max(0.0, min(1.0, float(v)))  # type: ignore[arg-type]
        except (TypeError, ValueError):
            return 0.5

    @field_validator("techniques", mode="before")
    @classmethod
    def coerce_techniques(cls, v: object) -> list[str]:
        if isinstance(v, str):
            return [v]
        if not isinstance(v, list):
            return []
        return [str(t) for t in v]


class GradedSeed(BaseModel):
    """A probe that scored above threshold and has been analyzed by the grader agent."""

    probe: Probe
    score: float
    analysis: SeedAnalysis
    iteration: int

    @property
    def weight(self) -> float:
        """Composite weight used for weighted sampling during mutation."""
        return self.score * self.analysis.novelty * self.analysis.transferability
