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"}]