# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

from datetime import datetime
from typing import Any, Literal
from uuid import UUID

from pydantic import BaseModel, Field


class NumberTriplet(BaseModel):
    """Triplet of baseline, target, and delta values."""

    baseline: float | None = Field(None, description="Baseline value")
    target: float | None = Field(None, description="Target value")
    delta: float | None = Field(None, description="Delta (target - baseline)")


class FunctionalTestsDiff(BaseModel):
    """Functional tests diff section."""

    summary: dict[str, NumberTriplet] = Field(
        ..., description="Summary statistics as triplets (baseline, target, delta)"
    )
    regressions: list[str] | None = Field(None, description="List of test names that regressed (passed -> failed)")
    fixes: list[str] | None = Field(None, description="List of test names that were fixed (failed -> passed)")


class BenchmarkDiff(BaseModel):
    """Individual benchmark diff result."""

    name: str = Field(..., description="Benchmark name")
    metric: str = Field(..., description="Metric name")
    baseline: float | None = Field(None, description="Baseline metric value")
    target: float | None = Field(None, description="Target metric value")
    delta: float | None = Field(None, description="Delta (target - baseline)")
    subset_metrics: dict[str, NumberTriplet] | None = Field(None, description="Subset-specific metric triplets")


class ReportDiffV1(BaseModel):
    """Report diff schema v1.0 for comparing two evaluation runs."""

    schema_version: Literal["1.0"] = Field("1.0", description="Schema version")
    generated_at: datetime = Field(..., description="Diff generation timestamp (RFC 3339)")
    baseline_run_id: UUID = Field(..., description="Baseline run UUID")
    target_run_id: UUID = Field(..., description="Target run UUID")
    functional_tests: FunctionalTestsDiff
    benchmarks: list[BenchmarkDiff] = Field(default_factory=list, description="Benchmark diff results")

    class Config:
        json_encoders = {
            datetime: lambda v: v.isoformat(),
            UUID: lambda v: str(v),
        }

    def to_json_dict(self) -> dict[str, Any]:
        """Convert to JSON-serializable dictionary."""
        data = self.model_dump(by_alias=True, exclude_none=False)
        # Convert UUIDs to strings for JSON serialization
        if isinstance(data.get("baseline_run_id"), UUID):
            data["baseline_run_id"] = str(data["baseline_run_id"])
        if isinstance(data.get("target_run_id"), UUID):
            data["target_run_id"] = str(data["target_run_id"])
        # Convert datetime to ISO string for JSON serialization
        if isinstance(data.get("generated_at"), datetime):
            data["generated_at"] = data["generated_at"].isoformat()
        return data
