import asyncio
from copy import deepcopy
import json
from types import SimpleNamespace

from langchain_core.messages import AIMessageChunk, HumanMessage, ToolMessage
from langchain_core.tools import tool
import pytest

from artemis.agents.checker import checker
from artemis.config.llm import LLM
from artemis.llm import final_review
from artemis.llm.diagnostics import ACTIVE_ATTEMPT
from artemis.llm.final_review import (
    ACTIVE_FINAL_REVIEW,
    FinalReviewBudget,
    FinalReviewBudgetExceeded,
)
from artemis.llm.reliability import LLMExhaustedError
from artemis.llm.router import ModelEndpoint
from artemis.services import llm


class RecordingEngine:
    current_session_id = "review-test"
    current_step_id = None

    def __init__(self):
        self.records = []

    def record_trace(self, **record):
        self.records.append(deepcopy(record))

    def _publish(self, *_):
        pass

    def latest(self, kind):
        return next(r["payload"] for r in reversed(self.records) if r["type"] == kind)


@pytest.fixture
def engine(monkeypatch, tmp_path):
    engine = RecordingEngine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path))
    monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
    )
    llm._ENDPOINT_BREAKER._states.clear()
    yield engine
    llm._ENDPOINT_BREAKER._states.clear()
    assert ACTIVE_FINAL_REVIEW.get() is None
    assert ACTIVE_ATTEMPT.get() is None


def budget(engine, *, timeout=1, retries=1):
    return FinalReviewBudget(
        engine, attempt_id="final#1", request_timeout=timeout, max_retries=retries, retry_delay=0
    )


def test_budget_uses_effective_checker_config_and_respects_overall_cap(monkeypatch):
    monkeypatch.setattr(llm.settings, "OPENAI_TIMEOUT_SECONDS", 300)
    monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 5)
    monkeypatch.setattr(llm.settings, "OPENAI_RETRY_INTERVAL_SECONDS", 10)
    monkeypatch.setattr(llm.settings, "OPENAI_RATE_LIMIT_RETRY_SECONDS", 60)
    cfg = LLM(model="test", timeout=20, max_retries=2)
    ctx = SimpleNamespace(llm_config=SimpleNamespace(get_agent=lambda name: cfg), data_engine=None)
    review = llm.final_review_budget(ctx, attempt_id="final#1")
    assert review.total_seconds == 280  # three normal calls + two retries + waits/tool reads
    assert review.data["request_timeout_seconds"] == 20
    assert review.data["max_retries"] == 2

    cfg.timeout = 300
    cfg.max_retries = 100
    assert llm.final_review_budget(ctx, attempt_id="final#2").total_seconds == 3600


@pytest.mark.asyncio
async def test_checker_retries_only_failed_request_with_existing_tool_evidence(engine, monkeypatch):
    calls = []
    evidence_reads = 0
    stream_closed = asyncio.Event()

    @tool
    async def read_evidence() -> str:
        """Read the already-recorded page evidence."""
        nonlocal evidence_reads
        evidence_reads += 1
        return "实际为中文搜索，预期为英文搜索"

    class ScriptedModel:
        def bind_tools(self, *args, **_):
            return self

        async def astream(self, messages, **_):
            calls.append(deepcopy(messages))
            if len(calls) == 1:
                yield AIMessageChunk(
                    content="",
                    tool_calls=[{"name": "read_evidence", "args": {}, "id": "read-1"}],
                )
            elif len(calls) == 2:
                try:
                    yield AIMessageChunk(content='{"test_result":{"status":"passed","reason":"')
                    await asyncio.Event().wait()  # the *request* timer triggers a retry
                finally:
                    stream_closed.set()
            else:
                yield AIMessageChunk(
                    content=json.dumps(
                        {
                            "test_result": {
                                "status": "failed",
                                "reason": "中文搜索不满足英文预期",
                                "evidence": ["read_evidence：实际为中文搜索"],
                            }
                        }
                    )
                )

    wrapper = llm.RobustChatModelWrapper(
        ScriptedModel(),
        endpoint=ModelEndpoint(model_name="test", timeout_seconds=0.02, max_retries=1),
    )
    monkeypatch.setattr(checker, "get_llm", lambda **_: wrapper)
    monkeypatch.setattr(checker, "trace_langchain_tool", lambda tool, ctx: tool)
    events = []
    monkeypatch.setattr(llm, "_record_llm_event", lambda *args, **kwargs: events.append(args))
    ctx = SimpleNamespace(execution_setup=SimpleNamespace(checker_max_iterations=4))
    review = budget(engine)
    report = await review.run(
        lambda: checker._run_check_loop(
            ctx, [HumanMessage(content="核实英文搜索")], [read_evidence], []
        )
    )

    assert report.test_result.status == "failed"  # partial `passed` was never a verdict
    assert evidence_reads == 1
    assert len(calls) == 3
    assert calls[1] == calls[2]
    assert any(isinstance(m, ToolMessage) and "中文搜索" in m.content for m in calls[2])
    assert stream_closed.is_set()
    assert any(e[0] == "llm_stream_reset" and e[1]["action"] == "discard" for e in events)
    stored = engine.latest("final_review_diagnostic")
    assert stored["status"] == "completed"
    assert stored["request_count"] == 2
    assert stored["retry_count"] == 1
    assert stored["last_failure_category"] == "timeout"
    assert engine.latest("llm_diagnostic")["final_review"]["review_id"] == stored["review_id"]


