import asyncio
from contextlib import aclosing
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest

from src.agent.collect import validated_reply
from src.agent.llm import LLMTimeoutError, OpenAILLMService
from src.backends.chat import ChatInput, TextDelta
from src.backends.langgraph import LangGraphBackend
from src.services.app_service import AccidentAppService
from src.services.repository import InMemorySessionRepository


class Graph:
    async def ainvoke(self, data, *, config, context):
        raw, _ = await validated_reply(context.llm, 'system', 'user', '1002')
        return {'response_text': raw, 'turn_patch': {}}


@pytest.mark.asyncio
async def test_valid_prefix_reaches_client_before_model_finishes():
    release = asyncio.Event()
    finished = asyncio.Event()
    class LLM:
        async def stream_tokens(self, *args, **kwargs):
            yield '<sta'
            yield 'te>1002</state>'
            await release.wait()
            yield 'tail'
            finished.set()
    backend = LangGraphBackend(graph=Graph(), llm=LLM())
    async with aclosing(backend.stream(ChatInput('s', 'hello'))) as events:
        first = await asyncio.wait_for(anext(events), 1)
        assert first == TextDelta('<state>1002</state>')
        assert not finished.is_set()
        release.set()
        remaining = [e async for e in events]
    assert ''.join(e.text for e in remaining if isinstance(e, TextDelta)) == 'tail'
    assert finished.is_set()


@pytest.mark.asyncio
async def test_closing_application_stream_cancels_model_before_unlocking():
    cancelled = asyncio.Event()
    class LLM:
        async def stream_tokens(self, *args, **kwargs):
            yield '<state>1002</state>'
            try:
                await asyncio.Future()
            finally:
                cancelled.set()
    service = AccidentAppService(LangGraphBackend(graph=Graph(), llm=LLM()), InMemorySessionRepository())
    async with aclosing(service.stream(ChatInput('s', 'hello'))) as events:
        await asyncio.wait_for(anext(events), 1)
    assert cancelled.is_set()
    assert not service._locks._entries
    assert await service.set_info('s', 'hphm1', 'test') == ('200', '')


@pytest.mark.asyncio
@pytest.mark.parametrize('stream', [False, True])
async def test_timeout_covers_whole_attempt_and_retries_are_bounded(stream):
    responses = []
    class HangingResponse:
        closed = False
        def __aiter__(self): return self.generate()
        async def generate(self):
            await asyncio.Future()
            yield None
        async def aclose(self): self.closed = True
    async def invoke(*args, **kwargs):
        await asyncio.Future()
    def create_stream(*args, **kwargs):
        response = HangingResponse()
        responses.append(response)
        return response
    llm = OpenAILLMService(model='m', api_key='audit', timeout_seconds=0.01, max_retries=2)
    call = Mock(side_effect=create_stream) if stream else AsyncMock(side_effect=invoke)
    llm._client = SimpleNamespace(astream=call) if stream else SimpleNamespace(ainvoke=call)
    llm._retry_delay = AsyncMock()
    # 先加载 SDK 错误类；外层看门狗不测模块惰性导入耗时。
    assert llm._retryable(TimeoutError())
    assert isinstance(llm._error(TimeoutError()), LLMTimeoutError)
    async with asyncio.timeout(10):
        with pytest.raises(LLMTimeoutError):
            if stream:
                async for _ in llm.stream_tokens('system', 'user'): pass
            else:
                await llm.complete('system', 'user')
    assert call.call_count == 3
    assert llm._retry_delay.await_count == 2
    assert all(response.closed for response in responses)