# 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, model_validator


class EvalTestCase(BaseModel):
    """Individual test case result."""

    name: str = Field(..., description="Fully-qualified pytest id")
    status: str = Field(
        ...,
        description="Test outcome",
        pattern="^(passed|failed|xfailed|xpassed|skipped)$",
    )
    duration_ms: float | None = Field(None, ge=0, description="Wall-clock runtime in milliseconds")
    error: str | None = Field(None, description="Traceback/message for failures only")


class FunctionalTestsSummary(BaseModel):
    """Summary statistics for functional tests."""

    total: int = Field(..., ge=0, description="Total number of tests")
    passed: int = Field(..., ge=0, description="Number of passed tests")
    failed: int = Field(..., ge=0, description="Number of failed tests")
    xfailed: int = Field(..., ge=0, description="Number of expected failures")
    xpassed: int = Field(..., ge=0, description="Number of unexpected passes")
    skipped: int = Field(..., ge=0, description="Number of skipped tests")
    pass_rate: float | None = Field(None, ge=0, le=100, description="Pass rate percentage")

    @model_validator(mode="before")
    @classmethod
    def calculate_pass_rate(cls, values):
        """Calculate pass rate if not provided."""
        if isinstance(values, dict):
            if values.get("pass_rate") is None:
                total = values.get("total", 0)
                passed = values.get("passed", 0)
                values["pass_rate"] = (passed / total * 100) if total > 0 else 0.0
        return values


class FunctionalTests(BaseModel):
    """Functional tests section of the report."""

    summary: FunctionalTestsSummary
    cases: list[EvalTestCase] = Field(default_factory=list, description="Individual test case results")


class BenchmarkResult(BaseModel):
    """Individual benchmark result."""

    name: str = Field(..., description="Dataset/benchmark identifier")
    topline_metric_name: str = Field(..., description="Primary metric name")
    metrics: dict[str, float] = Field(..., description="Top-level metrics")
    subset_metrics: dict[str, dict[str, float]] | None = Field(None, description="Domain/subset-specific metrics")
    config: dict[str, Any] | None = Field(None, description="Benchmark configuration parameters")


class BenchmarksSummary(BaseModel):
    """Summary statistics for benchmarks."""

    total: int = Field(..., ge=0, description="Number of benchmarks executed")


class Benchmarks(BaseModel):
    """Benchmarks section of the report."""

    summary: BenchmarksSummary
    benchmarks: list[BenchmarkResult] = Field(default_factory=list, description="Individual benchmark results")


class ReportV1(BaseModel):
    """Main report schema v1.0 for Llama Stack Evals."""

    schema_version: Literal["1.0"] = Field("1.0", description="Schema version")
    generated_at: datetime = Field(..., description="Report generation timestamp (RFC 3339)")
    run_id: UUID = Field(..., description="Globally unique run identifier")
    model: str = Field(..., description="Model name/alias")
    model_checkpoint: str = Field(..., description="Model checkpoint identifier")
    provider: str = Field(..., description="Hosting provider or cluster")
    metadata: dict[str, Any] | None = Field(
        None, description="Optional provider-specific or forward-compatibility fields"
    )
    functional_tests: FunctionalTests
    benchmarks: Benchmarks

    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 UUID to string for JSON serialization
        if isinstance(data.get("run_id"), UUID):
            data["run_id"] = str(data["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
