# 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 pathlib import Path

import pytest

from llama_verifications.reporting.unified_reporter import UnifiedReporter


class TestUnifiedReporter:
    @pytest.fixture
    def temp_results_dir(self):
        """Create a temporary results directory with test data."""
        with tempfile.TemporaryDirectory() as temp_dir:
            results_dir = Path(temp_dir)

            # Create test provider directory structure
            provider_dir = results_dir / "test_provider"
            provider_dir.mkdir(parents=True)

            # Create functional tests data
            tests_dir = provider_dir / "tests"
            tests_dir.mkdir()

            test_data = {
                "summary": {
                    "total": 10,
                    "passed": 8,
                    "failed": 1,
                    "xfailed": 0,
                    "xpassed": 0,
                    "skipped": 1,
                },
                "tests": [
                    {
                        "nodeid": "test_example::test_pass",
                        "outcome": "passed",
                        "duration": 0.15,
                    },
                    {
                        "nodeid": "test_example::test_fail",
                        "outcome": "failed",
                        "duration": 0.05,
                        "call": {"longrepr": "AssertionError: Test failed"},
                    },
                    {
                        "nodeid": "test_example::test_skip",
                        "outcome": "skipped",
                        "duration": 0.0,
                    },
                ],
                "run_timestamp": 1640995200,  # 2022-01-01 00:00:00
            }

            with open(tests_dir / "output.json", "w") as f:
                json.dump(test_data, f)

            # Create benchmarks data
            benchmarks_dir = provider_dir / "benchmarks"
            benchmarks_dir.mkdir()

            benchmark_data = {
                "test_model": {
                    "topline_metric_name": "accuracy",
                    "metrics": {"accuracy": 0.85, "f1": 0.82},
                    "subset_metrics": {
                        "subset1": {"accuracy": 0.87, "f1": 0.84},
                        "subset2": {"accuracy": 0.83, "f1": 0.80},
                    },
                    "config": {"n_shot": 5, "temperature": 0.0},
                }
            }

            with open(benchmarks_dir / "mmlu.json", "w") as f:
                json.dump(benchmark_data, f)

            yield results_dir

    def test_generate_report_with_data(self, temp_results_dir):
        """Test generating a report with functional tests and benchmarks data."""
        reporter = UnifiedReporter(temp_results_dir)

        report = reporter.generate_report(
            provider="test_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
            metadata={"test_key": "test_value"},
        )

        assert report.schema_version == "1.0"
        assert report.model == "test_model"
        assert report.model_checkpoint == "test_checkpoint"
        assert report.provider == "test_provider"
        assert report.metadata["test_key"] == "test_value"

        # Check functional tests
        assert report.functional_tests.summary.total == 10
        assert report.functional_tests.summary.passed == 8
        assert report.functional_tests.summary.failed == 1
        assert report.functional_tests.summary.skipped == 1
        assert report.functional_tests.summary.pass_rate == 80.0
        assert len(report.functional_tests.cases) == 3

        # Check test cases
        test_cases = {case.name: case for case in report.functional_tests.cases}
        assert "test_example::test_pass" in test_cases
        assert test_cases["test_example::test_pass"].status == "passed"
        assert test_cases["test_example::test_pass"].duration_ms == 150.0

        assert "test_example::test_fail" in test_cases
        assert test_cases["test_example::test_fail"].status == "failed"
        assert test_cases["test_example::test_fail"].error == "AssertionError: Test failed"

        # Check benchmarks
        assert report.benchmarks.summary.total == 1
        assert len(report.benchmarks.benchmarks) == 1

        benchmark = report.benchmarks.benchmarks[0]
        assert benchmark.name == "mmlu"
        assert benchmark.topline_metric_name == "accuracy"
        assert benchmark.metrics["accuracy"] == 0.85
        assert benchmark.subset_metrics["subset1"]["accuracy"] == 0.87
        assert benchmark.config["n_shot"] == 5

    def test_generate_report_no_data(self, temp_results_dir):
        """Test generating a report with no existing data."""
        reporter = UnifiedReporter(temp_results_dir)

        report = reporter.generate_report(
            provider="nonexistent_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
        )

        assert report.schema_version == "1.0"
        assert report.model == "test_model"
        assert report.provider == "nonexistent_provider"

        # Should have empty functional tests
        assert report.functional_tests.summary.total == 0
        assert report.functional_tests.summary.passed == 0
        assert report.functional_tests.summary.pass_rate == 0.0
        assert len(report.functional_tests.cases) == 0

        # Should have empty benchmarks
        assert report.benchmarks.summary.total == 0
        assert len(report.benchmarks.benchmarks) == 0

    def test_generate_report_with_output_file(self, temp_results_dir):
        """Test generating a report and saving to file."""
        reporter = UnifiedReporter(temp_results_dir)
        output_path = temp_results_dir / "test_report.json"

        report = reporter.generate_report(
            provider="test_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
            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["model"] == "test_model"
        assert saved_data["run_id"] == str(report.run_id)

    def test_generate_report_with_custom_run_id(self, temp_results_dir):
        """Test generating a report with custom run ID."""
        reporter = UnifiedReporter(temp_results_dir)
        custom_run_id = str(uuid.uuid4())

        report = reporter.generate_report(
            provider="test_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
            run_id=custom_run_id,
        )

        assert str(report.run_id) == custom_run_id

    def test_list_available_runs(self, temp_results_dir):
        """Test listing available runs for a provider."""
        reporter = UnifiedReporter(temp_results_dir)

        runs = reporter.list_available_runs("test_provider")

        assert len(runs) == 2  # functional_tests and benchmarks

        # Check functional tests run
        ft_run = next((r for r in runs if r["type"] == "functional_tests"), None)
        assert ft_run is not None
        assert ft_run["timestamp"] == 1640995200

        # Check benchmarks run
        bench_run = next((r for r in runs if r["type"] == "benchmarks"), None)
        assert bench_run is not None
        assert len(bench_run["files"]) == 1

    def test_list_available_runs_nonexistent_provider(self, temp_results_dir):
        """Test listing runs for nonexistent provider."""
        reporter = UnifiedReporter(temp_results_dir)

        runs = reporter.list_available_runs("nonexistent_provider")

        assert runs == []

    def test_load_functional_tests_invalid_json(self, temp_results_dir):
        """Test loading functional tests with invalid JSON."""
        # Create invalid JSON file
        provider_dir = temp_results_dir / "invalid_provider"
        tests_dir = provider_dir / "tests"
        tests_dir.mkdir(parents=True)

        with open(tests_dir / "output.json", "w") as f:
            f.write("invalid json content")

        reporter = UnifiedReporter(temp_results_dir)

        report = reporter.generate_report(
            provider="invalid_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
        )

        # Should fall back to empty functional tests
        assert report.functional_tests.summary.total == 0
        assert len(report.functional_tests.cases) == 0

    def test_load_benchmarks_invalid_json(self, temp_results_dir):
        """Test loading benchmarks with invalid JSON."""
        # Create invalid JSON file
        provider_dir = temp_results_dir / "invalid_provider"
        benchmarks_dir = provider_dir / "benchmarks"
        benchmarks_dir.mkdir(parents=True)

        with open(benchmarks_dir / "invalid.json", "w") as f:
            f.write("invalid json content")

        reporter = UnifiedReporter(temp_results_dir)

        report = reporter.generate_report(
            provider="invalid_provider",
            model="test_model",
            model_checkpoint="test_checkpoint",
        )

        # Should fall back to empty benchmarks
        assert report.benchmarks.summary.total == 0
        assert len(report.benchmarks.benchmarks) == 0

    def test_load_benchmarks_model_not_in_file(self, temp_results_dir):
        """Test loading benchmarks when model is not in the benchmark file."""
        reporter = UnifiedReporter(temp_results_dir)

        report = reporter.generate_report(
            provider="test_provider",
            model="nonexistent_model",  # Model not in benchmark data
            model_checkpoint="test_checkpoint",
        )

        # Should have empty benchmarks since model not found
        assert report.benchmarks.summary.total == 0
        assert len(report.benchmarks.benchmarks) == 0
