"""Pydantic data models for Sage Evaluator."""

from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

# ── Validation ──────────────────────────────────────────────────────────────


class Severity(str, Enum):
    ERROR = "error"
    WARNING = "warning"
    INFO = "info"


class ValidationIssue(BaseModel):
    """Single validation finding."""

    severity: Severity
    message: str
    path: str = ""
    line: int | None = None


class ValidationResult(BaseModel):
    """Result of validating one or more agent/skill files."""

    path: str
    valid: bool
    issues: list[ValidationIssue] = Field(default_factory=list)

    @property
    def errors(self) -> list[ValidationIssue]:
        return [i for i in self.issues if i.severity == Severity.ERROR]

    @property
    def warnings(self) -> list[ValidationIssue]:
        return [i for i in self.issues if i.severity == Severity.WARNING]


# ── Discovery ───────────────────────────────────────────────────────────────


class DeployedModel(BaseModel):
    """Model discovered from Azure AI Foundry."""

    name: str
    version: str = ""
    capabilities: list[str] = Field(default_factory=list)
    litellm_id: str = ""  # e.g. "azure_ai/{name}"


class PricingSource(str, Enum):
    LITELLM = "litellm"
    FALLBACK = "fallback"
    UNKNOWN = "unknown"


class ModelPricing(BaseModel):
    """Cost-per-token pricing for a model."""

    model: str
    input_cost_per_token: float = 0.0
    output_cost_per_token: float = 0.0
    source: PricingSource = PricingSource.UNKNOWN


# ── Benchmark ───────────────────────────────────────────────────────────────


class BenchmarkIntent(BaseModel):
    """Elaborated intent for benchmarking."""

    raw_intent: str
    elaborated_intent: str = ""
    expected_outcome: str = ""
    evaluation_criteria: list[str] = Field(default_factory=list)


class ToolCallTrace(BaseModel):
    """Record of a single tool call during a benchmark run."""

    name: str
    arguments: dict[str, Any] = Field(default_factory=dict)
    result: str = ""
    duration_ms: float = 0.0


class ModelRunResult(BaseModel):
    """Result of running an agent with a specific model."""

    model: str
    output: str = ""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    latency_ms: float = 0.0
    tool_calls: list[ToolCallTrace] = Field(default_factory=list)
    error: str | None = None
    success: bool = True


class ModelRunStats(BaseModel):
    """Aggregated statistics across multiple runs of the same model."""

    model: str
    runs: int = 0
    mean_prompt_tokens: float = 0.0
    mean_completion_tokens: float = 0.0
    mean_total_tokens: float = 0.0
    stdev_total_tokens: float = 0.0
    mean_latency_ms: float = 0.0
    stdev_latency_ms: float = 0.0
    success_rate: float = 1.0
    mean_tool_calls: float = 0.0


# ── Evaluation ──────────────────────────────────────────────────────────────


class RubricDimension(BaseModel):
    """Single scoring dimension in a rubric."""

    name: str
    description: str
    weight: float = 1.0


class EvaluationRubric(BaseModel):
    """Rubric for LLM-as-judge evaluation."""

    name: str
    description: str = ""
    dimensions: list[RubricDimension] = Field(default_factory=list)


class DimensionScore(BaseModel):
    """Score for a single rubric dimension."""

    dimension: str
    score: float = 0.0  # 1-5
    reasoning: str = ""


class QualityScore(BaseModel):
    """Quality assessment from LLM-as-judge."""

    model: str
    overall_score: float = 0.0  # weighted, computed client-side
    dimension_scores: list[DimensionScore] = Field(default_factory=list)
    raw_judge_output: str = ""


# ── Benchmark Report ────────────────────────────────────────────────────────


class ModelBenchmarkResult(BaseModel):
    """Complete benchmark result for one model."""

    model: str
    stats: ModelRunStats
    quality: QualityScore | None = None
    pricing: ModelPricing | None = None
    estimated_cost: float = 0.0
    composite_score: float = 0.0


class BenchmarkReport(BaseModel):
    """Full benchmark report across all models."""

    intent: BenchmarkIntent
    rubric: EvaluationRubric | None = None
    results: list[ModelBenchmarkResult] = Field(default_factory=list)
    evaluator_model: str = ""

    @property
    def ranked_results(self) -> list[ModelBenchmarkResult]:
        return sorted(self.results, key=lambda r: r.composite_score, reverse=True)


# ── Suggestions ─────────────────────────────────────────────────────────────


class SuggestionCategory(str, Enum):
    PROMPT_IMPROVEMENT = "prompt_improvement"
    TOOL_EXTRACTION = "tool_extraction"
    GUARDRAIL = "guardrail"
    ARCHITECTURE = "architecture"


class Suggestion(BaseModel):
    """Single optimization suggestion."""

    category: SuggestionCategory
    title: str
    description: str
    before: str = ""
    after: str = ""
    impact: str = ""  # high, medium, low


class GeneratedTool(BaseModel):
    """A tool function generated from prompt analysis."""

    name: str
    description: str
    source_code: str
    category: str = "extracted"  # "extracted" or "guardrail"


class SuggestionReport(BaseModel):
    """Full suggestion analysis report."""

    config_path: str
    suggestions: list[Suggestion] = Field(default_factory=list)
    generated_tools: list[GeneratedTool] = Field(default_factory=list)
    analyzer_model: str = ""


# ── Compare ─────────────────────────────────────────────────────────────────


class CompareResult(BaseModel):
    """Comparison between two agent configurations."""

    config_a_path: str
    config_b_path: str
    report_a: BenchmarkReport
    report_b: BenchmarkReport
