"""通过真实 ChatOpenAI/OpenAI SDK + HTTP mock 验证迁移边界。"""

import asyncio
import json
from contextlib import aclosing
from typing import TypedDict
from unittest.mock import AsyncMock

import httpx
import pytest
from langchain_core.callbacks import AsyncCallbackHandler
from langchain_core.messages import AIMessage
from langgraph.graph import END, START, StateGraph

from src.agent.llm import LLMAPIError, LLMTimeoutError, OpenAILLMService
from src.core.config import Config


def _chunk(text):
    data = {
        'id': 'response', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'probe',
        'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
    }
    return ('data: ' + json.dumps(data) + '\n\n').encode()


def _response(text, *, stream=False):
    if stream:
        return httpx.Response(200, headers={'content-type': 'text/event-stream'},
                              content=_chunk(text) + b'data: [DONE]\n\n')
    return httpx.Response(200, json={
        'id': 'response', 'object': 'chat.completion', 'created': 0, 'model': 'probe',
        'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': text},
                     'finish_reason': 'stop'}],
        'usage': {'prompt_tokens': 7, 'completion_tokens': 2, 'total_tokens': 9},
    })


def _service(monkeypatch, handler, **kwargs):
    monkeypatch.setattr(Config, 'LLM_HTTP_TRUST_ENV', False)
    service = OpenAILLMService(
        model='probe', api_key='offline-test', base_url='http://probe.local/v1',
        extra_body={'reasoning_effort': 'none'}, **kwargs,
    )
    # 只替换 HTTP 传输，模型工厂、消息转换、callbacks 和 SDK 全部真实执行。
    model = service._get_client()
    model.http_async_client._transport = httpx.MockTransport(handler)
    return service


@pytest.mark.asyncio
@pytest.mark.parametrize('stream', [False, True])
@pytest.mark.parametrize('temperature', [None, 0.0])
async def test_wire_contract_preserves_profile_messages_and_chat_completions(monkeypatch, stream, temperature):
    monkeypatch.setenv('OPENAI_PROXY', 'http://invalid-proxy.local:1')
    monkeypatch.setenv('HTTP_PROXY', 'http://invalid-proxy.local:1')
    monkeypatch.setenv('HTTPS_PROXY', 'http://invalid-proxy.local:1')
    requests = []
    def handler(request):
        requests.append(request)
        return _response('收到', stream=stream)
    llm = _service(monkeypatch, handler, temperature=temperature)
    http_async = llm._client.http_async_client
    http_sync = llm._client.http_client
    history = [{'role': 'user', 'content': '之前的输入'}, {'role': 'assistant', 'content': '之前的回复'}]
    try:
        if stream:
            result = ''.join([t async for t in llm.stream_tokens('系统', '本轮', history=history)])
        else:
            result = await llm.complete('系统', '本轮', history=history)
        assert result == '收到'
        assert len(requests) == 1
        request = requests[0]
        payload = json.loads(request.content)
        assert request.url.path == '/v1/chat/completions'
        assert payload['model'] == 'probe'
        assert payload['reasoning_effort'] == 'none'
        assert payload['stream'] is stream
        assert 'stream_options' not in payload
        if temperature is None:
            assert 'temperature' not in payload
        else:
            assert payload['temperature'] == temperature
        assert payload['messages'] == [
            {'role': 'system', 'content': '系统'}, *history, {'role': 'user', 'content': '本轮'},
        ]
        assert history == [{'role': 'user', 'content': '之前的输入'}, {'role': 'assistant', 'content': '之前的回复'}]
    finally:
        await llm.aclose()
    assert http_async.is_closed and http_sync.is_closed
    assert llm._client is None
    await llm.aclose()


class TrackingBody(httpx.AsyncByteStream):
    def __init__(self, text='partial', *, fail=False, emit=True):
        self.text = text
        self.fail = fail
        self.emit = emit
        self.closed = False
        self.started = asyncio.Event()
        self.release = asyncio.Event()

    async def __aiter__(self):
        self.started.set()
        if self.emit:
            yield _chunk(self.text)
        if self.fail:
            raise httpx.ReadError('broken transport')
        await self.release.wait()
        yield _chunk('tail')
        yield b'data: [DONE]\n\n'

    async def aclose(self):
        self.closed = True


def _stream_response(body):
    return httpx.Response(200, headers={'content-type': 'text/event-stream'}, stream=body)


@pytest.mark.asyncio
async def test_early_close_is_immediate_and_does_not_close_another_session(monkeypatch):
    bodies = {}
    def handler(request):
        text = json.loads(request.content)['messages'][-1]['content']
        body = bodies[text] = TrackingBody(text)
        return _stream_response(body)
    llm = _service(monkeypatch, handler)
    first = llm.stream_tokens('system', 'first')
    second = llm.stream_tokens('system', 'second')
    try:
        # anext 运行在其他任务中，aclose 在当前任务中；上下文也必须安全复原。
        assert await asyncio.wait_for(anext(first), 5) == 'first'
        assert await asyncio.wait_for(anext(second), 5) == 'second'
        await first.aclose()
        assert bodies['first'].closed
        assert not bodies['second'].closed
        bodies['second'].release.set()
        assert [t async for t in second] == ['tail']
        assert bodies['second'].closed
    finally:
        await first.aclose()
        await second.aclose()
        await llm.aclose()


