"""Tests for obliteratus.capability_check."""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pytest


class TestCapabilityCheckImport:
    def test_import(self):
        from obliteratus.capability_check import capability_check  # noqa: F401

    def test_quick_subjects_defined(self):
        from obliteratus.capability_check import QUICK_SUBJECTS

        assert len(QUICK_SUBJECTS) >= 3
        assert all(s.startswith("mmlu_") for s in QUICK_SUBJECTS)


class TestRunLmEval:
    @patch("obliteratus.capability_check.subprocess.run")
    def test_calls_lm_eval(self, mock_run, tmp_path):
        from obliteratus.capability_check import _run_lm_eval

        # Create fake results
        results_dir = tmp_path / "results" / "fake_model"
        results_dir.mkdir(parents=True)
        results_file = results_dir / "results_2026.json"
        results_file.write_text(json.dumps({
            "results": {"mmlu": {"acc,none": 0.85}},
        }))

        mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="")

        result = _run_lm_eval("fake/model", "mmlu", 5, "cpu", str(tmp_path / "results"))
        assert result["results"]["mmlu"]["acc,none"] == 0.85
        mock_run.assert_called_once()

    @patch("obliteratus.capability_check.subprocess.run")
    def test_raises_on_failure(self, mock_run, tmp_path):
        from obliteratus.capability_check import _run_lm_eval

        mock_run.return_value = MagicMock(returncode=1, stdout="", stderr="error")

        with pytest.raises(RuntimeError, match="lm-eval exited"):
            _run_lm_eval("fake/model", "mmlu", 5, "cpu", str(tmp_path))


class TestCapabilityCheck:
    @patch("obliteratus.capability_check._run_lm_eval")
    def test_computes_delta(self, mock_lm_eval, tmp_path):
        from obliteratus.capability_check import capability_check

        mock_lm_eval.side_effect = [
            {"results": {"mmlu": {"acc,none": 0.81}}},  # abliterated
            {"results": {"mmlu": {"acc,none": 0.87}}},  # stock
        ]

        result = capability_check(
            "fake/abliterated", "fake/stock",
            device="cpu", output_dir=str(tmp_path),
        )

        assert result["abliterated_acc"] == 0.81
        assert result["stock_acc"] == 0.87
        assert result["delta_pp"] == -6.0

    @patch("obliteratus.capability_check._run_lm_eval")
    def test_quick_mode(self, mock_lm_eval, tmp_path):
        from obliteratus.capability_check import QUICK_SUBJECTS, capability_check

        mock_lm_eval.side_effect = [
            {"results": {s: {"acc,none": 0.8} for s in QUICK_SUBJECTS}},
            {"results": {s: {"acc,none": 0.9} for s in QUICK_SUBJECTS}},
        ]

        result = capability_check(
            "fake/abliterated", "fake/stock",
            device="cpu", quick=True, output_dir=str(tmp_path),
        )

        assert result["tasks"] == ",".join(QUICK_SUBJECTS)
        assert abs(result["abliterated_acc"] - 0.8) < 0.01

    @patch("obliteratus.capability_check._run_lm_eval")
    def test_saves_summary(self, mock_lm_eval, tmp_path):
        from obliteratus.capability_check import capability_check

        mock_lm_eval.side_effect = [
            {"results": {"mmlu": {"acc,none": 0.85}}},
            {"results": {"mmlu": {"acc,none": 0.87}}},
        ]

        capability_check(
            "fake/abliterated", "fake/stock",
            device="cpu", output_dir=str(tmp_path),
        )

        summary_path = tmp_path / "capability_summary.json"
        assert summary_path.exists()
        summary = json.loads(summary_path.read_text())
        assert "delta_pp" in summary
        assert summary["method"] == "lm-eval-harness 0-shot log-likelihood"

    @patch("obliteratus.capability_check._run_lm_eval")
    def test_custom_subjects(self, mock_lm_eval, tmp_path):
        from obliteratus.capability_check import capability_check

        mock_lm_eval.side_effect = [
            {"results": {"mmlu_physics": {"acc,none": 0.75}}},
            {"results": {"mmlu_physics": {"acc,none": 0.80}}},
        ]

        result = capability_check(
            "fake/abliterated", "fake/stock",
            device="cpu", subjects=["mmlu_physics"],
            output_dir=str(tmp_path),
        )

        assert result["tasks"] == "mmlu_physics"
        assert result["abliterated_acc"] == 0.75

    @patch("obliteratus.capability_check._run_lm_eval")
    def test_individual_subjects_mean(self, mock_lm_eval, tmp_path):
        """Test mean computation when no aggregate mmlu key exists."""
        from obliteratus.capability_check import capability_check

        # Results don't have "mmlu" aggregate key or match tasks.split(",")[0]
        mock_lm_eval.side_effect = [
            {"results": {"sub_a": {"acc,none": 0.6}, "sub_b": {"acc,none": 0.8}}},
            {"results": {"sub_a": {"acc,none": 0.7}, "sub_b": {"acc,none": 0.9}}},
        ]

        result = capability_check(
            "fake/abliterated", "fake/stock",
            device="cpu", subjects=["mmlu_a", "mmlu_b"],
            output_dir=str(tmp_path),
        )

        assert abs(result["abliterated_acc"] - 0.7) < 0.01
        assert abs(result["stock_acc"] - 0.8) < 0.01


class TestMainCLI:
    @patch("obliteratus.capability_check.capability_check")
    def test_main_runs(self, mock_check):
        from obliteratus.capability_check import main

        mock_check.return_value = {
            "stock_acc": 0.87, "abliterated_acc": 0.81,
            "delta_pp": -6.0, "tasks": "mmlu", "method": "test",
        }

        import sys
        old_argv = sys.argv
        sys.argv = ["prog", "--abliterated", "fake/abl", "--stock", "fake/stock", "--device", "cpu"]
        try:
            main()
        finally:
            sys.argv = old_argv

        mock_check.assert_called_once()

    @patch("obliteratus.capability_check._run_lm_eval")
    def test_cli_dispatch(self, mock_lm_eval, tmp_path):
        """Test that 'obliteratus capability-check' dispatches correctly."""
        from obliteratus.cli import main as cli_main

        mock_lm_eval.side_effect = [
            {"results": {"mmlu": {"acc,none": 0.81}}},
            {"results": {"mmlu": {"acc,none": 0.87}}},
        ]

        cli_main(["capability-check",
                   "--abliterated", "fake/abl",
                   "--stock", "fake/stock",
                   "--device", "cpu",
                   "--output-dir", str(tmp_path)])