"""Tests for the core monkey-patching module (record/replay/stop)."""

import json
import tempfile
import threading
from pathlib import Path
from unittest.mock import MagicMock

import httpx
import pytest

import litmus.patch as patch
from litmus.models import Trace, LLMCall, LLMRequest, LLMResponse, TraceMetadata


def _reset_patch_state():
    """Reset all patch module global state between tests."""
    patch._unpatch_sdks()
    patch._active_trace = None
    patch._mode = None
    patch._replay_trace = None
    patch._replay_index = 0
    patch._patched = False
    patch._atexit_registered = False


def _make_trace_file(tmpdir: str, n_steps: int = 2) -> str:
    """Create a trace file with n_steps LLM calls."""
    trace = Trace(agent_id="test", metadata=TraceMetadata())
    for i in range(n_steps):
        trace.add_step(LLMCall(
            request=LLMRequest(
                provider="anthropic", model="claude-sonnet-4-20250514",
                method="POST", path="/v1/messages",
                body={"model": "claude-sonnet-4-20250514", "messages": [{"role": "user", "content": f"test {i}"}]},
            ),
            response=LLMResponse(
                status_code=200,
                body={
                    "content": [{"type": "text", "text": f"response {i}"}],
                    "stop_reason": "end_turn",
                    "model": "claude-sonnet-4-20250514",
                    "usage": {"input_tokens": 10, "output_tokens": 5},
                },
                is_stream=False,
            ),
            latency_ms=100,
        ))
    trace.complete()
    path = str(Path(tmpdir) / "test.trace.json")
    trace.save(path)
    return path


class TestRecord:
    def setup_method(self):
        _reset_patch_state()

    def teardown_method(self):
        _reset_patch_state()

    def test_record_returns_trace(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace = patch.record(output_dir=tmpdir)
            assert isinstance(trace, Trace)
            assert trace.trace_id.startswith("lt-")
            patch.stop()

    def test_record_sets_mode(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            patch.record(output_dir=tmpdir)
            assert patch._mode == "record"
            patch.stop()
            assert patch._mode is None

    def test_record_patches_httpx(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            patch.record(output_dir=tmpdir)
            assert patch._patched
            assert "httpx_sync_send" in patch._original_transports
            patch.stop()
            assert not patch._patched

    def test_stop_saves_trace_file(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            patch.record(output_dir=tmpdir)
            path = patch.stop()
            assert path is not None
            assert Path(path).exists()
            loaded = Trace.load(path)
            assert loaded.status.value == "completed"

    def test_stop_idempotent(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            patch.record(output_dir=tmpdir)
            path1 = patch.stop()
            path2 = patch.stop()
            assert path1 is not None
            assert path2 is None  # second call returns None

    def test_record_with_agent_id(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace = patch.record(output_dir=tmpdir, agent_id="my-agent")
            assert trace.agent_id == "my-agent"
            patch.stop()


class TestReplay:
    def setup_method(self):
        _reset_patch_state()

    def teardown_method(self):
        _reset_patch_state()

    def test_replay_loads_trace(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir, 2)
            loaded = patch.replay(trace_path)
            assert len(loaded.steps) == 2
            patch.stop()

    def test_replay_sets_mode(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir)
            patch.replay(trace_path)
            assert patch._mode == "replay"
            patch.stop()

    def test_replay_file_not_found(self):
        with pytest.raises(FileNotFoundError, match="Trace file not found"):
            patch.replay("/nonexistent/trace.json")

    def test_replay_invalid_json(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            bad_path = str(Path(tmpdir) / "bad.trace.json")
            Path(bad_path).write_text("not json")
            with pytest.raises(ValueError, match="Failed to load trace"):
                patch.replay(bad_path)

    def test_replay_with_faults(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir, 2)
            loaded = patch.replay(trace_path, faults=["llm_refuse:step=0"])
            # Step 0 should be mutated to a refusal
            step0 = loaded.steps[0]
            assert isinstance(step0, LLMCall)
            assert "sorry" in step0.response.body["content"][0]["text"].lower()
            # Step 1 should be unchanged
            step1 = loaded.steps[1]
            assert step1.response.body["content"][0]["text"] == "response 1"
            patch.stop()


class TestReplayNextStep:
    def setup_method(self):
        _reset_patch_state()

    def teardown_method(self):
        _reset_patch_state()

    def test_serves_recorded_responses_in_order(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir, 3)
            patch.replay(trace_path)

            # Create a fake httpx request
            request = MagicMock()
            request.url = "https://api.anthropic.com/v1/messages"

            for i in range(3):
                resp = patch._replay_next_step(request)
                assert resp.status_code == 200
                body = json.loads(resp.content)
                assert body["content"][0]["text"] == f"response {i}"

            patch.stop()

    def test_exhaustion_returns_500(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir, 1)
            patch.replay(trace_path)

            request = MagicMock()

            # First call succeeds
            resp1 = patch._replay_next_step(request)
            assert resp1.status_code == 200

            # Second call exhausted
            resp2 = patch._replay_next_step(request)
            assert resp2.status_code == 500
            body = json.loads(resp2.content)
            assert "exhausted" in body["error"].lower()

            patch.stop()


class TestThreadSafety:
    def setup_method(self):
        _reset_patch_state()

    def teardown_method(self):
        _reset_patch_state()

    def test_concurrent_replay_no_race(self):
        """Multiple threads replaying should each get sequential steps."""
        with tempfile.TemporaryDirectory() as tmpdir:
            trace_path = _make_trace_file(tmpdir, 10)
            patch.replay(trace_path)

            request = MagicMock()
            results = []
            errors = []

            def fetch():
                try:
                    resp = patch._replay_next_step(request)
                    if resp.status_code == 200:
                        body = json.loads(resp.content)
                        results.append(body["content"][0]["text"])
                except Exception as e:
                    errors.append(str(e))

            threads = [threading.Thread(target=fetch) for _ in range(10)]
            for t in threads:
                t.start()
            for t in threads:
                t.join()

            assert len(errors) == 0
            # All 10 responses should be served (no duplicates, no skips)
            assert len(results) == 10
            assert set(results) == {f"response {i}" for i in range(10)}

            patch.stop()


class TestGetTrace:
    def setup_method(self):
        _reset_patch_state()

    def teardown_method(self):
        _reset_patch_state()

    def test_get_trace_during_recording(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            trace = patch.record(output_dir=tmpdir)
            assert patch.get_trace() is trace
            patch.stop()
            assert patch.get_trace() is None

    def test_get_trace_when_inactive(self):
        assert patch.get_trace() is None
