import asyncio
from collections.abc import Callable
from types import SimpleNamespace
from typing import Any

import pytest

from tests.conftest import FakeComfy


async def _run(tools: Any, **kw: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    events: list[dict[str, Any]] = []

    async def emit(e: dict[str, Any]) -> None:
        events.append(e)

    result = await tools.comfier_generate_image(
        positive_prompt="a marble",
        seed="5",
        size="64x64",
        __request__=SimpleNamespace(),
        __user__={"id": "u"},
        __event_emitter__=emit,
        **kw,
    )
    return result, events


async def test_success(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
    tools, env = make_tools()
    result, events = await _run(tools)
    assert result["status"] == "success"
    assert len(result["images"]) == 1 and len(env.stored) == 1
    assert comfy.submitted[0]["520:517"]["inputs"]["seed"] == 5
    assert comfy.submitted[0]["520:518"]["inputs"]["value"] == 64
    assert events[-1]["type"] == "files"


async def test_execution_error(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
    comfy.script = [
        {
            "type": "execution_error",
            "data": {"prompt_id": "pid-1", "exception_message": "boom"},
        }
    ]
    tools, _ = make_tools()
    with pytest.raises(RuntimeError, match="boom"):
        await _run(tools)


async def test_ws_drop_retry_declined(
    make_tools: Callable[[], Any], comfy: FakeComfy
) -> None:
    comfy.drop_ws = True
    tools, _ = make_tools()
    asked: list[dict[str, Any]] = []

    async def call(e: dict[str, Any]) -> bool:
        asked.append(e)
        return False

    with pytest.raises(ConnectionError):
        await _run(tools, __event_call__=call)
    assert len(asked) == 1


async def test_ws_drop_retry_accepted(
    make_tools: Callable[[], Any], comfy: FakeComfy
) -> None:
    comfy.drop_ws = True
    tools, _ = make_tools()

    async def call(e: dict[str, Any]) -> bool:
        comfy.drop_ws = False
        return True

    result, _ = await _run(tools, __event_call__=call)
    assert result["status"] == "success"


async def test_cancel_interrupts_backend(
    make_tools: Callable[[], Any], comfy: FakeComfy
) -> None:
    comfy.hang = True
    tools, _ = make_tools()
    task = asyncio.create_task(_run(tools))
    while not comfy.submitted:
        await asyncio.sleep(0.05)
    await asyncio.sleep(0.1)
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert comfy.interrupts == [{"prompt_id": "pid-1"}]


@pytest.fixture
def fast_idle(monkeypatch: pytest.MonkeyPatch) -> None:
    """Shrink the 1s WebSocket idle timeout so stall detection takes ms, not seconds."""
    import comfier

    real_wait_for = asyncio.wait_for

    async def wait_for(aw: Any, timeout: float | None = None) -> Any:
        return await real_wait_for(aw, 0.01 if timeout == 1.0 else timeout)

    monkeypatch.setattr(comfier.asyncio, "wait_for", wait_for)


async def test_stalled_backend_lost_prompt(
    make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
) -> None:
    comfy.hang = True
    comfy.queue_has_prompt = False  # not in /queue, and /history is empty
    tools, _ = make_tools()
    with pytest.raises(ConnectionError, match="lost prompt pid-1"):
        await asyncio.wait_for(_run(tools), timeout=5)


async def test_stalled_backend_still_queued_keeps_waiting(
    make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
) -> None:
    comfy.hang = True  # silent, but /queue still lists the prompt
    tools, _ = make_tools()
    task = asyncio.create_task(_run(tools))
    while not comfy.submitted:
        await asyncio.sleep(0.01)
    await asyncio.sleep(0.3)  # many idle ticks and liveness checks pass
    assert not task.done()
    await comfy.release()
    result, _ = await asyncio.wait_for(task, timeout=5)
    assert result["status"] == "success"