"""Tests for litellm migration in proposer and trajectory summarizer.

Unit tests mock litellm.completion to verify call patterns and response parsing.
Integration smoke test (marked with pytest.mark.integration) makes a real API call.
"""

from __future__ import annotations

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

import pytest


# ---------------------------------------------------------------------------
# Helpers: build a fake litellm response
# ---------------------------------------------------------------------------

def _make_litellm_response(content: str):
    """Build a mock object matching litellm.completion() return shape."""
    message = SimpleNamespace(content=content)
    choice = SimpleNamespace(message=message)
    return SimpleNamespace(choices=[choice])


# ---------------------------------------------------------------------------
# Proposer._call_llm
# ---------------------------------------------------------------------------

class TestProposerCallLlm:
    def test_parses_new_skill_tags(self):
        from clawbench.proposer import Proposer

        fake_resp = _make_litellm_response(
            "Here is the improved skill:\n"
            "<new_skill>\n# Better Skill\nDo things better.\n</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")

        assert result == "# Better Skill\nDo things better."
        mock_ll.completion.assert_called_once()
        call_kw = mock_ll.completion.call_args
        assert call_kw.kwargs["model"] == proposer.model
        assert call_kw.kwargs["messages"][0]["content"] == "improve this"

    def test_returns_full_text_without_tags(self):
        from clawbench.proposer import Proposer

        fake_resp = _make_litellm_response("# Just raw skill content\nNo tags here.")
        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.return_value = fake_resp
            result = Proposer()._call_llm("improve this")

        assert result == "# Just raw skill content\nNo tags here."

    def test_returns_empty_on_exception(self):
        from clawbench.proposer import Proposer

        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.side_effect = RuntimeError("API down")
            result = Proposer()._call_llm("improve this")

        assert result == ""

    def test_writes_error_file_on_exception(self):
        from clawbench.proposer import Proposer

        with tempfile.TemporaryDirectory() as tmpdir:
            run_dir = Path(tmpdir) / "proposals"
            with patch("clawbench.proposer.litellm") as mock_ll:
                mock_ll.completion.side_effect = RuntimeError("timeout")
                proposer = Proposer(run_dir=run_dir)
                proposer._call_llm("improve this")

            error_file = run_dir / "error_0.txt"
            assert error_file.exists()
            assert "timeout" in error_file.read_text()

    def test_handles_none_content(self):
        from clawbench.proposer import Proposer

        fake_resp = _make_litellm_response(None)
        # content is None — our code does `or ""`
        fake_resp.choices[0].message.content = None
        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.return_value = fake_resp
            result = Proposer()._call_llm("improve this")

        assert result == ""

    def test_custom_model_override(self):
        from clawbench.proposer import Proposer

        fake_resp = _make_litellm_response("<new_skill>ok</new_skill>")
        with patch("clawbench.proposer.litellm") as mock_ll:
            mock_ll.completion.return_value = fake_resp
            proposer = Proposer(model="anthropic/claude-sonnet-4-6")
            proposer._call_llm("test")

        assert mock_ll.completion.call_args.kwargs["model"] == "anthropic/claude-sonnet-4-6"


# ---------------------------------------------------------------------------
# _llm_summarize_trajectory
# ---------------------------------------------------------------------------

class TestLlmSummarizeTrajectory:
    """Tests for _llm_summarize_trajectory.

    The function uses a local `import litellm`, so we patch `litellm.completion`
    at the module level rather than on the runner namespace.
    """

    def test_returns_summary(self):
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        fake_resp = _make_litellm_response("The agent tried X and failed because Y.")
        with patch("litellm.completion", return_value=fake_resp) as mock_comp:
            result = _llm_summarize_trajectory(
                "[Bash]: gws gmail list", "test-task", 0.0
            )

        assert result == "The agent tried X and failed because Y."
        call_kw = mock_comp.call_args.kwargs
        assert call_kw["max_tokens"] == 1024
        assert "test-task" in call_kw["messages"][0]["content"]

    def test_returns_none_on_exception(self):
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        with patch("litellm.completion", side_effect=RuntimeError("network error")):
            result = _llm_summarize_trajectory(
                "[Bash]: curl ...", "test-task", 1.0
            )

        assert result is None

    def test_returns_none_on_empty_content(self):
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        fake_resp = _make_litellm_response("")
        with patch("litellm.completion", return_value=fake_resp):
            result = _llm_summarize_trajectory(
                "[Bash]: echo hi", "test-task", 1.0
            )

        # Empty string -> strip -> falsy -> None
        assert result is None

    def test_returns_none_on_none_content(self):
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        fake_resp = _make_litellm_response(None)
        fake_resp.choices[0].message.content = None
        with patch("litellm.completion", return_value=fake_resp):
            result = _llm_summarize_trajectory(
                "[Bash]: echo hi", "test-task", 1.0
            )

        assert result is None


# ---------------------------------------------------------------------------
# DEFAULT_OPTIMIZER_MODEL wiring
# ---------------------------------------------------------------------------

