import asyncio
import json
import os
import subprocess
import sys
from pathlib import Path
from contextlib import aclosing, asynccontextmanager
from dataclasses import replace
from uuid import uuid4

import pytest

from src.agent.graph import build_graph
from src.backends.chat import ChatBackendConflictError, ChatInput, FormUpdate
from src.backends.langgraph import LangGraphBackend
from src.services.app_service import AccidentAppService
from src.storage.runtime import StorageSettings, open_storage


def run(coroutine):
    # psycopg 在 Windows 要求 SelectorEventLoop；Linux 使用标准 runner。
    factory = asyncio.SelectorEventLoop if os.name == 'nt' else None
    with asyncio.Runner(loop_factory=factory) as runner:
        return runner.run(coroutine)


@pytest.fixture(params=['sqlite', 'postgresql'])
def settings(request, tmp_path):
    if request.param == 'postgresql':
        dsn = os.getenv('AGENT_TEST_POSTGRES_DSN')
        if not dsn:
            pytest.skip('Set AGENT_TEST_POSTGRES_DSN to an isolated PostgreSQL database')
        return StorageSettings(kind='postgresql', postgres_dsn=dsn, auto_migrate=True, pool_size=2)
    return StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'agent.sqlite3'), auto_migrate=True)


class RecordingLLM:
    def __init__(self):
        self.histories = []

    async def complete(self, system, user, *, history=None):
        if '【字段提取】' in system:
            return '{"has_injury":false,"vehicle_count":0}'
        self.histories.append(history or [])
        return '<state>1002</state>已记录'

    async def stream_tokens(self, system, user, *, history=None):
        self.histories.append(history or [])
        yield '<state>1002</state>'
        yield '已记录'


@asynccontextmanager
async def service_for(settings, llm=None):
    async with open_storage(settings) as storage:
        backend = LangGraphBackend(build_graph(storage.checkpointer), llm=llm or RecordingLLM())
        yield AccidentAppService(backend, storage.repository, locks=storage.locks), storage


async def session_row(storage, session):
    async with storage.repository.database.transaction() as conn:
        return await conn.one('SELECT * FROM agent_sessions WHERE session_id=?', (session,))


def test_restart_keeps_fields_phase_and_history(settings):
    async def scenario():
        session = uuid4().hex
        async with service_for(settings) as (service, storage):
            await service.set_info(session, 'hphm1', 'test-plate')
            first = await service.chat(ChatInput(session, '新对话'))
            assert first.content.startswith('<state>1001</state>')
            await service.chat(ChatInput(session, '继续办理'))
            values = await storage.repository.load(session)
            assert values['ywrysw'] is False
            assert type(values['jdcsl']) is int and values['jdcsl'] == 0
        llm = RecordingLLM()
        async with service_for(replace(settings, auto_migrate=False), llm) as (service, storage):
            assert json.loads(await service.get_info(session, 'hphm1')) == 'test-plate'
            result = await service.chat(ChatInput(session, '重启后的新输入'))
            assert result.content.startswith('<state>1002</state>')
            assert any(m['content'] == '继续办理' for m in llm.histories[0])
            row = await session_row(storage, session)
            graph = json.loads(row['graph_state'])
            assert graph['turn_count'] == 3
            assert graph['dialogue'][-2]['content'] == '重启后的新输入'
            assert graph['state_version'] == row['version']
    run(scenario())


@pytest.mark.parametrize('stream', [False, True])
def test_idempotency_replay_after_restart_and_conflict(settings, stream):
    async def scenario():
        session = uuid4().hex
        request = ChatInput(session, '继续办理', request_key='stable-key', need_form_update=True)
        async with service_for(settings) as (service, storage):
            await service.chat(ChatInput(session, '新对话'))
            if stream:
                first = [event async for event in service.stream(request)]
                assert any(isinstance(event, FormUpdate) for event in first)
            else:
                first = await service.chat(request)
            before = await session_row(storage, session)
        llm = RecordingLLM()
        async with service_for(settings, llm) as (service, storage):
            result = await service.chat(request)
            assert result.content == '<state>1002</state>已记录'
            assert llm.histories == []
            assert (await session_row(storage, session))['version'] == before['version']
            with pytest.raises(ChatBackendConflictError):
                await service.chat(replace(request, text='different'))
    run(scenario())


