"""A completed/cancelled graph must not leave SSE alive on metadata I/O."""

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest

from deerflow.runtime import journal as journal_module
from deerflow.runtime.runs import worker
from deerflow.runtime.runs.manager import RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge


class FinishedAgent:
    async def astream(self, *args, **kwargs):
        yield {"messages": []}


async def make_run(monkeypatch, phase, *, graph=None, ignore_cancellation=False):
    monkeypatch.setattr(worker, "_RUN_FINALIZATION_TIMEOUT_SECONDS", 0.03)
    entered = asyncio.Event()
    cancelled = asyncio.Event()
    release = asyncio.Event()
    drained = asyncio.Event()

    async def blocked(*args, **kwargs):
        entered.set()
        try:
            await release.wait()
        except asyncio.CancelledError:
            cancelled.set()
            if ignore_cancellation:
                await release.wait()
            else:
                raise
        finally:
            drained.set()

    fake_journal = SimpleNamespace(flush=AsyncMock(), get_completion_data=lambda: {})
    monkeypatch.setattr(journal_module, "RunJournal", lambda **kwargs: fake_journal)
    manager = RunManager()
    record = await manager.create("cleanup-test")
    bridge = MemoryStreamBridge()
    # Preserve the stream for assertions and avoid leaving a 60s sleeper.
    monkeypatch.setattr(bridge, "cleanup", AsyncMock())
    thread_store = SimpleNamespace(update_display_name=AsyncMock(), update_status=AsyncMock())
    checkpoint_calls = 0

    async def checkpoint(*args):
        nonlocal checkpoint_calls
        checkpoint_calls += 1
        # Snapshot and terminal inspection precede final metadata cleanup.
        if (phase == "checkpoint" and checkpoint_calls == 3) or (phase == "terminal_checkpoint" and checkpoint_calls == 2):
            await blocked()
        return SimpleNamespace(checkpoint={"channel_values": {"title": "Test"}}, config={})

    checkpointer = SimpleNamespace(aget_tuple=checkpoint)
    if phase == "journal":
        fake_journal.flush = blocked
    elif phase == "completion":
        monkeypatch.setattr(manager, "update_run_completion", blocked)
    elif phase == "title":
        thread_store.update_display_name = blocked
    elif phase == "thread_status":
        thread_store.update_status = blocked
    elif phase == "terminal_status":
        async def persist_status(run_id, status, **kwargs):
            if status != "running":
                await blocked()
        manager._store = SimpleNamespace(update_status=persist_status, update_run_completion=AsyncMock())
    task = asyncio.create_task(worker.run_agent(
        bridge, manager, record,
        ctx=worker.RunContext(checkpointer=checkpointer, event_store=object(), thread_store=thread_store),
        agent_factory=lambda **kwargs: graph or FinishedAgent(),
        graph_input={}, config={},
    ))
    record.task = task
    return SimpleNamespace(
        task=task, manager=manager, record=record, bridge=bridge,
        entered=entered, cancelled=cancelled, release=release, drained=drained,
        journal=fake_journal, thread_store=thread_store,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("phase", ["journal", "completion", "checkpoint", "title", "thread_status", "terminal_status", "terminal_checkpoint"])
async def test_metadata_stall_always_ends_sse(monkeypatch, phase, caplog):
    run = await make_run(monkeypatch, phase)
    await asyncio.wait_for(run.task, 1)
    await asyncio.wait_for(run.cancelled.wait(), 1)
    assert run.entered.is_set()
    assert run.record.status == (RunStatus.error if phase == "terminal_checkpoint" else RunStatus.success)
    assert run.bridge._streams[run.record.run_id].ended
    assert "RUN_COMPLETION_TIMEOUT" in caplog.text
    if phase == "terminal_checkpoint":
        assert any(event.event == "error" for event in run.bridge._streams[run.record.run_id].events)


@pytest.mark.asyncio
async def test_delayed_cancellation_acknowledgement_does_not_hold_stream(monkeypatch):
    run = await make_run(monkeypatch, "journal", ignore_cancellation=True)
    try:
        await asyncio.wait_for(run.task, 1)
        assert run.bridge._streams[run.record.run_id].ended
        await asyncio.wait_for(run.cancelled.wait(), 1)
    finally:
        run.release.set()
        await asyncio.wait_for(run.drained.wait(), 1)
        await asyncio.sleep(0)
    run.thread_store.update_display_name.assert_not_awaited()
    run.thread_store.update_status.assert_not_awaited()


@pytest.mark.asyncio
async def test_cancelled_graph_ends_sse_even_when_cleanup_stalls(monkeypatch):
    graph_started = asyncio.Event()

    class WaitingAgent:
        async def astream(self, *args, **kwargs):
            graph_started.set()
            await asyncio.Event().wait()
            yield {}

    run = await make_run(monkeypatch, "journal", graph=WaitingAgent())
    await asyncio.wait_for(graph_started.wait(), 1)
    await run.manager.cancel(run.record.run_id)
    await asyncio.wait_for(run.task, 1)
    assert run.record.status == RunStatus.interrupted
    assert run.bridge._streams[run.record.run_id].ended


@pytest.mark.asyncio
async def test_cancellation_during_cleanup_preserves_terminator(monkeypatch):
    run = await make_run(monkeypatch, "journal")
    await asyncio.wait_for(run.entered.wait(), 1)
    run.task.cancel()
    await asyncio.wait_for(run.task, 1)
    assert run.bridge._streams[run.record.run_id].ended
    await asyncio.wait_for(run.cancelled.wait(), 1)


@pytest.mark.asyncio
async def test_healthy_completion_still_persists_metadata(monkeypatch):
    run = await make_run(monkeypatch, "healthy")
    await asyncio.wait_for(run.task, 1)
    run.journal.flush.assert_awaited_once()
    run.thread_store.update_display_name.assert_awaited_once_with("cleanup-test", "Test")
    run.thread_store.update_status.assert_awaited_once_with("cleanup-test", "idle")
    assert run.record.status == RunStatus.success
    assert run.bridge._streams[run.record.run_id].ended


@pytest.mark.asyncio
async def test_older_cleanup_does_not_overwrite_newer_run_metadata(monkeypatch):
    run = await make_run(monkeypatch, "journal")
    await asyncio.wait_for(run.entered.wait(), 1)
    await run.manager.create(run.record.thread_id)
    run.release.set()
    await asyncio.wait_for(run.task, 1)
    run.thread_store.update_display_name.assert_not_awaited()
    run.thread_store.update_status.assert_not_awaited()


@pytest.mark.asyncio
async def test_completion_budget_does_not_time_out_executing_graph(monkeypatch):
    graph_started = asyncio.Event()
    graph_release = asyncio.Event()

    class LongAgent:
        async def astream(self, *args, **kwargs):
            graph_started.set()
            await graph_release.wait()
            yield {}

    run = await make_run(monkeypatch, "healthy", graph=LongAgent())
    try:
        await asyncio.wait_for(graph_started.wait(), 1)
        done, _ = await asyncio.wait({run.task}, timeout=0.06)
        assert not done
        assert run.record.status == RunStatus.running
        assert not run.bridge._streams[run.record.run_id].ended
    finally:
        graph_release.set()
        await asyncio.wait_for(run.task, 1)