"""Lightweight tests for the GEPA engine path.

These tests verify wiring, config shape, artifact output, error handling,
and litellm integration for the temporary --engine gepa path.
Tests mock GEPA internals — no real optimization runs.
"""

from __future__ import annotations

import json
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_gepa_mocks():
    """Build mock GEPA modules and a capturing optimize_anything."""
    mock_gepa = MagicMock()
    captured = {}

    def capture_optimize_anything(seed_candidate, evaluator, dataset, valset, config):
        captured["seed_candidate"] = seed_candidate
        captured["dataset"] = dataset
        captured["valset"] = valset
        captured["config"] = config
        return SimpleNamespace(
            best_idx=0,
            best_candidate=seed_candidate,
            candidates=[seed_candidate],
            val_subscores=[{}],
            val_aggregate_scores=[0.5],
            total_metric_calls=3,
        )

    mock_gepa.optimize_anything.optimize_anything = capture_optimize_anything
    mock_gepa.optimize_anything.GEPAConfig = MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw))
    mock_gepa.optimize_anything.EngineConfig = MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw))
    mock_gepa.optimize_anything.ReflectionConfig = MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw))
    mock_gepa.optimize_anything.TrackingConfig = MagicMock(side_effect=lambda **kw: SimpleNamespace(**kw))

    return mock_gepa, captured


def _run_with_mocks(mock_gepa, tmpdir, **overrides):
    """Run _run_gepa_engine with mocked GEPA and sensible defaults."""
    skill_file = Path(tmpdir) / "SKILL.md"
    skill_file.write_text("# Test Skill\nDo stuff.")
    output_dir = Path(tmpdir) / "output"
    output_dir.mkdir()

    tasks = [
        {"task_name": "task-a", "task_path": "/fake/task-a"},
        {"task_name": "task-b", "task_path": "/fake/task-b"},
    ]
    defaults = dict(
        skill_path=skill_file,
        train_tasks=tasks,
        val_tasks=tasks,
        test_tasks=[],
        agent_name="gemini-cli",
        model_name="google/gemini-3.1-flash-lite-preview",
        environment_type="docker",
        max_iter=20,
        n_workers=1,
        output_dir=output_dir,
        seed=42,
        agent_import_path=None,
    )
    defaults.update(overrides)

    with patch.dict("sys.modules", {
        "gepa": mock_gepa,
        "gepa.optimize_anything": mock_gepa.optimize_anything,
    }):
        with patch("clawbench.optimize._create_fitness_fn", return_value=MagicMock()):
            from clawbench.optimize import _run_gepa_engine
            return _run_gepa_engine(**defaults), output_dir


# ---------------------------------------------------------------------------
# Config shape tests
# ---------------------------------------------------------------------------

class TestGepaConfigWiring:
    """Verify _run_gepa_engine passes correct config to optimize_anything."""

    def test_engine_config_fields(self):
        mock_gepa, captured = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _run_with_mocks(mock_gepa, tmpdir, max_iter=30, seed=99)

        engine = captured["config"].engine
        assert engine.max_candidate_proposals == 30
        # Safety cap: max_iter * minibatch_size(=n_workers=1..but 2 tasks) * 4
        # With 2 tasks and n_workers=1: minibatch_size=1, so safety = 30*1*4 = 120
        assert engine.max_metric_calls == 120
        assert engine.seed == 99
        assert engine.parallel is False  # n_workers=1 -> not parallel

    def test_reflection_config_uses_default_optimizer_model(self):
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL
        mock_gepa, captured = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _run_with_mocks(mock_gepa, tmpdir)

        assert captured["config"].reflection.reflection_lm == DEFAULT_OPTIMIZER_MODEL

    def test_seed_candidate_contains_skill_content(self):
        mock_gepa, captured = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _run_with_mocks(mock_gepa, tmpdir)

        assert captured["seed_candidate"]["skills"] == "# Test Skill\nDo stuff."

    def test_proposer_is_proposer(self):
        mock_gepa, captured = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _run_with_mocks(mock_gepa, tmpdir)

        proposer = captured["config"].reflection.custom_candidate_proposer
        assert type(proposer).__name__ == "Proposer"


# ---------------------------------------------------------------------------
# Artifact output tests
# ---------------------------------------------------------------------------