def test_checkpoint_ahead_of_business_commit_is_discarded(settings):
    async def scenario():
        session = uuid4().hex
        async with service_for(settings) as (service, storage):
            await service.chat(ChatInput(session, '新对话'))
            original = storage.repository.commit_turn
            async def fail(*args): raise ConnectionError('business commit failed')
            async def unavailable(*args): raise ConnectionError('journal cleanup failed')
            storage.repository.commit_turn = fail
            storage.repository.abort_turn = unavailable
            with pytest.raises(ConnectionError):
                await service.chat(ChatInput(session, '未提交的输入', request_key='retry'))
            storage.repository.commit_turn = original
        llm = RecordingLLM()
        async with service_for(settings, llm) as (service, storage):
            await service.chat(ChatInput(session, '未提交的输入', request_key='retry'))
            assert all(m['content'] != '未提交的输入' for m in llm.histories[0])
            assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 2
    run(scenario())


def test_committed_business_state_restores_missing_checkpoint(settings):
    async def scenario():
        session = uuid4().hex
        async with service_for(settings) as (service, storage):
            await service.chat(ChatInput(session, '新对话'))
            await service.chat(ChatInput(session, '继续办理'))
            await storage.checkpointer.adelete_thread(session)
        async with service_for(settings) as (service, storage):
            reply = await service.chat(ChatInput(session, '再次输入'))
            assert reply.content.startswith('<state>1002</state>')
            assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 3
    run(scenario())


def test_disconnect_after_commit_can_replay_without_advancing(settings):
    async def scenario():
        session = uuid4().hex
        request = ChatInput(session, '继续办理', request_key='disconnect')
        async with service_for(settings) as (service, storage):
            await service.chat(ChatInput(session, '新对话'))
            initial = await session_row(storage, session)
            async with aclosing(service.stream(request)) as events:
                await anext(events)
                async with asyncio.timeout(5):
                    while (await session_row(storage, session))['version'] == initial['version']:
                        await asyncio.sleep(0.01)
                # 此时图及业务提交已完成，但消费者没有读到回合末尾。
            before = await session_row(storage, session)
        llm = RecordingLLM()
        async with service_for(settings, llm) as (service, storage):
            await service.chat(request)
            assert llm.histories == []
            assert (await session_row(storage, session))['version'] == before['version']
    run(scenario())


def test_delete_then_restart_and_stale_commit_cannot_revive_data(settings):
    async def scenario():
        session = uuid4().hex
        async with service_for(settings) as (service, storage):
            await service.set_info(session, 'hphm1', 'old-plate')
            await service.chat(ChatInput(session, '新对话', request_key='old-key'))
            ticket = await storage.repository.begin_turn(session, None, 'pending')
            await service.delete_session(session, service._backend.reset_checkpoint)
            with pytest.raises(ChatBackendConflictError):
                await storage.repository.commit_turn(ticket, {'hphm1':'revived'}, {})
            assert await storage.checkpointer.aget_tuple({'configurable': {'thread_id':session}}) is None
        async with service_for(settings) as (service, storage):
            assert json.loads(await service.get_info(session, 'hphm1')) == ''
            reply = await service.chat(ChatInput(session, '新会话', request_key='old-key'))
            assert reply.content.startswith('<state>1001</state>')
            assert json.loads((await session_row(storage, session))['graph_state'])['turn_count'] == 1
    run(scenario())


def test_business_version_conflict_rejects_stale_turn(settings):
    async def scenario():
        session = uuid4().hex
        async with service_for(settings) as (service, storage):
            ticket = await storage.repository.begin_turn(session, None, 'pending')
            await storage.repository.apply_update(session, {'jdcsl':2})
            with pytest.raises(ChatBackendConflictError):
                await storage.repository.commit_turn(ticket, {'jdcsl':1}, {})
            assert (await storage.repository.load(session))['jdcsl'] == 2
    run(scenario())


@pytest.mark.parametrize('boundary', ['before', 'after'])
def test_abrupt_process_exit_at_commit_boundary(settings, boundary):
    session = uuid4().hex
    async def initialize():
        async with service_for(settings) as (service, storage):
            await service.chat(ChatInput(session, '新对话'))
    run(initialize())
    worker = Path(__file__).with_name('crash_worker.py')
    result = subprocess.run([sys.executable, str(worker), settings.kind, settings.sqlite_path, session, boundary],
                            capture_output=True, timeout=45)
    assert result.returncode == (23 if boundary=='before' else 24), result.stderr.decode(errors='replace')
    async def recover():
        llm = RecordingLLM()
        async with service_for(replace(settings,auto_migrate=False),llm) as (service,storage):
            reply=await service.chat(ChatInput(session,'继续办理',request_key='crash-key'))
            assert reply.content.startswith('<state>1002</state>')
            assert bool(llm.histories) is (boundary=='before')
            row=await session_row(storage,session)
            assert json.loads(row['graph_state'])['turn_count']==2
    run(recover())