import asyncio
import json
import pathlib
from collections.abc import AsyncIterator, Callable
from typing import Any

import pytest
import pytest_asyncio
from aiohttp import WSMsgType, web

from devkit import FakeEnv, install

ROOT = pathlib.Path(__file__).parent.parent
PNG = bytes.fromhex(
    "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
    "0000000d49444154789c6360000002000001e221bc330000000049454e44ae426082"
)


# Make `import comfier` work at collection time; fixtures re-install per test.
install(FakeEnv(comfyui_url="", workflow="{}"))


class FakeComfy:
    """Minimal in-process ComfyUI. `script` lists WS messages sent after /prompt."""

    def __init__(self) -> None:
        self.script: list[dict[str, Any]] = []
        self.submitted: list[dict[str, Any]] = []
        self.interrupts: list[Any] = []
        self.drop_ws = False
        self.hang = False
        self.url = ""
        self._ws: web.WebSocketResponse | None = None
        self._ws_ready = asyncio.Event()

    def success_script(self, prompt_id: str) -> list[dict[str, Any]]:
        return [
            {"type": "executing", "data": {"node": "1", "prompt_id": prompt_id}},
            {
                "type": "progress",
                "data": {"value": 1, "max": 2, "prompt_id": prompt_id},
            },
            {
                "type": "executed",
                "data": {
                    "prompt_id": prompt_id,
                    "output": {"images": [{"filename": "a.png", "type": "output"}]},
                },
            },
            {"type": "executing", "data": {"node": None, "prompt_id": prompt_id}},
        ]

    async def ws(self, request: web.Request) -> web.WebSocketResponse:
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        self._ws = ws
        self._ws_ready.set()
        async for msg in ws:
            if msg.type == WSMsgType.ERROR:
                break
        return ws

    async def prompt(self, request: web.Request) -> web.Response:
        body = await request.json()
        self.submitted.append(body["prompt"])
        pid = "pid-1"
        await self._ws_ready.wait()
        assert self._ws is not None
        ws = self._ws
        asyncio.get_running_loop().call_soon(
            lambda: asyncio.ensure_future(self._play(ws, pid))
        )
        return web.json_response({"prompt_id": pid})

    async def _play(self, ws: web.WebSocketResponse, pid: str) -> None:
        if self.hang:
            return
        if self.drop_ws:
            await ws.close()
            return
        for msg in self.script or self.success_script(pid):
            await ws.send_str(json.dumps(msg))

    async def queue(self, request: web.Request) -> web.Response:
        if request.method == "POST":
            return web.json_response({})
        return web.json_response({"queue_running": [[0, "pid-1"]], "queue_pending": []})

    async def interrupt(self, request: web.Request) -> web.Response:
        self.interrupts.append(await request.json())
        return web.json_response({})

    async def history(self, request: web.Request) -> web.Response:
        return web.json_response({})

    async def view(self, request: web.Request) -> web.Response:
        return web.Response(body=PNG, content_type="image/png")


@pytest_asyncio.fixture
async def comfy() -> AsyncIterator[FakeComfy]:
    fake = FakeComfy()
    app = web.Application()
    app.add_routes(
        [
            web.get("/ws", fake.ws),
            web.post("/prompt", fake.prompt),
            web.get("/queue", fake.queue),
            web.post("/queue", fake.queue),
            web.post("/interrupt", fake.interrupt),
            web.get("/history/{id}", fake.history),
            web.get("/view", fake.view),
        ]
    )
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", 0)
    await site.start()
    port = site._server.sockets[0].getsockname()[1]  # type: ignore[union-attr]
    fake.url = f"http://127.0.0.1:{port}"
    yield fake
    await runner.cleanup()


@pytest.fixture
def make_tools(
    tmp_path: pathlib.Path, comfy: FakeComfy
) -> Callable[[], tuple[Any, FakeEnv]]:
    def factory() -> tuple[Any, FakeEnv]:
        env = FakeEnv(
            comfyui_url=comfy.url,
            workflow=(ROOT / "fixtures" / "workflow.json").read_text(encoding="utf-8"),
            nodes=json.loads(
                (ROOT / "fixtures" / "nodes.json").read_text(encoding="utf-8")
            ),
            output_dir=tmp_path,
        )
        install(env)
        from comfier import Tools

        return Tools(), env

    return factory