@pytest.mark.asyncio
@pytest.mark.parametrize("phase", ["request", "retry_wait", "evidence"])
async def test_total_budget_expiry_is_not_user_cancellation_or_a_new_retry(
    engine, monkeypatch, tmp_path, phase
):
    monkeypatch.setattr(final_review, "MAX_FINAL_REVIEW_SECONDS", 0.02)
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 60, max_attempts=6)
    )
    cancelled = asyncio.Event()
    calls = 0

    async def request():
        nonlocal calls
        calls += 1
        if phase == "retry_wait":
            raise ConnectionError("connection reset")
        try:
            await asyncio.Event().wait()
        finally:
            cancelled.set()

    review = budget(engine, retries=5)

    async def operation():
        if phase == "evidence":
            return await request()
        return await llm._run_with_recovery(request, request_timeout=1, max_retries=5)

    with pytest.raises(FinalReviewBudgetExceeded, match="总预算已用尽"):
        await review.run(operation)
    assert calls == 1
    assert not (tmp_path / "pause").exists()
    assert engine.latest("final_review_diagnostic")["status"] == "budget_exhausted"
    if phase != "retry_wait":
        assert cancelled.is_set()
    if phase != "evidence":
        stored = engine.latest("llm_diagnostic")
        assert stored["phase"] == "failed"
        assert stored["category"] == "final_review_budget_exhausted"
        assert stored["timeout_scope"] == "final_review"
        assert "cancel_reason" not in stored


@pytest.mark.asyncio
async def test_exhausted_final_review_request_ends_without_hidden_pause(
    engine, monkeypatch, tmp_path
):
    calls = 0

    async def fails():
        nonlocal calls
        calls += 1
        raise TimeoutError()

    pause = []
    monkeypatch.setattr(llm, "_handle_llm_pause_and_resume", lambda *_: pause.append(True))
    with pytest.raises(LLMExhaustedError):
        await budget(engine).run(
            lambda: llm._run_with_recovery(fails, request_timeout=1, max_retries=1)
        )
    assert calls == 2
    assert not pause
    assert not (tmp_path / "pause").exists()
    assert engine.latest("final_review_diagnostic")["status"] == "request_retries_exhausted"


@pytest.mark.asyncio
async def test_user_cancellation_closes_review_and_child_without_retry(engine):
    started = asyncio.Event()
    cancelled = asyncio.Event()
    calls = 0

    async def request():
        nonlocal calls
        calls += 1
        started.set()
        try:
            await asyncio.Event().wait()
        finally:
            cancelled.set()

    task = asyncio.create_task(
        budget(engine).run(
            lambda: llm._run_with_recovery(request, request_timeout=1, max_retries=5)
        )
    )
    await started.wait()
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert calls == 1 and cancelled.is_set()
    assert engine.latest("final_review_diagnostic")["status"] == "cancelled"
    assert engine.latest("llm_diagnostic")["category"] == "cancelled"