class TestGepaArtifacts:
    """Verify _run_gepa_engine writes expected dashboard-compatible artifacts."""

    def test_writes_best_skill(self):
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            best = output_dir / "best_SKILL.md"
            assert best.exists()
            assert "Test Skill" in best.read_text()

    def test_writes_optimization_summary(self):
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            summary_path = output_dir / "optimization_summary.json"
            assert summary_path.exists()
            summary = json.loads(summary_path.read_text())
            assert "best_idx" in summary
            assert "total_metric_calls" in summary
            assert summary["total_metric_calls"] == 3
            assert "config" in summary
            assert summary["config"]["agent_name"] == "gemini-cli"

    def test_writes_baseline_skill(self):
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            assert (output_dir / "skill_iter_0" / "SKILL.md").exists(), \
                "baseline skill (seed) should be saved as skill_iter_0/SKILL.md"
            baseline = (output_dir / "skill_iter_0" / "SKILL.md").read_text()
            assert "Test Skill" in baseline

    def test_writes_iter_results_and_skill_dirs(self):
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            # 1 candidate -> iter_1
            assert (output_dir / "iter_1_results.json").exists()
            assert (output_dir / "skill_iter_1" / "SKILL.md").exists()

    def test_gepa_iter_results_have_null_train(self):
        """GEPA doesn't track train independently — train should be null."""
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            iter_data = json.loads((output_dir / "iter_1_results.json").read_text())
            assert iter_data["train"] is None, "GEPA should emit train: null"
            assert iter_data["val"] is not None

    def test_writes_results_json(self):
        mock_gepa, _ = _make_gepa_mocks()
        with tempfile.TemporaryDirectory() as tmpdir:
            _, output_dir = _run_with_mocks(mock_gepa, tmpdir)
            results_path = output_dir / "results.json"
            assert results_path.exists()
            results = json.loads(results_path.read_text())
            assert results["engine"] == "gepa"
            assert "iterations_run" in results


# ---------------------------------------------------------------------------
# Error handling
# ---------------------------------------------------------------------------

class TestGepaErrorHandling:
    """Verify graceful errors for the GEPA engine path."""

    def test_missing_gepa_gives_click_exception(self):
        """Missing gepa package should raise ClickException, not raw ImportError."""
        import click

        # Ensure gepa is NOT importable by hiding it
        import sys
        gepa_mod = sys.modules.pop("gepa", None)
        gepa_oa_mod = sys.modules.pop("gepa.optimize_anything", None)

        try:
            with patch.dict("sys.modules", {"gepa": None, "gepa.optimize_anything": None}):
                # Need to reload to pick up the patched sys.modules
                import importlib
                import clawbench.optimize as ho
                importlib.reload(ho)

                with tempfile.TemporaryDirectory() as tmpdir:
                    skill_file = Path(tmpdir) / "SKILL.md"
                    skill_file.write_text("# Test")
                    output_dir = Path(tmpdir) / "output"
                    output_dir.mkdir()

                    with pytest.raises(click.ClickException, match="gepa"):
                        ho._run_gepa_engine(
                            skill_path=skill_file,
                            train_tasks=[{"task_name": "t1", "task_path": "/fake/t1"}],
                            val_tasks=[],
                            test_tasks=[],
                            agent_name="claude-code",
                            model_name="test/model",
                            environment_type="docker",
                            max_iter=5,
                            n_workers=1,
                            output_dir=output_dir,
                            seed=1,
                            agent_import_path=None,
                        )
        finally:
            # Restore
            if gepa_mod is not None:
                sys.modules["gepa"] = gepa_mod
            if gepa_oa_mod is not None:
                sys.modules["gepa.optimize_anything"] = gepa_oa_mod


# ---------------------------------------------------------------------------
# GEPA + litellm integration
# ---------------------------------------------------------------------------

class TestGepaLitellmIntegration:
    """Verify the GEPA engine path routes LLM calls through litellm."""

    def test_proposer_uses_litellm_not_anthropic(self):
        """Proposer (used by GEPA as custom_candidate_proposer) calls litellm."""
        from clawbench.proposer import Proposer

        fake_resp = MagicMock()
        fake_resp.choices = [MagicMock()]
        fake_resp.choices[0].message.content = "<new_skill># Better\nImproved.</new_skill>"

        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.return_value = fake_resp
            proposer = Proposer()
            result = proposer._call_llm("improve this")

        mock_ll.completion.assert_called_once()
        call_kw = mock_ll.completion.call_args.kwargs
        assert "model" in call_kw
        assert "messages" in call_kw
        assert result == "# Better\nImproved."

    def test_proposer_model_is_default_optimizer_model(self):
        """Proposer should use DEFAULT_OPTIMIZER_MODEL for litellm calls."""
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL
        from clawbench.proposer import Proposer

        fake_resp = MagicMock()
        fake_resp.choices = [MagicMock()]
        fake_resp.choices[0].message.content = "<new_skill># Skill</new_skill>"

        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.return_value = fake_resp
            proposer = Proposer()
            proposer._call_llm("test")

        call_kw = mock_ll.completion.call_args.kwargs
        assert call_kw["model"] == DEFAULT_OPTIMIZER_MODEL

    def test_gepa_real_imports_not_missing_deps(self):
        """Smoke-test that gepa and its transitive deps (cloudpickle, etc.) import cleanly."""
        import gepa  # noqa: F401
        from gepa.optimize_anything import optimize_anything, GEPAConfig, EngineConfig  # noqa: F401
        # cloudpickle is used by gepa internally for state saving
        import cloudpickle  # noqa: F401

    def test_trajectory_summarizer_uses_litellm(self):
        """_llm_summarize_trajectory should call litellm, not anthropic."""
        import litellm as real_litellm
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        fake_resp = MagicMock()
        fake_resp.choices = [MagicMock()]
        fake_resp.choices[0].message.content = "Agent did stuff."

        with patch.object(real_litellm, "completion", return_value=fake_resp) as mock_completion:
            result = _llm_summarize_trajectory("Agent ran gws commands", "task-a", 0.5)

        mock_completion.assert_called_once()
        assert result == "Agent did stuff."
