"""并行节点和原生流事件的时序、提交与取消契约。"""

import asyncio
from contextlib import aclosing
import json

import httpx
import pytest

from src.agent.collect import CollectReplyError
from src.agent.graph import build_graph
from src.agent.llm import OpenAILLMService
from src.api import endpoints
from src.backends.chat import ChatInput, FormUpdate, ReplyFinished, TextDelta
from src.backends.langgraph import LangGraphBackend
from src.core.config import Config
from src.schemas.models import ProcessRequest_chat
from src.services.app_service import AccidentAppService
from src.storage.runtime import StorageSettings, open_storage
from src.services.repository import InMemorySessionRepository


class ParallelLLM:
    def __init__(self, patch=None):
        self.patch = patch if patch is not None else {'has_injury': False}
        self.extraction_started = asyncio.Event()
        self.reply_started = asyncio.Event()
        self.release_extraction = asyncio.Event()
        self.extraction_cancelled = asyncio.Event()
        self.reply_done = asyncio.Event()

    async def complete(self, system, user, **kwargs):
        if '【字段提取】' in system:
            self.extraction_started.set()
            try:
                await self.release_extraction.wait()
            except asyncio.CancelledError:
                self.extraction_cancelled.set()
                raise
            return json.dumps(self.patch)
        self.reply_started.set()
        # 若是串行实现，会卡在这里；两个原生图节点必须同时运行。
        await self.extraction_started.wait()
        self.reply_done.set()
        return '<state>1002</state>请问您还在事故现场吗？'

    async def stream_tokens(self, system, user, **kwargs):
        self.reply_started.set()
        await self.extraction_started.wait()
        yield '<state>1002</state>'
        yield '请问您还在事故现场吗？'
        self.reply_done.set()


async def service_for(llm):
    repository = InMemorySessionRepository()
    backend = LangGraphBackend(llm=llm)
    service = AccidentAppService(backend, repository)
    await service.chat(ChatInput('session', 'start'))
    return service, repository


@pytest.mark.asyncio
async def test_non_stream_nodes_fan_out_before_join_and_commit():
    llm = ParallelLLM()
    service, repo = await service_for(llm)
    task = asyncio.create_task(service.chat(ChatInput('session', '继续')))
    try:
        await asyncio.wait_for(llm.reply_done.wait(), 3)
        assert llm.extraction_started.is_set()
        assert not task.done()
        assert await repo.load('session') == {}
        llm.release_extraction.set()
        await asyncio.wait_for(task, 3)
        assert (await repo.load('session'))['ywrysw'] is False
    finally:
        llm.release_extraction.set()
        task.cancel()
        await asyncio.gather(task, return_exceptions=True)


@pytest.mark.asyncio
async def test_reply_finished_precedes_slow_extraction_and_commit():
    llm = ParallelLLM()
    service, repo = await service_for(llm)
    events = []
    async with aclosing(service.stream(ChatInput('session', '继续', need_form_update=True))) as source:
        while True:
            event = await asyncio.wait_for(anext(source), 3)
            events.append(event)
            if isinstance(event, ReplyFinished):
                break
        assert llm.reply_done.is_set()
        assert not llm.release_extraction.is_set()
        assert all(not isinstance(event, FormUpdate) for event in events)
        assert await repo.load('session') == {}
        llm.release_extraction.set()
        events.extend([event async for event in source])
    assert isinstance(events[-1], FormUpdate)
    assert (await repo.load('session'))['ywrysw'] is False
    assert not service._locks._entries


