import asyncio
from types import SimpleNamespace

import pytest
from fastapi import HTTPException

from app.gateway.worker_pool.admission import Admission, bounded_worker_run, run_with_worker_admission


class Bridge:
    def __init__(self): self.events = []
    async def publish(self, run_id, event, data): self.events.append((run_id, event, data))
    async def publish_end(self, run_id): self.events.append((run_id, "end", None))
    async def cleanup(self, *args, **kwargs): pass


class Manager:
    def __init__(self): self.records = {}
    async def set_status(self, run_id, status, **kwargs): self.records[run_id].status = status


async def prepare(monkeypatch, capacity=4, queue=128, timeout=10):
    async def owner(thread): return "owner"
    monkeypatch.setattr("deerflow.runtime.thread_paths.aresolve_path_user_id", owner)
    gate = Admission(capacity, queue, timeout); bridge = Bridge(); manager = Manager()
    app = SimpleNamespace(state=SimpleNamespace(worker_admission=gate, worker_deferred_admission=True, stream_bridge=bridge, run_manager=manager))
    return gate, bridge, manager, SimpleNamespace(app=app)


@pytest.mark.asyncio
async def test_70_streams_get_real_run_metadata_before_waiting_for_capacity(monkeypatch):
    gate, bridge, manager, request = await prepare(monkeypatch)
    release = asyncio.Event(); peak = 0
    async def runner(**kwargs):
        nonlocal peak
        peak = max(peak, gate.active)
        await release.wait()
        kwargs["record"].status = "success"
    @bounded_worker_run
    async def start(body, thread_id, request):
        record = SimpleNamespace(run_id=thread_id, thread_id=thread_id, status="pending")
        manager.records[thread_id] = record
        record.task = asyncio.create_task(run_with_worker_admission(runner, bridge=bridge, run_manager=manager, record=record))
        return record
    records = await asyncio.wait_for(asyncio.gather(*(start(None, str(i), request) for i in range(70))), 1)
    await asyncio.sleep(.02)
    assert len([e for e in bridge.events if e[1] == "metadata"]) == 70
    assert gate.active == 4 and gate.waiting == 66
    assert all(not r.task.done() for r in records)
    release.set(); await asyncio.gather(*(r.task for r in records)); await asyncio.sleep(0)
    assert peak == 4 and gate.active == gate.waiting == 0


@pytest.mark.asyncio
async def test_queued_cancel_timeout_and_overflow_finish_stream_without_capacity_leak(monkeypatch):
    gate, bridge, manager, request = await prepare(monkeypatch, capacity=1, queue=2, timeout=.04)
    release = asyncio.Event()
    async def runner(**kwargs): await release.wait()
    @bounded_worker_run
    async def start(body, thread_id, request):
        record = SimpleNamespace(run_id=thread_id, thread_id=thread_id, status="pending")
        manager.records[thread_id] = record
        record.task = asyncio.create_task(run_with_worker_admission(runner, bridge=bridge, run_manager=manager, record=record))
        return record
    first = await start(None, "one", request)
    second = await start(None, "two", request)
    third = await start(None, "three", request)
    with pytest.raises(HTTPException) as exc: await start(None, "overflow", request)
    assert exc.value.status_code == 429
    await asyncio.sleep(.01); second.task.cancel()
    await asyncio.gather(second.task, third.task, return_exceptions=True)
    assert manager.records["two"].status.value == "interrupted"
    assert manager.records["three"].status.value == "error"
    assert any(e[0] == "three" and e[1] == "end" for e in bridge.events)
    assert gate.active == 1 and gate.waiting == 0
    release.set(); await first.task; await asyncio.sleep(0)
    assert gate.active == 0


@pytest.mark.asyncio
async def test_preparation_exception_releases_reserved_place(monkeypatch):
    gate, bridge, manager, request = await prepare(monkeypatch)
    @bounded_worker_run
    async def fail(*args): raise HTTPException(409, "conflict")
    with pytest.raises(HTTPException): await fail(None, "one", request)
    assert gate.active == gate.waiting == 0