import asyncio
from types import SimpleNamespace

import pytest

from deerflow.config.checkpointer_config import CheckpointerConfig
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
from deerflow.runtime.store.async_provider import make_store
from deerflow.runtime.runs.manager import RunRecord
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus
from deerflow.runtime.runs.worker import _finalize_run


@pytest.mark.asyncio
async def test_concurrent_sqlite_initialization_and_shared_store(tmp_path):
    config = SimpleNamespace(checkpointer=CheckpointerConfig(type="sqlite", connection_string=str(tmp_path / "chat.db")))

    async def node(index):
        async with make_checkpointer(config) as saver, make_store(config) as store:
            assert (await (await saver.conn.execute("PRAGMA busy_timeout")).fetchone())[0] == 30000
            await store.aput(("qa",), str(index), {"index": index})
            assert (await store.aget(("qa",), str(index))).value == {"index": index}

    await asyncio.gather(*(node(i) for i in range(4)))
    async with make_store(config) as store:
        assert len(await store.asearch(("qa",))) == 4


@pytest.mark.asyncio
async def test_pg_disabled_checkpoint_and_store_use_same_sqlite(tmp_path, monkeypatch):
    monkeypatch.setenv("DEER_FLOW_USE_PG_CHECKPOINTER", "false")
    path = str(tmp_path / "fallback.db")
    config = SimpleNamespace(checkpointer=CheckpointerConfig(type="postgres", connection_string="postgresql://unused", sqlite_connection_string=path))
    async with make_checkpointer(config) as saver, make_store(config) as store:
        saver_path = (await (await saver.conn.execute("PRAGMA database_list")).fetchone())[2]
        store_path = (await (await store.conn.execute("PRAGMA database_list")).fetchone())[2]
        assert saver_path == store_path
        await store.aput(("qa",), "fallback", {"ok": True})
        assert (await store.aget(("qa",), "fallback")).value == {"ok": True}


@pytest.mark.asyncio
async def test_run_finalization_does_not_wait_for_title_checkpoint_read():
    title_read_started = asyncio.Event()
    release_title_read = asyncio.Event()
    status_updates: list[tuple[str, str]] = []
    titles: list[tuple[str, str]] = []

    class Checkpointer:
        async def aget_tuple(self, _config):
            title_read_started.set()
            await release_title_read.wait()
            return SimpleNamespace(checkpoint={"channel_values": {"title": "异步标题"}})

    class ThreadStore:
        async def update_status(self, thread_id: str, status: str) -> None:
            status_updates.append((thread_id, status))

        async def update_display_name(self, thread_id: str, title: str) -> None:
            titles.append((thread_id, title))

    class Manager:
        def is_latest(self, _record: RunRecord) -> bool:
            return True

    record = RunRecord(
        run_id="run-title-test",
        thread_id="thread-title-test",
        assistant_id="assistant",
        status=RunStatus.success,
        on_disconnect=DisconnectMode.continue_,
    )
    await asyncio.wait_for(
        _finalize_run(
            journal=None,
            run_manager=Manager(),
            record=record,
            ctx=SimpleNamespace(checkpointer=Checkpointer(), thread_store=ThreadStore(), app_config=None),
        ),
        timeout=0.1,
    )

    assert status_updates == [("thread-title-test", "idle")]
    await asyncio.wait_for(title_read_started.wait(), timeout=0.1)
    assert titles == []
    release_title_read.set()
    for _ in range(10):
        if titles:
            break
        await asyncio.sleep(0)
    assert titles == [("thread-title-test", "异步标题")]