"""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"