@pytest.mark.asyncio
@pytest.mark.parametrize('patch', [{}, {'has_injury': False}])
@pytest.mark.parametrize('visible', [False, True])
async def test_sse_flushes_final_sentence_before_extraction_with_unchanged_public_events(monkeypatch, patch, visible):
    llm = ParallelLLM(patch)
    service, repo = await service_for(llm)
    monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph')
    monkeypatch.setattr(endpoints, 'get_app_service', lambda: service)
    response = await endpoints.chat(
        ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续',
                            needFormUpdate=visible, useTextChunk=True),
        stream=True, backend=service._backend, idempotency_key=None,
    )
    received = []
    async with aclosing(response.body_iterator) as source:
        while True:
            event = await asyncio.wait_for(anext(source), 3)
            received.append(event)
            if 'event: text_delta' in event:
                break
        assert '请问您还在事故现场吗？' in received[-1]
        assert llm.extraction_started.is_set()
        assert not llm.release_extraction.is_set()
        assert not any('event: done' in event for event in received)
        assert await repo.load('session') == {}
        llm.release_extraction.set()
        received.extend([event async for event in source])
    body = ''.join(received)
    assert body.count('event: text_delta') == 1
    assert body.count('event: done') == 1
    assert ('event: formUpdate' in body) is (visible and bool(patch))
    assert 'reply_finished' not in body
    if not visible:
        assert 'has_injury' not in body


@pytest.mark.asyncio
async def test_closing_after_reply_finished_cancels_extraction_before_session_unlock():
    llm = ParallelLLM()
    service, repo = await service_for(llm)
    async with aclosing(service.stream(ChatInput('session', '继续'))) as source:
        while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished):
            pass
    assert llm.extraction_cancelled.is_set()
    assert not service._locks._entries
    assert await repo.load('session') == {}


@pytest.mark.asyncio
async def test_invalid_full_reply_never_emits_finished_or_commits_and_cancels_parallel_extraction():
    class Invalid(ParallelLLM):
        async def stream_tokens(self, *args, **kwargs):
            await self.extraction_started.wait()
            yield '<state>1002</state>partial'
            yield '<state>9999</state>'
    llm = Invalid()
    service, repo = await service_for(llm)
    events = []
    with pytest.raises(CollectReplyError):
        async for event in service.stream(ChatInput('session', '继续')):
            events.append(event)
    assert any(isinstance(event, TextDelta) for event in events)
    assert not any(isinstance(event, ReplyFinished) for event in events)
    assert llm.extraction_cancelled.is_set()
    assert await repo.load('session') == {}
    assert not service._locks._entries


@pytest.mark.asyncio
async def test_native_graph_nodes_are_visible_and_join_once():
    llm = ParallelLLM()
    llm.release_extraction.set()
    graph = build_graph()
    from src.agent.graph import GraphDeps
    names = []
    async for event in graph.astream_events(
        {'session_id': 'native', 'input_text': '继续', 'stage_code': '1002'},
        config={'configurable': {'thread_id': 'native'}}, context=GraphDeps(llm=llm), version='v2',
    ):
        if event['event'] == 'on_chain_end':
            names.append(event['name'])
    assert names.count('collect') == names.count('extract_accident_fields') == names.count('merge_collect') == 1


@pytest.mark.asyncio
async def test_confirmation_reply_finishes_before_its_dependent_extraction():
    class LLM:
        def __init__(self):
            self.release = asyncio.Event()
        async def complete(self, *args, **kwargs):
            await self.release.wait()
            return '{}'
        async def stream_tokens(self, *args, **kwargs):
            yield '<state>3002</state>请问您的姓名是什么？'
    llm = LLM()
    backend = LangGraphBackend(llm=llm)
    async with aclosing(backend.stream(ChatInput(
        'confirm', '是的', injected_fields={'stage_code': '3002'},
    ))) as source:
        received = []
        while True:
            event = await asyncio.wait_for(anext(source), 3)
            received.append(event)
            if isinstance(event, ReplyFinished):
                break
        assert ''.join(event.text for event in received if isinstance(event, TextDelta)) == '<state>3002</state>请问您的姓名是什么？'
        assert not llm.release.is_set()
        llm.release.set()
        assert [event async for event in source] == []


