import asyncio

import pytest
import httpx
from msgspec import DecodeError, Struct
from openai import APIError, OpenAIError
from types import SimpleNamespace

from aceai.llm.errors import (
    AceAIConfigurationError,
    AceAIRuntimeError,
    AceAIValidationError,
    LLMContextWindowExceededError,
    LLMProviderError,
)
from aceai.llm.models import (
    LLMInput,
    LLMMessage,
    LLMMessagePart,
    LLMProviderBase,
    LLMProviderModality,
    LLMResponse,
    LLMStreamEvent,
)
from aceai.llm.service import LLMService, LLM_RETRY_EXHAUSTED_MESSAGE


class AttrMessagePart(dict):
    @property
    def type(self):
        return self["type"]


class Payload(Struct):
    value: int


class RecordingProvider(LLMProviderBase):
    def __init__(
        self,
        *,
        responses: list[LLMResponse] | None = None,
        stream_events: list[LLMStreamEvent] | None = None,
        default_model: str = "gpt-4o",
        default_stream_model: str = "gpt-4o-mini",
        modality: LLMProviderModality | None = None,
    ) -> None:
        self._responses = list(responses or [LLMResponse(text="ok")])
        self._stream_events = list(stream_events or [])
        self._default_model = default_model
        self._default_stream_model = default_stream_model
        self._modality = modality or LLMProviderModality()
        self.complete_requests: list[LLMInput] = []
        self.stream_requests: list[LLMInput] = []

    async def complete(self, request: LLMInput) -> LLMResponse:
        self.complete_requests.append(request)
        return self._responses.pop(0)

    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            for event in self._stream_events:
                yield event

        return iterator()

    @property
    def default_model(self) -> str:
        return self._default_model

    @property
    def default_stream_model(self) -> str:
        return self._default_stream_model

    @property
    def modality(self) -> LLMProviderModality:
        return self._modality

    async def stt(
        self, filename, file, *, model: str, prompt: str | None = None
    ) -> str:  # pragma: no cover - not used
        return "transcript"


class ErroringProvider(LLMProviderBase):
    def __init__(self) -> None:
        self._default_model = "gpt-4o"
        self._default_stream_model = "gpt-4o-mini"

    async def complete(self, request: LLMInput) -> LLMResponse:
        raise OpenAIError("boom")

    def stream(self, request: LLMInput):
        raise AceAIRuntimeError("unused stream")

    @property
    def default_model(self) -> str:
        return self._default_model

    @property
    def default_stream_model(self) -> str:
        return self._default_stream_model

    async def stt(
        self, filename, file, *, model: str, prompt: str | None = None
    ) -> str:  # pragma: no cover - not used
        return ""


class FlakyCompleteProvider(RecordingProvider):
    def __init__(self) -> None:
        super().__init__(responses=[LLMResponse(text="recovered")])
        self.failures_remaining = 1

    async def complete(self, request: LLMInput) -> LLMResponse:
        self.complete_requests.append(request)
        if self.failures_remaining > 0:
            self.failures_remaining -= 1
            raise httpx.ConnectError("connection dropped")
        return self._responses.pop(0)


class FlakyStreamProvider(RecordingProvider):
    def __init__(self) -> None:
        super().__init__()
        self.failures_remaining = 1

    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            if self.failures_remaining > 0:
                self.failures_remaining -= 1
                raise httpx.RemoteProtocolError(
                    "peer closed connection without sending complete message body"
                )
            yield LLMStreamEvent(
                event_type="response.completed",
                response=LLMResponse(text="recovered"),
            )

        return iterator()


class FlakyAPIErrorStreamProvider(RecordingProvider):
    def __init__(self) -> None:
        super().__init__()
        self.failures_remaining = 1

    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            if self.failures_remaining > 0:
                self.failures_remaining -= 1
                raise APIError(
                    "Our servers are currently overloaded. Please try again later.",
                    request=httpx.Request(
                        "POST",
                        "https://api.openai.com/v1/responses",
                    ),
                    body={
                        "type": "error",
                        "error": {
                            "type": "service_unavailable_error",
                            "code": "server_is_overloaded",
                            "message": (
                                "Our servers are currently overloaded. "
                                "Please try again later."
                            ),
                        },
                    },
                )
            yield LLMStreamEvent(
                event_type="response.completed",
                response=LLMResponse(text="recovered"),
            )

        return iterator()


class ContextWindowCompleteProvider(RecordingProvider):
    async def complete(self, request: LLMInput) -> LLMResponse:
        self.complete_requests.append(request)
        raise APIError(
            "Your input exceeds the context window of this model. "
            "Please adjust your input and try again.",
            request=httpx.Request(
                "POST",
                "https://api.openai.com/v1/responses",
            ),
            body={
                "error": {
                    "code": "context_length_exceeded",
                    "message": "Your input exceeds the context window of this model.",
                },
            },
        )