@pytest.mark.asyncio
async def test_transport_failure_after_text_closes_immediately_without_retry(monkeypatch):
    body = TrackingBody(fail=True)
    calls = []
    def handler(request):
        calls.append(request)
        return _stream_response(body)
    llm = _service(monkeypatch, handler)
    llm._retry_delay = AsyncMock()
    tokens = []
    try:
        with pytest.raises(LLMAPIError):
            async for token in llm.stream_tokens('system', 'user'):
                tokens.append(token)
        assert tokens == ['partial']
        assert body.closed
        assert len(calls) == 1
        llm._retry_delay.assert_not_awaited()
    finally:
        await llm.aclose()


@pytest.mark.asyncio
async def test_stream_cancellation_closes_http_response_before_returning(monkeypatch):
    body = TrackingBody(emit=False)
    llm = _service(monkeypatch, lambda request: _stream_response(body))
    llm._retry_delay = AsyncMock()
    async def consume():
        async with aclosing(llm.stream_tokens('system', 'user')) as stream:
            async for _ in stream:
                pass
    task = asyncio.create_task(consume())
    try:
        await asyncio.wait_for(body.started.wait(), 5)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert body.closed
        llm._retry_delay.assert_not_awaited()
    finally:
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)
        await llm.aclose()


@pytest.mark.asyncio
async def test_stream_timeout_closes_each_attempt_before_retry(monkeypatch):
    bodies = []
    def handler(request):
        assert all(body.closed for body in bodies)
        body = TrackingBody(emit=False)
        bodies.append(body)
        return _stream_response(body)
    llm = _service(monkeypatch, handler, timeout_seconds=0.1, max_retries=2)
    llm._retry_delay = AsyncMock()
    try:
        with pytest.raises(LLMTimeoutError):
            async for _ in llm.stream_tokens('system', 'user'):
                pass
        assert len(bodies) == 3
        assert all(body.closed for body in bodies)
        assert llm._retry_delay.await_count == 2
    finally:
        await llm.aclose()


class TraceRecorder(AsyncCallbackHandler):
    def __init__(self):
        self.chains = {}
        self.models = []
        self.results = []

    async def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
        self.chains[run_id] = kwargs.get('name')

    async def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
        self.models.append({'run_id': run_id, **kwargs})

    async def on_llm_end(self, response, **kwargs):
        self.results.append(response)


@pytest.mark.asyncio
@pytest.mark.parametrize('stream', [False, True])
async def test_model_callbacks_inherit_graph_parent_metadata_and_tags(monkeypatch, stream):
    llm = _service(monkeypatch, lambda request: _response('ok', stream=stream))
    recorder = TraceRecorder()
    class State(TypedDict):
        text: str
    async def dialogue(state):
        if stream:
            text = ''.join([t async for t in llm.stream_tokens('system', state['text'])])
        else:
            text = await llm.complete('system', state['text'])
        return {'text': text}
    graph = StateGraph(State)
    graph.add_node('dialogue', dialogue)
    graph.add_edge(START, 'dialogue')
    graph.add_edge('dialogue', END)
    try:
        result = await graph.compile().ainvoke({'text': 'hello'}, config={
            'callbacks': [recorder], 'metadata': {'session_id': 'trace-session'}, 'tags': ['probe'],
        })
        assert result['text'] == 'ok'
        assert len(recorder.models) == 1
        model = recorder.models[0]
        assert recorder.chains[model['parent_run_id']] == 'dialogue'
        assert model['name'] == ('llm.stream' if stream else 'llm.complete')
        assert model['metadata']['session_id'] == 'trace-session'
        assert model['metadata']['llm_attempt'] == 1
        assert 'probe' in model['tags']
        assert len(recorder.results) == 1
        if not stream:
            usage = recorder.results[0].generations[0][0].message.usage_metadata
            assert usage['input_tokens'] == 7 and usage['output_tokens'] == 2
    finally:
        await llm.aclose()


@pytest.mark.parametrize('content, expected', [
    ('正文', '正文'),
    ([{'type': 'reasoning', 'reasoning': '内部内容'}, {'type': 'text', 'text': '正文'}], '正文'),
    ([{'type': 'reasoning', 'reasoning': '内部内容'}], ''),
])
def test_only_text_content_is_forwarded(content, expected):
    assert OpenAILLMService._text(AIMessage(content=content)) == expected


@pytest.mark.asyncio
async def test_shared_model_serves_concurrent_calls_without_mixing_messages(monkeypatch):
    async def handler(request):
        payload = json.loads(request.content)
        await asyncio.sleep(0)
        return _response(payload['messages'][-1]['content'], stream=payload['stream'])
    llm = _service(monkeypatch, handler)
    async def request(index):
        text = f'input-{index}'
        if index % 2:
            return ''.join([t async for t in llm.stream_tokens('system', text)])
        return await llm.complete('system', text)
    try:
        assert await asyncio.gather(*(request(i) for i in range(12))) == [f'input-{i}' for i in range(12)]
    finally:
        await llm.aclose()