class TestDefaultOptimizerModel:
    def test_constant_value(self):
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL

        assert DEFAULT_OPTIMIZER_MODEL == "gemini/gemini-3.1-flash-lite-preview"

    def test_proposer_uses_default(self):
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL
        from clawbench.proposer import Proposer

        proposer = Proposer()
        assert proposer.model == DEFAULT_OPTIMIZER_MODEL

    def test_summarizer_uses_default(self):
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        fake_resp = _make_litellm_response("summary")
        with patch("litellm.completion", return_value=fake_resp) as mock_comp:
            _llm_summarize_trajectory("[Bash]: ls", "t", 1.0)

        assert mock_comp.call_args.kwargs["model"] == DEFAULT_OPTIMIZER_MODEL


# ---------------------------------------------------------------------------
# GEPA engine path: verify it uses litellm via shared components
# ---------------------------------------------------------------------------

class TestGepaLitellmWiring:
    """Verify the GEPA engine path uses DEFAULT_OPTIMIZER_MODEL and litellm."""

    def test_gepa_reflection_lm_uses_default_model(self):
        """GEPA ReflectionConfig.reflection_lm should be DEFAULT_OPTIMIZER_MODEL."""
        from clawbench._env import DEFAULT_OPTIMIZER_MODEL

        # Mock GEPA imports so we can inspect config without GEPA installed
        mock_gepa = MagicMock()
        captured_config = {}

        def capture_optimize_anything(seed_candidate, evaluator, dataset, valset, config):
            captured_config["reflection_lm"] = config.reflection.reflection_lm
            # Return a fake result
            result = SimpleNamespace(
                best_idx=0,
                best_candidate=seed_candidate,
                candidates=[seed_candidate],
                val_subscores=[{}],
                val_aggregate_scores=[0.5],
                total_metric_calls=1,
            )
            return result

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

        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

                with tempfile.TemporaryDirectory() as tmpdir:
                    # Create a minimal seed skill file
                    skill_file = Path(tmpdir) / "SKILL.md"
                    skill_file.write_text("# Test Skill")
                    output_dir = Path(tmpdir) / "output"
                    output_dir.mkdir()

                    tasks = [{"task_name": "t1", "task_path": "/fake/t1"}]
                    _run_gepa_engine(
                        skill_path=skill_file,
                        train_tasks=tasks,
                        val_tasks=tasks,
                        test_tasks=[],
                        agent_name="claude-code",
                        model_name="anthropic/claude-sonnet-4-6",
                        environment_type="docker",
                        max_iter=20,
                        n_workers=1,
                        output_dir=output_dir,
                        seed=42,
                        agent_import_path=None,
                    )

        assert captured_config["reflection_lm"] == DEFAULT_OPTIMIZER_MODEL

    def test_gepa_proposer_calls_litellm(self):
        """GEPA's custom_candidate_proposer (Proposer) uses litellm."""
        fake_resp = _make_litellm_response("<new_skill># Improved\nBetter.</new_skill>")

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

            proposer = Proposer()
            result = proposer._call_llm("improve this skill")

        assert result == "# Improved\nBetter."
        mock_ll.completion.assert_called_once()
        # Verify it's using litellm, not anthropic
        call_kw = mock_ll.completion.call_args.kwargs
        assert "model" in call_kw
        assert "messages" in call_kw

    def test_gepa_engine_does_not_import_anthropic(self):
        """optimize.py and proposer.py should not import anthropic."""
        import clawbench.optimize as ho_mod
        import clawbench.proposer as hp_mod

        for mod in [ho_mod, hp_mod]:
            source = Path(mod.__file__).read_text()
            assert "import anthropic" not in source, f"{mod.__name__} still imports anthropic"


# ---------------------------------------------------------------------------
# Integration smoke test (real API call, skipped without key)
# ---------------------------------------------------------------------------

def _ensure_google_api_key() -> bool:
    """Load key from project env files if not already in environment."""
    if os.environ.get("GOOGLE_API_KEY"):
        return True
    from clawbench._env import load_local_env_files
    load_local_env_files()
    return bool(os.environ.get("GOOGLE_API_KEY"))


@pytest.mark.integration
@pytest.mark.skipif(
    not _ensure_google_api_key(),
    reason="GOOGLE_API_KEY not set and .gemini/.env not found",
)
class TestLitellmSmoke:
    def test_proposer_real_call(self):
        """Verify litellm actually talks to the Gemini API."""
        from clawbench.proposer import Proposer

        proposer = Proposer()
        result = proposer._call_llm("Say 'hello' and nothing else.")
        assert len(result) > 0
        assert "hello" in result.lower()

    def test_summarizer_real_call(self):
        """Verify trajectory summarizer works end-to-end."""
        from clawbench.trajectory import llm_summarize_trajectory as _llm_summarize_trajectory

        result = _llm_summarize_trajectory(
            "[Bash]: echo hello\n  → hello",
            "smoke-test",
            1.0,
        )
        assert result is not None
        assert len(result) > 0
