# 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.

import json
import uuid
from datetime import UTC, datetime

import pytest
from pydantic import ValidationError

from llama_verifications.schemas.report_diff_v1 import (
    BenchmarkDiff,
    FunctionalTestsDiff,
    NumberTriplet,
    ReportDiffV1,
)


class TestNumberTriplet:
    def test_valid_triplet(self):
        """Test creating a valid number triplet."""
        triplet = NumberTriplet(baseline=0.75, target=0.80, delta=0.05)
        assert triplet.baseline == 0.75
        assert triplet.target == 0.80
        assert triplet.delta == 0.05

    def test_triplet_with_nulls(self):
        """Test triplet with null values."""
        triplet = NumberTriplet(baseline=None, target=0.80, delta=None)
        assert triplet.baseline is None
        assert triplet.target == 0.80
        assert triplet.delta is None

    def test_all_null_triplet(self):
        """Test triplet with all null values."""
        triplet = NumberTriplet(baseline=None, target=None, delta=None)
        assert triplet.baseline is None
        assert triplet.target is None
        assert triplet.delta is None


class TestFunctionalTestsDiff:
    def test_valid_functional_tests_diff(self):
        """Test creating a valid functional tests diff."""
        summary = {
            "total": NumberTriplet(baseline=100, target=105, delta=5),
            "passed": NumberTriplet(baseline=95, target=98, delta=3),
            "failed": NumberTriplet(baseline=5, target=7, delta=2),
            "xfailed": NumberTriplet(baseline=0, target=0, delta=0),
            "xpassed": NumberTriplet(baseline=0, target=0, delta=0),
            "skipped": NumberTriplet(baseline=0, target=0, delta=0),
            "pass_rate": NumberTriplet(baseline=95.0, target=93.3, delta=-1.7),
        }

        diff = FunctionalTestsDiff(
            summary=summary,
            regressions=["test_example::test_regression"],
            fixes=["test_example::test_fix"],
        )

        assert diff.summary["total"].baseline == 100
        assert diff.summary["total"].target == 105
        assert diff.regressions == ["test_example::test_regression"]
        assert diff.fixes == ["test_example::test_fix"]

    def test_functional_tests_diff_without_changes(self):
        """Test functional tests diff with no regressions or fixes."""
        summary = {
            "total": NumberTriplet(baseline=100, target=100, delta=0),
            "passed": NumberTriplet(baseline=95, target=95, delta=0),
            "failed": NumberTriplet(baseline=5, target=5, delta=0),
            "xfailed": NumberTriplet(baseline=0, target=0, delta=0),
            "xpassed": NumberTriplet(baseline=0, target=0, delta=0),
            "skipped": NumberTriplet(baseline=0, target=0, delta=0),
            "pass_rate": NumberTriplet(baseline=95.0, target=95.0, delta=0.0),
        }

        diff = FunctionalTestsDiff(summary=summary)

        assert diff.regressions is None
        assert diff.fixes is None


class TestBenchmarkDiff:
    def test_valid_benchmark_diff(self):
        """Test creating a valid benchmark diff."""
        subset_metrics = {
            "humanities": NumberTriplet(baseline=0.81, target=0.83, delta=0.02),
            "stem": NumberTriplet(baseline=0.75, target=0.77, delta=0.02),
        }

        diff = BenchmarkDiff(
            name="mmlu",
            metric="accuracy",
            baseline=0.78,
            target=0.80,
            delta=0.02,
            subset_metrics=subset_metrics,
        )

        assert diff.name == "mmlu"
        assert diff.metric == "accuracy"
        assert diff.baseline == 0.78
        assert diff.target == 0.80
        assert diff.delta == 0.02
        assert diff.subset_metrics["humanities"].delta == 0.02

    def test_benchmark_diff_with_nulls(self):
        """Test benchmark diff with null values."""
        diff = BenchmarkDiff(
            name="new_benchmark",
            metric="accuracy",
            baseline=None,
            target=0.85,
            delta=None,
        )

        assert diff.name == "new_benchmark"
        assert diff.baseline is None
        assert diff.target == 0.85
        assert diff.delta is None
        assert diff.subset_metrics is None