@pytest.mark.asyncio
async def test_sqlite_disconnect_after_reply_finishes_preserves_committed_snapshot(tmp_path):
    llm = ParallelLLM()
    settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'parallel.sqlite3'), auto_migrate=True)
    async with open_storage(settings) as storage:
        backend = LangGraphBackend(build_graph(storage.checkpointer), llm=llm)
        service = AccidentAppService(backend, storage.repository, locks=storage.locks)
        await service.chat(ChatInput('session', 'start'))
        async with aclosing(service.stream(ChatInput('session', '继续', request_key='interrupted'))) as source:
            while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished):
                pass
        assert llm.extraction_cancelled.is_set()
        assert not storage.repository.database.connection.in_transaction
        assert not storage.checkpointer.conn.in_transaction
        async with storage.repository.database.transaction() as conn:
            row = await conn.one('SELECT version,graph_state FROM agent_sessions WHERE session_id=?', ('session',))
            assert row['version'] == 1
            assert json.loads(row['graph_state'])['turn_count'] == 1
        llm.release_extraction.set()
        await service.chat(ChatInput('session', '继续', request_key='interrupted'))
        assert (await storage.repository.load('session'))['ywrysw'] is False
        assert not service._locks._entries


@pytest.mark.asyncio
async def test_commit_failure_after_early_reply_still_emits_error_without_form_update_or_done(monkeypatch):
    llm = ParallelLLM()
    service, repo = await service_for(llm)
    monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph')
    monkeypatch.setattr(endpoints, 'get_app_service', lambda: service)
    async def fail(*args):
        raise OSError('commit unavailable')
    monkeypatch.setattr(repo, 'apply_update', fail)
    response = await endpoints.chat(
        ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续',
                            needFormUpdate=True, useTextChunk=True),
        stream=True, backend=service._backend, idempotency_key=None,
    )
    received = []
    async with aclosing(response.body_iterator) as source:
        while True:
            event = await asyncio.wait_for(anext(source), 3)
            received.append(event)
            if 'event: text_delta' in event:
                break
        llm.release_extraction.set()
        received.extend([event async for event in source])
    body = ''.join(received)
    assert 'event: text_delta' in body and 'event: error' in body
    assert 'event: done' not in body and 'event: formUpdate' not in body
    assert await repo.load('session') == {}


@pytest.mark.asyncio
async def test_real_chatopenai_native_events_do_not_leak_extraction_or_force_streaming(monkeypatch):
    requests = []
    def handler(request):
        payload = json.loads(request.content)
        requests.append(payload)
        extraction = '【字段提取】' in payload['messages'][0]['content']
        if extraction:
            assert payload['stream'] is False
            return httpx.Response(200, json={
                'id': 'extract', 'object': 'chat.completion', 'created': 0, 'model': 'probe',
                'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': '{"has_injury":false}'}, 'finish_reason': 'stop'}],
            })
        assert payload['stream'] is True
        data = {'id': 'reply', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'probe',
                'choices': [{'index': 0, 'delta': {'content': '<state>1002</state>请问您还在事故现场吗？'}, 'finish_reason': None}]}
        return httpx.Response(200, headers={'content-type': 'text/event-stream'},
                              content='data: '+json.dumps(data)+'\n\ndata: [DONE]\n\n')
    monkeypatch.setattr(Config, 'LLM_HTTP_TRUST_ENV', False)
    llm = OpenAILLMService('probe', 'offline', base_url='http://probe.local/v1')
    llm._get_client().http_async_client._transport = httpx.MockTransport(handler)
    try:
        service, repo = await service_for(llm)
        events = [event async for event in service.stream(ChatInput('session', '继续', need_form_update=True))]
        text = ''.join(event.text for event in events if isinstance(event, TextDelta))
        assert text == '<state>1002</state>请问您还在事故现场吗？'
        assert len(requests) == 2
        assert sum(isinstance(event, ReplyFinished) for event in events) == 1
        assert isinstance(events[-1], FormUpdate)
    finally:
        await llm.aclose()