"""Tests for the LLM gateway: complete(), classified recovery, and fallback."""

from types import SimpleNamespace

from langchain_core.messages import AIMessage, AIMessageChunk
import pytest

from artemis.llm.reliability import LLMExhaustedError, LLMPermanentError
from artemis.services import llm as llm_service
from artemis.services.llm import (
    RobustChatModelWrapper,
    acomplete,
    acomplete_structured,
    with_fallback,
)


@pytest.fixture(autouse=True)
def _isolate_gateway_state(monkeypatch, tmp_path):
    llm_service._ENDPOINT_BREAKER._states.clear()
    llm_service._NON_STREAMING_ENDPOINTS.clear()
    monkeypatch.setattr(llm_service, "PAUSE_FILE", tmp_path / ".artemis_paused")
    monkeypatch.setattr(
        llm_service,
        "retry_policy_for",
        lambda category: SimpleNamespace(max_attempts=2, delay_for=lambda attempt: 0.0),
    )
    monkeypatch.setattr(llm_service.settings, "LLM_PAUSE_TIMEOUT_SECONDS", 1.0, raising=False)
    yield
    llm_service._ENDPOINT_BREAKER._states.clear()
    llm_service._NON_STREAMING_ENDPOINTS.clear()


@pytest.mark.asyncio
async def test_mid_stream_failure_retries_without_duplicating_chunks():
    class BreaksMidStream:
        def __init__(self):
            self.astream_calls = 0

        async def astream(self, *args, **kwargs):
            self.astream_calls += 1
            if self.astream_calls == 1:
                yield AIMessageChunk(content="Hello")
                raise RuntimeError("503 service unavailable")
            yield AIMessageChunk(content="Hello")
            yield AIMessageChunk(content=" world")

    base = BreaksMidStream()
    wrapper = RobustChatModelWrapper(base)
    result = await wrapper.complete([])

    assert base.astream_calls == 2
    # The partial first attempt was discarded entirely: no duplicated prefix.
    assert result.content == "Hello world"


@pytest.mark.asyncio
async def test_auth_error_fails_fast_without_retry_or_pause(tmp_path):
    class AuthFails:
        def __init__(self):
            self.astream_calls = 0
            self.ainvoke_calls = 0

        async def astream(self, *args, **kwargs):
            self.astream_calls += 1
            raise RuntimeError("401 unauthorized: invalid api key")
            yield  # pragma: no cover - makes this an async generator

        async def ainvoke(self, *args, **kwargs):
            self.ainvoke_calls += 1
            return AIMessage(content="should never happen")

    base = AuthFails()
    wrapper = RobustChatModelWrapper(base)
    with pytest.raises(LLMPermanentError):
        await wrapper.complete([])

    assert base.astream_calls == 1
    assert base.ainvoke_calls == 0
    assert not llm_service.PAUSE_FILE.exists()


@pytest.mark.asyncio
async def test_stream_unsupported_downgrades_loudly_and_is_remembered():
    class NoStreaming:
        def __init__(self):
            self.astream_calls = 0
            self.ainvoke_calls = 0

        async def astream(self, *args, **kwargs):
            self.astream_calls += 1
            raise NotImplementedError("streaming is not supported")
            yield  # pragma: no cover

        async def ainvoke(self, *args, **kwargs):
            self.ainvoke_calls += 1
            return AIMessage(content="ok")

    base = NoStreaming()
    wrapper = RobustChatModelWrapper(base)

    first = await wrapper.complete([])
    assert first.content == "ok"
    assert wrapper._endpoint_key() in llm_service._NON_STREAMING_ENDPOINTS

    second = await wrapper.complete([])
    assert second.content == "ok"
    # The capability was cached: streaming is not probed again.
    assert base.astream_calls == 1
    assert base.ainvoke_calls == 2


@pytest.mark.asyncio
async def test_exhausted_retries_pause_then_raise_on_deadline():
    class AlwaysUnavailable:
        def __init__(self):
            self.calls = 0

        async def astream(self, *args, **kwargs):
            self.calls += 1
            raise RuntimeError("503 service unavailable")
            yield  # pragma: no cover

    wrapper = RobustChatModelWrapper(AlwaysUnavailable())
    with pytest.raises(LLMExhaustedError):
        await wrapper.complete([])
    # The pause file was written for the pause window, then cleaned up.
    assert not llm_service.PAUSE_FILE.exists()


@pytest.mark.asyncio
async def test_with_fallback_takes_over_without_pausing():
    class AlwaysUnavailable:
        async def astream(self, *args, **kwargs):
            raise RuntimeError("503 service unavailable")
            yield  # pragma: no cover

    wrapper = RobustChatModelWrapper(AlwaysUnavailable())

    async def fallback_call():
        return AIMessage(content="fallback answer")

    result = await with_fallback(
        main_call=lambda: wrapper.complete([]),
        fallback_call=fallback_call,
    )
    assert result.content == "fallback answer"
    # Because a fallback existed, exhaustion handed over instead of pausing.
    assert not llm_service.PAUSE_FILE.exists()