class ContextWindowStreamProvider(RecordingProvider):
    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            raise APIError(
                "Your input exceeds the context window of this model. "
                "Please adjust your input and try again.",
                request=httpx.Request(
                    "POST",
                    "https://api.openai.com/v1/responses",
                ),
                body={
                    "error": {
                        "code": "context_length_exceeded",
                        "message": "Your input exceeds the context window of this model.",
                    },
                },
            )
            yield LLMStreamEvent(event_type="response.completed")

        return iterator()


class ContextWindowStreamErrorEventProvider(RecordingProvider):
    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            yield LLMStreamEvent(
                event_type="response.error",
                error=(
                    "Your input exceeds the context window of this model. "
                    "Please adjust your input and try again."
                ),
            )

        return iterator()


class ContextWindowFailedCompletionProvider(RecordingProvider):
    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            yield LLMStreamEvent(
                event_type="response.completed",
                response=LLMResponse(
                    text=(
                        "APIError: Your input exceeds the context window of this model. "
                        "Please adjust your input and try again."
                    ),
                    status="failed",
                ),
            )

        return iterator()


class HangingStreamProvider(RecordingProvider):
    def __init__(self) -> None:
        super().__init__()
        self.hangs_remaining = 1

    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            if self.hangs_remaining > 0:
                self.hangs_remaining -= 1
                await asyncio.Event().wait()
            yield LLMStreamEvent(
                event_type="response.completed",
                response=LLMResponse(text="recovered"),
            )

        return iterator()


class SlowAfterFirstStreamProvider(RecordingProvider):
    def __init__(self) -> None:
        super().__init__()
        self.attempts = 0

    def stream(self, request: LLMInput):
        self.stream_requests.append(request)
        self.attempts += 1
        attempt = self.attempts

        async def iterator():
            if attempt == 1:
                yield LLMStreamEvent(
                    event_type="response.output_text.delta",
                    text_delta="hello",
                )
                await asyncio.Event().wait()
            yield LLMStreamEvent(
                event_type="response.completed",
                response=LLMResponse(text="recovered"),
            )

        return iterator()


class AlwaysFailingStreamProvider(RecordingProvider):
    def stream(self, request: LLMInput):
        self.stream_requests.append(request)

        async def iterator():
            raise httpx.RemoteProtocolError("connection still down")
            yield LLMStreamEvent(event_type="response.completed")

        return iterator()


def test_llm_service_requires_providers() -> None:
    with pytest.raises(AceAIConfigurationError):
        LLMService([], timeout_seconds=1.0)


def test_llm_service_defaults_distinguish_stream_start_and_event_timeouts() -> None:
    service = LLMService([RecordingProvider()], timeout_seconds=1.0)

    assert service._stream_start_timeout_seconds == 30.0
    assert service._stream_event_timeout_seconds == 15.0


def test_llm_service_rotation_and_counts() -> None:
    providers = [RecordingProvider(), RecordingProvider()]
    service = LLMService(providers, timeout_seconds=1.0)

    assert service.get_provider_count() == 2
    assert service.has_provider is True
    assert service._get_current_provider() is providers[0]

    service._rotate_provider()

    assert service._get_current_provider() is providers[1]


@pytest.mark.anyio
async def test_llm_service_complete_passes_metadata_through() -> None:
    provider = RecordingProvider(responses=[LLMResponse(text="ok")])
    service = LLMService([provider], timeout_seconds=1.0)
    metadata = {"model": provider.default_model}

    await service.complete(
        messages=[LLMMessage.build("system", "Prompt")],
        metadata=metadata,
    )

    assert provider.complete_requests[0]["metadata"] is metadata


@pytest.mark.anyio
async def test_llm_service_complete_updates_last_response() -> None:
    provider = RecordingProvider(responses=[LLMResponse(text="fresh")])
    service = LLMService([provider], timeout_seconds=1.0)

    response = await service.complete(messages=[LLMMessage.build("system", "Prompt")])

    assert response.text == "fresh"


@pytest.mark.anyio
async def test_llm_service_complete_retries_retryable_provider_errors() -> None:
    provider = FlakyCompleteProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    response = await service.complete(messages=[LLMMessage.build("system", "Prompt")])

    assert response.text == "recovered"
    assert len(provider.complete_requests) == 2


@pytest.mark.anyio
async def test_llm_service_complete_does_not_retry_context_window_errors() -> None:
    provider = ContextWindowCompleteProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    with pytest.raises(LLMContextWindowExceededError):
        await service.complete(messages=[LLMMessage.build("system", "Prompt")])

    assert len(provider.complete_requests) == 1


