# 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 tempfile
import uuid
from datetime import UTC, datetime
from pathlib import Path

import pytest

from llama_verifications.reporting.diff_generator import DiffGenerator
from llama_verifications.schemas.report_v1 import (
    BenchmarkResult,
    Benchmarks,
    BenchmarksSummary,
    EvalTestCase,
    FunctionalTests,
    FunctionalTestsSummary,
    ReportV1,
)


class TestDiffGenerator:
    @pytest.fixture
    def baseline_report(self):
        """Create a baseline report for testing."""
        return ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="baseline-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=100, passed=95, failed=5, xfailed=0, xpassed=0, skipped=0),
                cases=[
                    EvalTestCase(name="test1", status="passed"),
                    EvalTestCase(name="test2", status="passed"),
                    EvalTestCase(name="test3", status="failed", error="Test error"),
                    EvalTestCase(name="test4", status="passed"),
                ],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=2),
                benchmarks=[
                    BenchmarkResult(
                        name="mmlu",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.78},
                        subset_metrics={
                            "humanities": {"accuracy": 0.81},
                            "stem": {"accuracy": 0.75},
                        },
                    ),
                    BenchmarkResult(
                        name="hellaswag",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.85},
                    ),
                ],
            ),
        )

    @pytest.fixture
    def target_report(self):
        """Create a target report for testing."""
        return ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="target-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=105, passed=98, failed=7, xfailed=0, xpassed=0, skipped=0),
                cases=[
                    EvalTestCase(name="test1", status="failed", error="New error"),  # regression
                    EvalTestCase(name="test2", status="passed"),
                    EvalTestCase(name="test3", status="passed"),  # fix
                    EvalTestCase(name="test4", status="passed"),
                    EvalTestCase(name="test5", status="passed"),  # new test
                ],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=3),
                benchmarks=[
                    BenchmarkResult(
                        name="mmlu",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.80},  # improved
                        subset_metrics={
                            "humanities": {"accuracy": 0.83},  # improved
                            "stem": {"accuracy": 0.77},  # improved
                        },
                    ),
                    BenchmarkResult(
                        name="hellaswag",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.83},  # decreased
                    ),
                    BenchmarkResult(
                        name="truthful_qa",  # new benchmark
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.59},
                    ),
                ],
            ),
        )

    def test_generate_diff(self, baseline_report, target_report):
        """Test generating a diff between two reports."""
        diff_generator = DiffGenerator()

        diff_report = diff_generator.generate_diff(baseline_report, target_report)

        assert diff_report.schema_version == "1.0"
        assert diff_report.baseline_run_id == baseline_report.run_id
        assert diff_report.target_run_id == target_report.run_id

        # Check functional tests diff
        ft_summary = diff_report.functional_tests.summary
        assert ft_summary["total"].baseline == 100
        assert ft_summary["total"].target == 105
        assert ft_summary["total"].delta == 5

        assert ft_summary["passed"].baseline == 95
        assert ft_summary["passed"].target == 98
        assert ft_summary["passed"].delta == 3

        assert ft_summary["failed"].baseline == 5
        assert ft_summary["failed"].target == 7
        assert ft_summary["failed"].delta == 2

        # Check pass rate calculation
        assert ft_summary["pass_rate"].baseline == 95.0
        assert abs(ft_summary["pass_rate"].target - 93.33) < 0.01  # 98/105 ≈ 93.33
        assert abs(ft_summary["pass_rate"].delta - (-1.67)) < 0.01

        # Check regressions and fixes
        assert diff_report.functional_tests.regressions == ["test1"]
        assert diff_report.functional_tests.fixes == ["test3"]

        # Check benchmarks diff
        assert len(diff_report.benchmarks) == 3

        # Check MMLU diff
        mmlu_diff = next(b for b in diff_report.benchmarks if b.name == "mmlu")
        assert mmlu_diff.metric == "accuracy"
        assert mmlu_diff.baseline == 0.78
        assert mmlu_diff.target == 0.80
        assert abs(mmlu_diff.delta - 0.02) < 0.001

        # Check subset metrics
        assert mmlu_diff.subset_metrics["humanities"].baseline == 0.81
        assert mmlu_diff.subset_metrics["humanities"].target == 0.83
        assert abs(mmlu_diff.subset_metrics["humanities"].delta - 0.02) < 0.001

        assert mmlu_diff.subset_metrics["stem"].baseline == 0.75
        assert mmlu_diff.subset_metrics["stem"].target == 0.77
        assert abs(mmlu_diff.subset_metrics["stem"].delta - 0.02) < 0.001

        # Check HellaSwag diff (decreased)
        hellaswag_diff = next(b for b in diff_report.benchmarks if b.name == "hellaswag")
        assert hellaswag_diff.baseline == 0.85
        assert hellaswag_diff.target == 0.83
        assert abs(hellaswag_diff.delta - (-0.02)) < 0.001

        # Check new benchmark
        truthful_qa_diff = next(b for b in diff_report.benchmarks if b.name == "truthful_qa")
        assert truthful_qa_diff.baseline is None
        assert truthful_qa_diff.target == 0.59
        assert truthful_qa_diff.delta is None

    def test_generate_diff_with_output_file(self, baseline_report, target_report):
        """Test generating a diff and saving to file."""
        diff_generator = DiffGenerator()

        with tempfile.TemporaryDirectory() as temp_dir:
            output_path = Path(temp_dir) / "test_diff.json"

            diff_generator.generate_diff(baseline_report, target_report, output_path=output_path)

            # Check that file was created
            assert output_path.exists()

            # Check that file contains valid JSON
            with open(output_path) as f:
                saved_data = json.load(f)

            assert saved_data["schema_version"] == "1.0"
            assert saved_data["baseline_run_id"] == str(baseline_report.run_id)
            assert saved_data["target_run_id"] == str(target_report.run_id)

    def test_generate_diff_from_files(self, baseline_report, target_report):
        """Test generating a diff from JSON files."""
        diff_generator = DiffGenerator()

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_dir = Path(temp_dir)

            # Save reports to files
            baseline_path = temp_dir / "baseline.json"
            target_path = temp_dir / "target.json"
            output_path = temp_dir / "diff.json"

            with open(baseline_path, "w") as f:
                json.dump(baseline_report.to_json_dict(), f)

            with open(target_path, "w") as f:
                json.dump(target_report.to_json_dict(), f)

            # Generate diff from files
            diff_report = diff_generator.generate_diff_from_files(baseline_path, target_path, output_path)

            assert diff_report.baseline_run_id == baseline_report.run_id
            assert diff_report.target_run_id == target_report.run_id
            assert output_path.exists()

    def test_generate_diff_no_changes(self):
        """Test generating a diff when there are no changes."""
        # Create identical reports
        base_report = ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="test-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=10, passed=9, failed=1, xfailed=0, xpassed=0, skipped=0),
                cases=[
                    EvalTestCase(name="test1", status="passed"),
                    EvalTestCase(name="test2", status="failed", error="Test error"),
                ],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=1),
                benchmarks=[
                    BenchmarkResult(
                        name="mmlu",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.78},
                    ),
                ],
            ),
        )

        target_report = ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="test-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=10, passed=9, failed=1, xfailed=0, xpassed=0, skipped=0),
                cases=[
                    EvalTestCase(name="test1", status="passed"),
                    EvalTestCase(name="test2", status="failed", error="Test error"),
                ],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=1),
                benchmarks=[
                    BenchmarkResult(
                        name="mmlu",
                        topline_metric_name="accuracy",
                        metrics={"accuracy": 0.78},
                    ),
                ],
            ),
        )

        diff_generator = DiffGenerator()
        diff_report = diff_generator.generate_diff(base_report, target_report)

        # Check that all deltas are zero
        ft_summary = diff_report.functional_tests.summary
        assert ft_summary["total"].delta == 0
        assert ft_summary["passed"].delta == 0
        assert ft_summary["failed"].delta == 0
        assert ft_summary["pass_rate"].delta == 0.0

        # No regressions or fixes
        assert diff_report.functional_tests.regressions is None
        assert diff_report.functional_tests.fixes is None

        # Benchmark delta should be zero
        assert len(diff_report.benchmarks) == 1
        assert diff_report.benchmarks[0].delta == 0.0

    def test_generate_diff_empty_reports(self):
        """Test generating a diff with empty reports."""
        empty_report1 = ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="test-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=0, passed=0, failed=0, xfailed=0, xpassed=0, skipped=0),
                cases=[],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=0),
                benchmarks=[],
            ),
        )

        empty_report2 = ReportV1(
            generated_at=datetime.now(UTC),
            run_id=uuid.uuid4(),
            model="Llama-3-70B",
            model_checkpoint="test-checkpoint",
            provider="Groq",
            functional_tests=FunctionalTests(
                summary=FunctionalTestsSummary(total=0, passed=0, failed=0, xfailed=0, xpassed=0, skipped=0),
                cases=[],
            ),
            benchmarks=Benchmarks(
                summary=BenchmarksSummary(total=0),
                benchmarks=[],
            ),
        )

        diff_generator = DiffGenerator()
        diff_report = diff_generator.generate_diff(empty_report1, empty_report2)

        # All values should be zero
        ft_summary = diff_report.functional_tests.summary
        assert ft_summary["total"].baseline == 0
        assert ft_summary["total"].target == 0
        assert ft_summary["total"].delta == 0

        assert len(diff_report.benchmarks) == 0

    def test_find_test_changes(self):
        """Test finding test regressions and fixes."""
        diff_generator = DiffGenerator()

        baseline_cases = [
            EvalTestCase(name="test1", status="passed"),
            EvalTestCase(name="test2", status="failed", error="Error"),
            EvalTestCase(name="test3", status="passed"),
            EvalTestCase(name="test4", status="skipped"),
        ]

        target_cases = [
            EvalTestCase(name="test1", status="failed", error="New error"),  # regression
            EvalTestCase(name="test2", status="passed"),  # fix
            EvalTestCase(name="test3", status="passed"),  # no change
            EvalTestCase(name="test4", status="skipped"),  # no change
            EvalTestCase(name="test5", status="passed"),  # new test (ignored)
        ]

        regressions, fixes = diff_generator._find_test_changes(baseline_cases, target_cases)

        assert regressions == ["test1"]
        assert fixes == ["test2"]

    def test_generate_subset_metrics_diff(self):
        """Test generating subset metrics diff."""
        diff_generator = DiffGenerator()

        baseline_benchmark = BenchmarkResult(
            name="mmlu",
            topline_metric_name="accuracy",
            metrics={"accuracy": 0.78},
            subset_metrics={
                "humanities": {"accuracy": 0.81},
                "stem": {"accuracy": 0.75},
            },
        )

        target_benchmark = BenchmarkResult(
            name="mmlu",
            topline_metric_name="accuracy",
            metrics={"accuracy": 0.80},
            subset_metrics={
                "humanities": {"accuracy": 0.83},
                "stem": {"accuracy": 0.77},
                "social_sciences": {"accuracy": 0.79},  # new subset
            },
        )

        subset_diff = diff_generator._generate_subset_metrics_diff(baseline_benchmark, target_benchmark, "accuracy")

        assert subset_diff is not None
        assert len(subset_diff) == 3

        # Check existing subsets
        assert subset_diff["humanities"].baseline == 0.81
        assert subset_diff["humanities"].target == 0.83
        assert abs(subset_diff["humanities"].delta - 0.02) < 0.001

        assert subset_diff["stem"].baseline == 0.75
        assert subset_diff["stem"].target == 0.77
        assert abs(subset_diff["stem"].delta - 0.02) < 0.001

        # Check new subset
        assert subset_diff["social_sciences"].baseline is None
        assert subset_diff["social_sciences"].target == 0.79
        assert subset_diff["social_sciences"].delta is None