@pytest.mark.asyncio
async def test_with_fallback_refuses_to_hide_bad_requests():
    class CodedError(Exception):
        code = 400

    calls = {"fallback": 0}

    async def main_call():
        raise CodedError("malformed tool schema")

    async def fallback_call():
        calls["fallback"] += 1
        return "hidden bug"

    with pytest.raises(CodedError):
        await with_fallback(main_call=main_call, fallback_call=fallback_call)
    assert calls["fallback"] == 0


@pytest.mark.asyncio
async def test_acomplete_structured_corrective_reask():
    class CorrectsItself:
        def __init__(self):
            self.calls = []

        async def ainvoke(self, messages, **kwargs):
            self.calls.append(messages)
            if len(self.calls) == 1:
                return AIMessage(content="oops, here is prose instead of JSON")
            return AIMessage(content='{"a": 1}')

    base = CorrectsItself()
    result = await acomplete_structured(base, [])
    assert result == {"a": 1}
    assert len(base.calls) == 2
    correction_text = str(base.calls[1][-1].content)
    assert "could not be parsed" in correction_text


@pytest.mark.asyncio
async def test_acomplete_compat_path_for_plain_models():
    class PlainModel:
        async def astream(self, *args, **kwargs):
            yield AIMessageChunk(content="a")
            yield AIMessageChunk(content="b")

    result = await acomplete(PlainModel(), [])
    assert result.content == "ab"

    class InvokeOnly:
        async def ainvoke(self, *args, **kwargs):
            return AIMessage(content="direct")

    result = await acomplete(InvokeOnly(), [])
    assert result.content == "direct"


@pytest.mark.asyncio
async def test_mid_stream_failure_records_stream_reset_payload(monkeypatch):
    recorded_events = []

    def mock_record_event(name, payload, *, status="retrying"):
        recorded_events.append((name, payload, status))

    monkeypatch.setattr(llm_service, "_record_llm_event", mock_record_event)

    class BreaksMidStream:
        def __init__(self):
            self.calls = 0

        async def astream(self, *args, **kwargs):
            self.calls += 1
            if self.calls == 1:
                yield AIMessageChunk(content="Partial thoughts...")
                raise RuntimeError("503 high demand spike")
            yield AIMessageChunk(content="Recovered output")

    base = BreaksMidStream()
    wrapper = RobustChatModelWrapper(base)
    result = await wrapper.complete([])

    assert result.content == "Recovered output"
    resets = [e for e in recorded_events if e[0] == "llm_stream_reset"]
    assert len(resets) == 1
    event_name, payload, status = resets[0]
    assert event_name == "llm_stream_reset"
    assert payload["action"] == "discard"
    assert payload["reason"] == "mid_stream_failure"
    assert "stream_exec_id" in payload
    assert "lower API priority" in payload["message"]


@pytest.mark.asyncio
async def test_configured_retry_budget_is_shared_by_identical_fallback(monkeypatch):
    from unittest.mock import AsyncMock
    from artemis.llm.reliability import retry_policy_for
    from artemis.llm.router import ModelEndpoint

    monkeypatch.setattr(llm_service, "retry_policy_for", retry_policy_for)
    monkeypatch.setattr(llm_service.settings, "OPENAI_MAX_RETRIES", 5)
    monkeypatch.setattr(llm_service.settings, "OPENAI_RETRY_INTERVAL_SECONDS", 10)
    monkeypatch.setattr(llm_service, "_wait_for_resume", AsyncMock(return_value=False))
    sleep = AsyncMock()
    monkeypatch.setattr(llm_service.asyncio, "sleep", sleep)

    class Fails:
        def __init__(self):
            self.calls = 0

        async def ainvoke(self, *args, **kwargs):
            self.calls += 1
            if self.calls % 2:
                raise ConnectionError("connection reset")
            raise TimeoutError("timeout")

    main, fallback = Fails(), Fails()
    endpoint = ModelEndpoint(model_name="test", api_base="https://example.test/v1", max_retries=5)
    primary = RobustChatModelWrapper(main, endpoint=endpoint)
    secondary = RobustChatModelWrapper(fallback, endpoint=endpoint.model_copy())
    with pytest.raises(LLMExhaustedError):
        await with_fallback(lambda: primary.ainvoke([]), lambda: secondary.ainvoke([]))
    assert main.calls == 6
    assert fallback.calls == 0
    assert [call.args[0] for call in sleep.call_args_list] == [10] * 5
    assert llm_service._FAILED_ENDPOINTS.get() is None


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

    monkeypatch.setattr(llm_service, "_wait_for_resume", AsyncMock(return_value=False))
    cancelled = asyncio.Event()

    async def pending():
        try:
            await asyncio.Event().wait()
        finally:
            cancelled.set()

    with pytest.raises(LLMExhaustedError):
        await llm_service._run_with_recovery(pending, request_timeout=0.005, max_retries=0)
    assert cancelled.is_set()