@pytest.mark.anyio
async def test_llm_service_stream_preserves_request_metadata() -> None:
    event = LLMStreamEvent(
        event_type="response.output_text.delta",
        text_delta="hello",
    )
    provider = RecordingProvider(stream_events=[event])
    service = LLMService([provider], timeout_seconds=1.0)
    metadata = {"model": provider.default_stream_model}

    received = []
    async for part in service.stream(
        messages=[LLMMessage.build("system", "s")],
        metadata=metadata,
    ):
        received.append(part.text_delta)

    assert received == ["hello"]
    assert provider.stream_requests[0]["metadata"] is metadata


@pytest.mark.anyio
async def test_llm_service_stream_emits_retry_progress() -> None:
    provider = FlakyStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    events: list[LLMStreamEvent] = []
    async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
        events.append(event)

    assert [event.event_type for event in events] == [
        "response.retrying",
        "response.completed",
    ]
    retry_event = events[0]
    assert retry_event.retry_count == 1
    assert retry_event.retry_max == 5
    assert retry_event.retry_delay_seconds == 0.0
    assert retry_event.error == (
        "RemoteProtocolError: peer closed connection without sending complete message body"
    )
    assert len(provider.stream_requests) == 2


@pytest.mark.anyio
async def test_llm_service_stream_retries_openai_api_error_events() -> None:
    provider = FlakyAPIErrorStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    events: list[LLMStreamEvent] = []
    async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
        events.append(event)

    assert [event.event_type for event in events] == [
        "response.retrying",
        "response.completed",
    ]
    assert events[0].error == (
        "APIError: Our servers are currently overloaded. Please try again later."
    )
    assert len(provider.stream_requests) == 2


@pytest.mark.anyio
async def test_llm_service_stream_does_not_retry_context_window_errors() -> None:
    provider = ContextWindowStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    with pytest.raises(LLMContextWindowExceededError):
        async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
            raise AssertionError(f"unexpected stream event: {event.event_type}")

    assert len(provider.stream_requests) == 1


@pytest.mark.anyio
async def test_llm_service_stream_maps_context_window_error_events() -> None:
    provider = ContextWindowStreamErrorEventProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    with pytest.raises(LLMContextWindowExceededError):
        async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
            raise AssertionError(f"unexpected stream event: {event.event_type}")

    assert len(provider.stream_requests) == 1


@pytest.mark.anyio
async def test_llm_service_stream_maps_failed_context_window_completion() -> None:
    provider = ContextWindowFailedCompletionProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
    )

    with pytest.raises(LLMContextWindowExceededError):
        async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
            raise AssertionError(f"unexpected stream event: {event.event_type}")

    assert len(provider.stream_requests) == 1


@pytest.mark.anyio
async def test_llm_service_stream_retries_when_next_event_times_out() -> None:
    provider = HangingStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
        stream_start_timeout_seconds=0.01,
        stream_event_timeout_seconds=0.01,
    )

    events: list[LLMStreamEvent] = []
    async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
        events.append(event)

    assert [event.event_type for event in events] == [
        "response.retrying",
        "response.completed",
    ]
    retry_event = events[0]
    assert retry_event.retry_count == 1
    assert retry_event.error == "TimeoutError: "
    assert len(provider.stream_requests) == 2


@pytest.mark.anyio
async def test_llm_service_stream_uses_event_timeout_after_first_event() -> None:
    provider = SlowAfterFirstStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        retry_initial_delay_seconds=0.0,
        stream_start_timeout_seconds=10.0,
        stream_event_timeout_seconds=0.01,
    )

    events: list[LLMStreamEvent] = []
    async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
        events.append(event)

    assert [event.event_type for event in events] == [
        "response.output_text.delta",
        "response.retrying",
        "response.completed",
    ]
    retry_event = events[1]
    assert retry_event.retry_count == 1
    assert retry_event.error == "TimeoutError: "
    assert len(provider.stream_requests) == 2


@pytest.mark.anyio
async def test_llm_service_stream_completes_with_user_message_after_retries_exhausted() -> (
    None
):
    provider = AlwaysFailingStreamProvider()
    service = LLMService(
        [provider],
        timeout_seconds=1.0,
        max_retries=2,
        retry_initial_delay_seconds=0.0,
    )

    events: list[LLMStreamEvent] = []
    async for event in service.stream(messages=[LLMMessage.build("system", "s")]):
        events.append(event)

    assert [event.event_type for event in events] == [
        "response.retrying",
        "response.retrying",
        "response.completed",
    ]
    final_response = events[-1].response
    assert isinstance(final_response, LLMResponse)
    assert final_response.text == LLM_RETRY_EXHAUSTED_MESSAGE
    assert final_response.status == "failed"
    assert len(provider.stream_requests) == 3


