import asyncio
import logging

import pytest

from artemis.services import llm


@pytest.mark.asyncio
@pytest.mark.parametrize("replace_pause", [False, True])
async def test_cancelled_wait_clears_only_its_own_pause(tmp_path, monkeypatch, replace_pause):
    pause = tmp_path / "pause"
    pause.write_text("old request")
    entered = asyncio.Event()

    async def wait(_seconds):
        entered.set()
        await asyncio.Event().wait()

    monkeypatch.setattr(llm.asyncio, "sleep", wait)
    task = asyncio.create_task(llm._wait_for_resume(pause))
    await entered.wait()
    if replace_pause:
        pause.write_text("a newer request owns this pause")
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert pause.exists() is replace_pause


@pytest.mark.asyncio
async def test_pause_deadline_clears_file(tmp_path, monkeypatch):
    from unittest.mock import AsyncMock

    pause = tmp_path / "pause"
    pause.write_text("failed request")
    monkeypatch.setattr(llm.settings, "LLM_PAUSE_TIMEOUT_SECONDS", 1)
    monkeypatch.setattr(llm.asyncio, "sleep", AsyncMock())
    assert await llm._wait_for_resume(pause) is False
    assert not pause.exists()


@pytest.mark.asyncio
async def test_timeout_wrapper_cancels_llm_call_when_caller_is_cancelled():
    started = asyncio.Event()
    child_cancelled = asyncio.Event()

    async def pending_llm_call():
        started.set()
        try:
            await asyncio.Event().wait()
        except asyncio.CancelledError:
            child_cancelled.set()
            raise

    invocation = asyncio.create_task(
        llm.invoke_llm_with_timeout_message(
            pending_llm_call(),
            timeout_seconds=60,
        )
    )
    await started.wait()
    invocation.cancel()

    with pytest.raises(asyncio.CancelledError):
        await invocation

    assert child_cancelled.is_set()


class _RecordingEngine:
    current_step_id = "step-7"
    current_session_id = "session-1"

    def __init__(self):
        self.trace_kwargs = None
        self.published = None

    def record_trace(self, **kwargs):
        self.trace_kwargs = kwargs
        return "trace-1"

    def _publish(self, event_type, payload):
        self.published = (event_type, payload)


def test_pause_error_is_persisted_and_published(monkeypatch, tmp_path):
    engine = _RecordingEngine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces"))

    pause_file = llm._handle_llm_pause_and_resume(RuntimeError("503 high demand"))

    assert pause_file.read_text() == "LLM Error: 503 high demand"
    assert engine.trace_kwargs == {
        "type": "llm_call",
        "name": "llm_pause",
        "payload": {"error": "503 high demand", "pause": True},
        "step_id": "step-7",
        "status": "failed",
    }
    event_type, payload = engine.published
    assert event_type == "task_paused"
    assert payload["error"] == "503 high demand"
    assert payload["step_id"] == "step-7"
    assert isinstance(payload["timestamp"], float)


def test_terminal_failure_aggregates_observed_retries_and_total_wait(monkeypatch, tmp_path):
    engine = _RecordingEngine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces"))
    monkeypatch.setattr(llm.time, "time", lambda: 120.0)
    retries = [
        {
            "error": "503 high demand",
            "delay": 10.0,
            "provider": "openai",
            "source": "provider_sdk",
            "request_id": "request-2",
            "scheduled_at": 101.0,
            "timestamp": 101.0,
        }
    ]
    request_token = llm._ACTIVE_LLM_REQUEST.set(
        {
            "request_id": "request-2",
            "provider": "openai",
            "started_at": 100.0,
            "retries": retries,
        }
    )
    try:
        llm._handle_llm_pause_and_resume(RuntimeError("503 high demand"))
    finally:
        llm._ACTIVE_LLM_REQUEST.reset(request_token)

    payload = engine.trace_kwargs["payload"]
    assert payload["request_id"] == "request-2"
    assert payload["provider"] == "openai"
    assert payload["waited_seconds"] == 20.0
    assert payload["retries"] == retries