"""Slow agent initialization must not monopolize the gateway event loop."""

import asyncio
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from contextvars import ContextVar
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest

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
from deerflow.tools.builtins.tool_search import DeferredToolRegistry, get_deferred_registry, reset_deferred_registry, set_deferred_registry
from deerflow.utils import background_init
from deerflow.utils.file_io import FileIOPool


def _background(factory):
    factory._deerflow_background_build = True
    return factory


@pytest.fixture
def build_pool(monkeypatch):
    pool = FileIOPool(1, 1, "test-agent-init")
    monkeypatch.setattr(background_init, "_INITIALIZATION_IO", pool)
    yield pool
    pool._pool.shutdown(wait=True)


@pytest.mark.asyncio
async def test_shared_initializer_forwards_arguments_and_commits_context(build_pool):
    from deerflow.mcp.cache import _owner_initialization_loop

    loop = asyncio.get_running_loop()
    caller_thread = threading.get_ident()
    state = ContextVar("shared-initializer-context", default="before")

    def initialize(value, *, suffix):
        assert threading.get_ident() != caller_thread
        assert _owner_initialization_loop.get() is loop
        assert state.get() == "before"
        state.set(None)
        return value + suffix

    assert await background_init.run_background_init(initialize, "tools", suffix=" ready") == "tools ready"
    assert state.get() is None
    assert _owner_initialization_loop.get() is None


@pytest.mark.asyncio
async def test_failed_shared_initializer_does_not_commit_context(build_pool):
    state = ContextVar("failed-shared-initializer-context", default="before")

    def fail():
        state.set("partially initialized")
        raise ValueError("initialization failed")

    with pytest.raises(ValueError, match="initialization failed"):
        await background_init.run_background_init(fail)
    assert state.get() == "before"
    assert await background_init.run_background_init(lambda: state.get()) == "before"


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


async def _make_run(monkeypatch, factory, *, config=None, app_config=None, thread_id="build-responsiveness"):
    manager = RunManager()
    record = await manager.create(thread_id)
    bridge = MemoryStreamBridge()
    monkeypatch.setattr(bridge, "cleanup", AsyncMock())
    task = asyncio.create_task(
        worker.run_agent(
            bridge,
            manager,
            record,
            ctx=worker.RunContext(checkpointer=None, app_config=app_config),
            agent_factory=factory,
            graph_input={},
            config=config or {},
        )
    )
    record.task = task
    return SimpleNamespace(task=task, manager=manager, record=record, bridge=bridge)