@pytest.mark.anyio
async def test_complete_json_validates_message_structure() -> None:
    service = LLMService([RecordingProvider()], timeout_seconds=1.0)

    with pytest.raises(AceAIValidationError):
        await service.complete_json(schema=Payload)

    with pytest.raises(AceAIValidationError):
        await service.complete_json(schema=Payload, messages=[])

    with pytest.raises(AceAIValidationError):
        await service.complete_json(
            schema=Payload,
            messages=[LLMMessage.build("user", "hi")],
        )


@pytest.mark.anyio
async def test_complete_json_inserts_hint_and_retries_on_error() -> None:
    responses = [
        LLMResponse(text="not json"),
        LLMResponse(text='{"value":7}'),
    ]
    provider = RecordingProvider(responses=responses)
    service = LLMService([provider], timeout_seconds=1.0, max_retries=2)
    messages = [LLMMessage.build("system", "start")]

    payload = await service.complete_json(
        schema=Payload,
        retries=2,
        messages=messages,
    )

    assert payload.value == 7
    assert len(provider.complete_requests) == 2
    schema_hint = provider.complete_requests[0]["messages"][1]
    assert "Return Format Advisory" in schema_hint.content[0]["data"]
    error_notice = provider.complete_requests[1]["messages"][-1]
    assert error_notice.role == "system"
    assert "Error handling notice" in error_notice.content[0]["data"]


@pytest.mark.anyio
async def test_complete_passes_image_parts_through_without_validation() -> None:
    provider = RecordingProvider(
        modality=LLMProviderModality(image_in=False),
        responses=[LLMResponse(text="ok")],
    )
    service = LLMService([provider], timeout_seconds=1.0)

    image_part = LLMMessagePart(type="image", binary=b"bytes", mime_type="image/png")
    message = LLMMessage.build("user", [image_part])

    response = await service.complete(messages=[message])

    assert response.text == "ok"
    forwarded = provider.complete_requests[0]["messages"][0].content[0]
    assert forwarded["type"] == "image"


def test_validate_messages_requires_messages() -> None:
    service = LLMService([RecordingProvider()], timeout_seconds=1.0)
    with pytest.raises(ValueError):
        service._validate_messages({"messages": []})


def test_validate_messages_requires_list_content() -> None:
    service = LLMService([RecordingProvider()], timeout_seconds=1.0)
    bad_message = SimpleNamespace(role="user", content="oops")
    with pytest.raises(TypeError):
        service._validate_messages({"messages": [bad_message]})


@pytest.mark.parametrize(
    ("part", "modality"),
    [
        (LLMMessagePart(type="text", data="t"), LLMProviderModality(text_in=False)),
        (
            LLMMessagePart(type="image", binary=b"raw"),
            LLMProviderModality(image_in=False),
        ),
        (
            LLMMessagePart(type="audio", binary=b"pcm"),
            LLMProviderModality(audio_in=False),
        ),
        (
            LLMMessagePart(type="file", binary=b"blob"),
            LLMProviderModality(file_in=False),
        ),
    ],
)
def test_validate_messages_enforces_modality_support(
    part: LLMMessagePart, modality: LLMProviderModality
) -> None:
    provider = RecordingProvider(modality=modality)
    service = LLMService([provider], timeout_seconds=1.0)
    message = LLMMessage.build("user", [part])

    with pytest.raises(LLMProviderError):
        service._validate_messages({"messages": [message]})


def test_validate_messages_rejects_unknown_part_type() -> None:
    provider = RecordingProvider()
    service = LLMService([provider], timeout_seconds=1.0)
    message = LLMMessage(
        role="user",
        content=[AttrMessagePart(type="binary", data="payload")],
    )

    with pytest.raises(ValueError):
        service._validate_messages({"messages": [message]})


@pytest.mark.anyio
async def test_complete_json_with_retry_raises_after_failures() -> None:
    responses = [
        LLMResponse(text="nah"),
        LLMResponse(text="still nah"),
    ]
    provider = RecordingProvider(responses=responses)
    service = LLMService([provider], timeout_seconds=1.0)
    messages = [LLMMessage.build("system", "start")]
    llm_input = {"messages": messages}

    with pytest.raises(DecodeError):
        await service._complete_json_with_retry(
            schema=Payload,
            retries=2,
            llm_input=llm_input,
        )


@pytest.mark.anyio
async def test_complete_json_without_retries_short_circuits() -> None:
    provider = RecordingProvider(responses=[LLMResponse(text='{"value":3}')])
    service = LLMService([provider], timeout_seconds=1.0)
    messages = [LLMMessage.build("system", "start")]

    payload = await service.complete_json(
        schema=Payload,
        retries=0,
        messages=messages,
    )

    assert payload.value == 3