class TestReportDiffV1:
    def test_valid_report_diff(self):
        """Test creating a valid report diff."""
        baseline_run_id = uuid.uuid4()
        target_run_id = uuid.uuid4()
        generated_at = datetime.now(UTC)

        functional_tests_diff = FunctionalTestsDiff(
            summary={
                "total": NumberTriplet(baseline=100, target=105, delta=5),
                "passed": NumberTriplet(baseline=95, target=98, delta=3),
                "failed": NumberTriplet(baseline=5, target=7, delta=2),
                "xfailed": NumberTriplet(baseline=0, target=0, delta=0),
                "xpassed": NumberTriplet(baseline=0, target=0, delta=0),
                "skipped": NumberTriplet(baseline=0, target=0, delta=0),
                "pass_rate": NumberTriplet(baseline=95.0, target=93.3, delta=-1.7),
            },
            regressions=["test_regression"],
            fixes=["test_fix"],
        )

        benchmarks_diff = [
            BenchmarkDiff(
                name="mmlu",
                metric="accuracy",
                baseline=0.78,
                target=0.80,
                delta=0.02,
            ),
            BenchmarkDiff(
                name="hellaswag",
                metric="accuracy",
                baseline=0.85,
                target=0.83,
                delta=-0.02,
            ),
        ]

        diff_report = ReportDiffV1(
            generated_at=generated_at,
            baseline_run_id=baseline_run_id,
            target_run_id=target_run_id,
            functional_tests=functional_tests_diff,
            benchmarks=benchmarks_diff,
        )

        assert diff_report.schema_version == "1.0"
        assert diff_report.baseline_run_id == baseline_run_id
        assert diff_report.target_run_id == target_run_id
        assert diff_report.functional_tests.summary["total"].delta == 5
        assert len(diff_report.benchmarks) == 2
        assert diff_report.benchmarks[0].name == "mmlu"
        assert diff_report.benchmarks[1].delta == -0.02

    def test_to_json_dict(self):
        """Test converting diff report to JSON dictionary."""
        baseline_run_id = uuid.uuid4()
        target_run_id = uuid.uuid4()

        diff_report = ReportDiffV1(
            generated_at=datetime.now(UTC),
            baseline_run_id=baseline_run_id,
            target_run_id=target_run_id,
            functional_tests=FunctionalTestsDiff(
                summary={
                    "total": NumberTriplet(baseline=100, target=100, delta=0),
                    "passed": NumberTriplet(baseline=95, target=95, delta=0),
                    "failed": NumberTriplet(baseline=5, target=5, delta=0),
                    "xfailed": NumberTriplet(baseline=0, target=0, delta=0),
                    "xpassed": NumberTriplet(baseline=0, target=0, delta=0),
                    "skipped": NumberTriplet(baseline=0, target=0, delta=0),
                    "pass_rate": NumberTriplet(baseline=95.0, target=95.0, delta=0.0),
                }
            ),
            benchmarks=[],
        )

        json_dict = diff_report.to_json_dict()

        assert json_dict["schema_version"] == "1.0"
        assert json_dict["baseline_run_id"] == str(baseline_run_id)
        assert json_dict["target_run_id"] == str(target_run_id)
        assert json_dict["functional_tests"]["summary"]["total"]["baseline"] == 100
        assert json_dict["benchmarks"] == []

        # Test that it's JSON serializable
        json_str = json.dumps(json_dict)
        assert isinstance(json_str, str)

    def test_missing_required_fields(self):
        """Test that missing required fields raise validation errors."""
        with pytest.raises(ValidationError):
            ReportDiffV1(
                # Missing required fields
                baseline_run_id=uuid.uuid4(),
            )

    def test_invalid_run_ids(self):
        """Test that invalid UUIDs raise validation errors."""
        with pytest.raises(ValidationError):
            ReportDiffV1(
                generated_at=datetime.now(UTC),
                baseline_run_id="not-a-uuid",  # Invalid UUID
                target_run_id=uuid.uuid4(),
                functional_tests=FunctionalTestsDiff(
                    summary={
                        "total": NumberTriplet(baseline=100, target=100, delta=0),
                        "passed": NumberTriplet(baseline=95, target=95, delta=0),
                        "failed": NumberTriplet(baseline=5, target=5, delta=0),
                        "xfailed": NumberTriplet(baseline=0, target=0, delta=0),
                        "xpassed": NumberTriplet(baseline=0, target=0, delta=0),
                        "skipped": NumberTriplet(baseline=0, target=0, delta=0),
                        "pass_rate": NumberTriplet(baseline=95.0, target=95.0, delta=0.0),
                    }
                ),
                benchmarks=[],
            )


class TestJSONSchemaCompliance:
    def test_example_diff_matches_schema(self):
        """Test that the example from the schema document is valid."""
        example_data = {
            "schema_version": "1.0",
            "generated_at": "2025-06-10T22:11:08Z",
            "baseline_run_id": "11111111-1111-1111-1111-111111111111",
            "target_run_id": "22222222-2222-2222-2222-222222222222",
            "functional_tests": {
                "summary": {
                    "passed": {"baseline": 130, "target": 127, "delta": -3},
                    "failed": {"baseline": 0, "target": 3, "delta": 3},
                    "total": {"baseline": 130, "target": 130, "delta": 0},
                    "xfailed": {"baseline": 0, "target": 0, "delta": 0},
                    "xpassed": {"baseline": 0, "target": 0, "delta": 0},
                    "skipped": {"baseline": 0, "target": 0, "delta": 0},
                    "pass_rate": {"baseline": 100, "target": 99.6, "delta": -0.4},
                },
                "regressions": ["tests.chat.test_basic::test_tool"],
                "fixes": ["tests.chat.test_streaming::test_token_stream"],
            },
            "benchmarks": [
                {
                    "name": "mmlu",
                    "metric": "accuracy",
                    "baseline": 0.792,
                    "target": 0.780,
                    "delta": -0.012,
                    "subset_metrics": {
                        "humanities": {
                            "baseline": 0.820,
                            "target": 0.810,
                            "delta": -0.010,
                        },
                        "stem": {"baseline": 0.760, "target": 0.750, "delta": -0.010},
                    },
                },
                {
                    "name": "hellaswag",
                    "metric": "accuracy",
                    "baseline": None,
                    "target": 0.844,
                    "delta": None,
                    "subset_metrics": {},
                },
            ],
        }

        # This should not raise any validation errors
        diff_report = ReportDiffV1(**example_data)

        assert diff_report.schema_version == "1.0"
        assert diff_report.functional_tests.summary["passed"].delta == -3
        assert len(diff_report.benchmarks) == 2
        assert diff_report.benchmarks[0].name == "mmlu"
        assert diff_report.benchmarks[0].delta == -0.012
        assert diff_report.benchmarks[1].baseline is None