@pytest.mark.asyncio
async def test_slow_agent_build_keeps_parallel_requests_responsive(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    entered = asyncio.Event()
    factory_finished = threading.Event()
    factory_started_at = []

    @_background
    def factory(*, config):
        factory_started_at.append(time.perf_counter())
        loop.call_soon_threadsafe(entered.set)
        deadline = time.perf_counter() + 0.18
        while time.perf_counter() < deadline:
            pass
        factory_finished.set()
        return FinishedAgent()

    run = await _make_run(monkeypatch, factory)
    await asyncio.wait_for(entered.wait(), 3)
    probe_delay = time.perf_counter() - factory_started_at[0]
    responsive = not factory_finished.is_set()
    await asyncio.wait_for(run.task, 3)

    assert run.record.status == RunStatus.success
    assert run.bridge._streams[run.record.run_id].ended
    print(f"agent build concurrent probe delay = {probe_delay:.3f}s")
    assert responsive, f"Agent initialization blocked parallel requests for {probe_delay:.3f}s"


@pytest.mark.asyncio
async def test_custom_factory_keeps_caller_loop_and_context_mutations(monkeypatch):
    """Third-party synchronous factories have no worker-safety declaration."""
    loop = asyncio.get_running_loop()
    caller_thread = threading.get_ident()
    context = ContextVar("custom-build-context", default="before")
    observed = []
    app_config = object()

    class Agent:
        async def astream(self, *args, **kwargs):
            observed.append(context.get())
            assert asyncio.get_running_loop() is loop
            yield {"messages": []}

    def factory(*, config, app_config):
        assert asyncio.get_running_loop() is loop
        assert threading.get_ident() == caller_thread
        assert app_config is expected_app_config
        assert config["context"]["app_config"] is app_config
        context.set("set by custom factory")
        return Agent()

    expected_app_config = app_config
    run = await _make_run(monkeypatch, factory, app_config=app_config)
    await asyncio.wait_for(run.task, 3)

    assert run.record.status == RunStatus.success
    assert observed == ["set by custom factory"]
    assert context.get() == "before", "Run-local state must not leak to its caller"


@pytest.mark.asyncio
async def test_factory_failure_ends_stream_and_does_not_poison_next_run(monkeypatch, build_pool):
    @_background
    def failing_factory(*, config):
        raise ValueError("synthetic initialization failure")

    failed = await _make_run(monkeypatch, failing_factory)
    await asyncio.wait_for(failed.task, 3)
    assert failed.record.status == RunStatus.error
    assert failed.record.error == "synthetic initialization failure"
    stream = failed.bridge._streams[failed.record.run_id]
    assert stream.ended
    assert any(event.event == "error" for event in stream.events)

    recovered = await _make_run(monkeypatch, _background(lambda **kwargs: FinishedAgent()))
    await asyncio.wait_for(recovered.task, 3)
    assert recovered.record.status == RunStatus.success
    assert recovered.bridge._streams[recovered.record.run_id].ended


@pytest.mark.asyncio
@pytest.mark.parametrize("clear_registry", [False, True])
async def test_background_build_preserves_user_config_and_deferred_registry(monkeypatch, build_pool, clear_registry):
    loop = asyncio.get_running_loop()
    loop_thread = threading.get_ident()
    request_user = ContextVar("build-request-user", default="outside")
    inherited_registry = DeferredToolRegistry()
    previous_registry = get_deferred_registry()
    set_deferred_registry(inherited_registry)
    observations = {}
    app_configs = {name: object() for name in ("alice", "bob")}
    registries = {name: DeferredToolRegistry() for name in app_configs}

    class Agent:
        def __init__(self, user):
            self.user = user

        async def astream(self, *args, config, **kwargs):
            assert asyncio.get_running_loop() is loop
            assert threading.get_ident() == loop_thread
            observations[self.user] = (request_user.get(), get_deferred_registry(), config["configurable"]["built_for"])
            yield {"messages": []}

    @_background
    def factory(*, config, app_config):
        user = request_user.get()
        assert threading.get_ident() != loop_thread
        assert app_config is app_configs[user]
        assert config["context"]["app_config"] is app_config
        assert get_deferred_registry() is inherited_registry
        config["configurable"]["built_for"] = user
        if clear_registry:
            reset_deferred_registry()
        else:
            set_deferred_registry(registries[user])
        return Agent(user)

    runs = []
    try:
        for user in app_configs:
            token = request_user.set(user)
            try:
                runs.append(await _make_run(monkeypatch, factory, app_config=app_configs[user], thread_id=user))
            finally:
                request_user.reset(token)
        await asyncio.wait_for(asyncio.gather(*(run.task for run in runs)), 3)
        assert all(run.record.status == RunStatus.success for run in runs)
        assert request_user.get() == "outside"
        assert get_deferred_registry() is inherited_registry
        for user in app_configs:
            assert observations[user] == (user, None if clear_registry else registries[user], user)
    finally:
        set_deferred_registry(previous_registry)


@pytest.mark.asyncio
async def test_build_queue_serializes_initialization_without_consuming_default_executor(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    first_entered = asyncio.Event()
    release = threading.Event()
    counter_lock = threading.Lock()
    active = peak = entered = 0
    default_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="test-general-worker")
    loop.set_default_executor(default_pool)

    @_background
    def factory(*, config):
        nonlocal active, peak, entered
        with counter_lock:
            active += 1
            entered += 1
            peak = max(active, peak)
            loop.call_soon_threadsafe(first_entered.set)
        try:
            assert release.wait(5), "Test did not release initialization worker"
            return FinishedAgent()
        finally:
            with counter_lock:
                active -= 1

    runs = [await _make_run(monkeypatch, factory, thread_id=f"queue-{index}") for index in range(4)]
    try:
        await asyncio.wait_for(first_entered.wait(), 3)
        assert await asyncio.wait_for(asyncio.to_thread(lambda: "authentication worker available"), 1) == "authentication worker available"
        await asyncio.sleep(0.03)
        assert entered == 1, "Cold builders must wait without competing for shared initialization state"
    finally:
        release.set()
        await asyncio.wait_for(asyncio.gather(*(run.task for run in runs)), 3)
        default_pool.shutdown(wait=True)

    assert peak == 1
    assert entered == 4
    assert all(run.record.status == RunStatus.success for run in runs)


@pytest.mark.asyncio
async def test_repeated_cancel_during_build_drains_worker_without_starting_graph(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    entered = asyncio.Event()
    release = threading.Event()
    graph_started = []
    finalization_registries = []
    original_finalize = worker._finalize_run
    previous_registry = get_deferred_registry()
    caller_registry = DeferredToolRegistry()
    set_deferred_registry(caller_registry)

    async def finalize(**kwargs):
        finalization_registries.append(get_deferred_registry())
        await original_finalize(**kwargs)

    monkeypatch.setattr(worker, "_finalize_run", finalize)

    class Agent:
        async def astream(self, *args, **kwargs):
            graph_started.append(True)
            yield {"messages": []}

    @_background
    def factory(*, config):
        set_deferred_registry(DeferredToolRegistry())
        loop.call_soon_threadsafe(entered.set)
        assert release.wait(5), "Test did not release cancelled builder"
        return Agent()

    run = await _make_run(monkeypatch, factory)
    try:
        await asyncio.wait_for(entered.wait(), 3)
        assert await run.manager.cancel(run.record.run_id)
        await asyncio.sleep(0)
        assert not run.task.done(), "Cancellation must not abandon a live initialization worker"
        run.task.cancel()
        await asyncio.sleep(0)
        assert not run.task.done(), "Repeated cancellation must still drain the builder"
    finally:
        release.set()
        await asyncio.wait_for(run.task, 3)
        set_deferred_registry(previous_registry)

    assert not graph_started
    assert run.record.status == RunStatus.interrupted
    assert run.bridge._streams[run.record.run_id].ended
    assert finalization_registries == [caller_registry], "Cancelled build context must not be committed to the run"


@pytest.mark.asyncio
async def test_queued_cancel_does_not_build_or_wait_for_active_factory(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    entered = asyncio.Event()
    release = threading.Event()
    queued_built = []

    @_background
    def first_factory(*, config):
        loop.call_soon_threadsafe(entered.set)
        assert release.wait(5)
        return FinishedAgent()

    @_background
    def queued_factory(*, config):
        queued_built.append(True)
        return FinishedAgent()

    first = await _make_run(monkeypatch, first_factory, thread_id="active")
    queued = None
    try:
        await asyncio.wait_for(entered.wait(), 3)
        queued = await _make_run(monkeypatch, queued_factory, thread_id="queued")
        await asyncio.sleep(0.02)
        await queued.manager.cancel(queued.record.run_id)
        await asyncio.wait_for(queued.task, 1)
        assert queued.record.status == RunStatus.interrupted
        assert queued.bridge._streams[queued.record.run_id].ended
        assert not queued_built
        assert not first.task.done()
    finally:
        release.set()
        await asyncio.wait_for(asyncio.gather(first.task, *([queued.task] if queued else [])), 3)

    assert first.record.status == RunStatus.success


@pytest.mark.asyncio
async def test_abort_flag_set_during_build_prevents_late_stream_start(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    entered = asyncio.Event()
    release = threading.Event()
    graph_started = []

    class Agent:
        async def astream(self, *args, **kwargs):
            graph_started.append(True)
            yield {"messages": []}

    @_background
    def factory(*, config):
        loop.call_soon_threadsafe(entered.set)
        assert release.wait(5)
        return Agent()

    run = await _make_run(monkeypatch, factory)
    try:
        await asyncio.wait_for(entered.wait(), 3)
        run.record.abort_event.set()
    finally:
        release.set()
        await asyncio.wait_for(run.task, 3)

    assert graph_started == []
    assert run.record.status == RunStatus.interrupted
    assert run.bridge._streams[run.record.run_id].ended


@pytest.mark.asyncio
async def test_existing_conversation_keeps_streaming_while_another_agent_initializes(monkeypatch, build_pool):
    loop = asyncio.get_running_loop()
    stream_started = asyncio.Event()
    continue_stream = asyncio.Event()
    build_entered = asyncio.Event()
    release_build = threading.Event()
    build_finished = threading.Event()

    class StreamingAgent:
        async def astream(self, *args, **kwargs):
            stream_started.set()
            yield {"messages": [], "answer": "first delta"}
            await continue_stream.wait()
            yield {"messages": [], "answer": "second delta"}

    @_background
    def slow_factory(*, config):
        loop.call_soon_threadsafe(build_entered.set)
        assert release_build.wait(5), "Test did not release the second conversation's initializer"
        build_finished.set()
        return FinishedAgent()

    existing = await _make_run(monkeypatch, lambda **kwargs: StreamingAgent(), thread_id="already-streaming")
    initializing = None
    try:
        await asyncio.wait_for(stream_started.wait(), 3)
        initializing = await _make_run(monkeypatch, slow_factory, thread_id="cold-conversation")
        await asyncio.wait_for(build_entered.wait(), 3)
        continue_stream.set()
        await asyncio.wait_for(existing.task, 1)

        stream = existing.bridge._streams[existing.record.run_id]
        answers = [event.data["answer"] for event in stream.events if event.event == "values"]
        assert answers == ["first delta", "second delta"]
        assert stream.ended
        assert existing.record.status == RunStatus.success
        assert not initializing.task.done()
        assert not build_finished.is_set(), "The existing answer must finish before the cold build is released"
    finally:
        continue_stream.set()
        release_build.set()
        await asyncio.wait_for(asyncio.gather(existing.task, *([initializing.task] if initializing else [])), 3)

    assert initializing.record.status == RunStatus.success


@pytest.mark.asyncio
async def test_builtin_lead_factory_uses_background_initializer(monkeypatch, build_pool):
    from deerflow.agents.lead_agent import agent as lead_agent

    loop_thread = threading.get_ident()
    app_config = object()
    observed_build_threads = []

    def build_body(config, *, app_config):
        assert app_config is expected_config
        observed_build_threads.append(threading.get_ident())
        return FinishedAgent()

    expected_config = app_config
    monkeypatch.setattr(lead_agent, "_make_lead_agent", build_body)
    # Exercise the real exported factory: a missing production opt-in must fail
    # even if all tests using synthetic marked factories continue to pass.
    run = await _make_run(monkeypatch, lead_agent.make_lead_agent, app_config=app_config)
    await asyncio.wait_for(run.task, 3)

    assert run.record.status == RunStatus.success
    assert len(observed_build_threads) == 1
    assert observed_build_threads[0] != loop_thread