"""Tests for image generation — cluster handlers, router, node agent, hardware probe."""
import asyncio
import base64
import json
import os
from contextlib import asynccontextmanager
from pathlib import Path
from unittest.mock import AsyncMock, patch

import httpx
import psutil
from fastapi.testclient import TestClient

import app as app_module
import cluster
import config
import db
import hardware
import node_agent.agent as agent
from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS


# ── helpers ──────────────────────────────────────────────────────────────

def _reset():
    cluster.CLUSTER_NODES.clear()
    cluster.CLUSTER_EVENTS.clear()
    cluster.CLUSTER_COORDINATOR = None
    cluster._pending_pings.clear()
    cluster._pending_image.clear()


_published = []


async def _fake_publish(exchange, routing_key, payload):
    _published.append((exchange, routing_key, payload))


def make_client(tmp_path: Path) -> TestClient:
    os.environ["CAIC_ADMIN_PIN"] = "1234"
    db.DB_PATH = tmp_path / "caic-image.db"
    hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json"
    SESSIONS.clear()
    PIN_ATTEMPTS.clear()
    RATE_EVENTS.clear()
    db.init_db()
    return TestClient(app_module.app, raise_server_exceptions=False)


def _guest_headers(client: TestClient) -> dict:
    sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
    return {"X-Session-ID": sid, "Origin": "http://testserver"}


def _admin_headers(client: TestClient) -> dict:
    resp = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"})
    sid = resp.json()["session_id"]
    return {"X-Session-ID": sid, "Origin": "http://testserver"}


class FakeMsg:
    def __init__(self, body_dict: dict):
        self.body = json.dumps(body_dict).encode()

    @asynccontextmanager
    async def process(self):
        yield


class FakeExchange:
    def __init__(self, name=""):
        self.name = name
        self.published = []

    async def publish(self, msg, routing_key):
        self.published.append((msg, routing_key))


class FakeChannel:
    def __init__(self):
        self.exchanges = {}
        self.is_closed = False

    async def declare_exchange(self, name, typ, durable=True):
        self.exchanges[name] = FakeExchange(name)
        return self.exchanges[name]

    async def declare_queue(self, name="", exclusive=True):
        return self

    async def bind(self, exchange, routing_key):
        pass


# ── 1. cluster.handle_image_generated resolves pending request ───────────


def test_handle_image_generated_resolves_pending(monkeypatch):
    _reset()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }

    event = asyncio.Event()
    cluster._pending_image["req-123"] = ("", event)

    asyncio.run(cluster.handle_image_generated(
        AMQP_EXCHANGE_SYSTEM, "node.corsair.image_generated",
        {"node_name": "corsair", "request_id": "req-123", "image_base64": "aW1hZ2U="},
    ))

    assert event.is_set()
    result = cluster._pending_image.get("req-123")
    assert result is not None
    assert result[0] == "aW1hZ2U="
    assert cluster.CLUSTER_NODES["corsair"]["last_seen"] is not None


# ── 2. cluster.handle_image_failed resolves pending request ─────────────


def test_handle_image_failed_resolves_pending(monkeypatch):
    _reset()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
    }

    event = asyncio.Event()
    cluster._pending_image["req-456"] = ("", event)

    asyncio.run(cluster.handle_image_failed(
        AMQP_EXCHANGE_SYSTEM, "node.corsair.image_failed",
        {"node_name": "corsair", "request_id": "req-456", "error": "timeout"},
    ))

    assert event.is_set()
    result = cluster._pending_image.get("req-456")
    assert result is not None
    assert result[0] == ""


# ── 3. cluster.handle_image_failed unknown node ─────────────────────────


def test_handle_image_failed_unknown_node(caplog, monkeypatch):
    _reset()
    caplog.set_level("WARNING")
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    asyncio.run(cluster.handle_image_failed(
        AMQP_EXCHANGE_SYSTEM, "node.ghost.image_failed",
        {"node_name": "ghost", "request_id": "x", "error": "boom"},
    ))

    assert not any("unknown node" in rec.message for rec in caplog.records)


# ── 4. cluster.request_image_generate publishes command ──────────────────


def test_request_image_generate_publishes_command(monkeypatch):
    _reset()
    _published.clear()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }

    # Simulate immediate completion
    async def fake_wait():
        cluster._pending_image.clear()

    original_wait_for = asyncio.wait_for

    async def patched_wait_for(coro, timeout):
        cluster._pending_image["fake-id"] = ("aW1hZ2U=", asyncio.Event())
        cluster._pending_image["fake-id"][1].set()
        return None

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

    result = asyncio.run(cluster.request_image_generate(
        "corsair", "a red dragon", width=512, height=512, steps=10,
    ))

    assert len(_published) == 1
    exchange, rk, payload = _published[0]
    assert exchange == AMQP_EXCHANGE_ADMIN
    assert rk == "node.corsair.cmd.image_generate"
    assert payload["prompt"] == "a red dragon"
    assert payload["width"] == 512
    assert payload["height"] == 512
    assert payload["steps"] == 10
    assert "request_id" in payload


# ── 5. cluster.request_image_generate unknown node ──────────────────────


def test_request_image_generate_unknown_node(monkeypatch):
    _reset()
    _published.clear()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    result = asyncio.run(cluster.request_image_generate("ghost", "prompt"))
    assert result is None
    assert len(_published) == 0


# ── 6. cluster.request_image_generate node lacks capability ─────────────


def test_request_image_generate_no_capability(monkeypatch):
    _reset()
    _published.clear()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    cluster.CLUSTER_NODES["jarvis"] = {
        "name": "jarvis", "type": "worker", "status": "active",
        "capabilities": ["llm"],
    }

    result = asyncio.run(cluster.request_image_generate("jarvis", "prompt"))
    assert result is None
    assert len(_published) == 0


# ── 7. cluster.request_image_generate timeout ───────────────────────────


def test_request_image_generate_timeout(monkeypatch):
    _reset()
    _published.clear()
    monkeypatch.setattr(cluster, "publish", _fake_publish)

    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }

    async def timeout_wait(coro, timeout):
        raise asyncio.TimeoutError()

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

    result = asyncio.run(cluster.request_image_generate("corsair", "prompt", timeout=1))
    assert result is None
    assert len(cluster._pending_image) == 0


# ── 8. _find_image_node selects active image_gen node ───────────────────


def test_find_image_node_selects_active():
    from routers.image import _find_image_node

    _reset()
    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }
    cluster.CLUSTER_NODES["jarvis"] = {
        "name": "jarvis", "type": "worker", "status": "active",
        "capabilities": ["llm"],
    }

    assert _find_image_node() == "corsair"


def test_find_image_node_skips_inactive():
    from routers.image import _find_image_node

    _reset()
    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "error",
        "capabilities": ["image_gen"],
    }

    assert _find_image_node() is None


def test_find_image_node_no_image_gen():
    from routers.image import _find_image_node

    _reset()
    cluster.CLUSTER_NODES["jarvis"] = {
        "name": "jarvis", "type": "worker", "status": "active",
        "capabilities": ["llm"],
    }

    assert _find_image_node() is None


# ── 9. POST /api/image/generate — no node available ─────────────────────


def test_image_generate_no_node_503(tmp_path):
    _reset()
    with make_client(tmp_path) as client:
        headers = _admin_headers(client)
        resp = client.post("/api/image/generate", json={"prompt": "test"}, headers=headers)
        assert resp.status_code == 503
        assert "No image generation service" in resp.json()["detail"]


# ── 10. POST /api/image/generate — empty prompt ─────────────────────────


def test_image_generate_empty_prompt_400(tmp_path):
    _reset()
    with make_client(tmp_path) as client:
        headers = _admin_headers(client)
        resp = client.post("/api/image/generate", json={"prompt": ""}, headers=headers)
        assert resp.status_code == 400
        assert "Prompt is required" in resp.json()["detail"]


# ── 11. POST /api/image/generate — happy path ──────────────────────────


def test_image_generate_happy_path(tmp_path, monkeypatch):
    _reset()
    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }

    fake_png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
    fake_b64 = base64.b64encode(fake_png).decode()

    async def fake_request_image_generate(**kwargs):
        return fake_b64

    monkeypatch.setattr("routers.image.request_image_generate", fake_request_image_generate)

    with make_client(tmp_path) as client:
        headers = _admin_headers(client)
        resp = client.post("/api/image/generate", json={
            "prompt": "a red dragon",
            "width": 512,
            "height": 512,
        }, headers=headers)
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "image/png"
        assert resp.content == fake_png


# ── 12. POST /api/image/generate — generation failed ───────────────────


def test_image_generate_timeout_504(tmp_path, monkeypatch):
    _reset()
    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
    }

    async def fake_request_image_generate(**kwargs):
        return None

    monkeypatch.setattr("routers.image.request_image_generate", fake_request_image_generate)

    with make_client(tmp_path) as client:
        headers = _admin_headers(client)
        resp = client.post("/api/image/generate", json={"prompt": "test"}, headers=headers)
        assert resp.status_code == 504


# ── 13. GET /api/image/status — available ──────────────────────────────


def test_image_status_available(tmp_path):
    _reset()
    cluster.CLUSTER_NODES["corsair"] = {
        "name": "corsair", "type": "worker", "status": "active",
        "capabilities": ["image_gen"],
        "load": {"gpu_pct": 30},
        "last_seen": "2026-07-27T00:00:00Z",
    }

    with make_client(tmp_path) as client:
        headers = _guest_headers(client)
        resp = client.get("/api/image/status", headers=headers)
        assert resp.status_code == 200
        data = resp.json()
        assert data["available"] is True
        assert len(data["nodes"]) == 1
        assert data["nodes"][0]["name"] == "corsair"


# ── 14. GET /api/image/status — no nodes ───────────────────────────────


def test_image_status_unavailable(tmp_path):
    _reset()
    with make_client(tmp_path) as client:
        headers = _guest_headers(client)
        resp = client.get("/api/image/status", headers=headers)
        assert resp.status_code == 200
        data = resp.json()
        assert data["available"] is False
        assert len(data["nodes"]) == 0


# ── 15. node_agent.detect_capabilities — ComfyUI present ───────────────


def test_detect_capabilities_comfyui_present(monkeypatch):
    cfg = agent.AgentConfig()
    cfg.comfyui_port = 8188

    monkeypatch.setattr(agent, "HAS_HTTPX", True)

    def fake_get(url, timeout=3):
        if "system_stats" in url:
            class R:
                status_code = 200
            return R()
        raise httpx.ConnectError("refused")

    monkeypatch.setattr(httpx, "get", fake_get)

    caps = agent.detect_capabilities(cfg)
    assert "image_gen" in caps
    assert "llm" in caps


# ── 16. node_agent.detect_capabilities — ComfyUI absent ────────────────


def test_detect_capabilities_comfyui_absent(monkeypatch):
    cfg = agent.AgentConfig()
    cfg.comfyui_port = 8188

    monkeypatch.setattr(agent, "HAS_HTTPX", True)
    monkeypatch.setattr(httpx, "get", lambda url, timeout=3: (_ for _ in ()).throw(httpx.ConnectError("refused")))

    caps = agent.detect_capabilities(cfg)
    assert "image_gen" not in caps
    assert "llm" in caps


# ── 17. node_agent.detect_capabilities — httpx not installed ────────────


def test_detect_capabilities_no_httpx(monkeypatch):
    cfg = agent.AgentConfig()
    monkeypatch.setattr(agent, "HAS_HTTPX", False)

    caps = agent.detect_capabilities(cfg)
    assert "image_gen" not in caps


# ── 18. node_agent.handle_image_generate — success ─────────────────────


def test_node_agent_handle_image_generate_success(monkeypatch):
    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
    monkeypatch.setattr(agent, "HAS_HTTPX", True)

    cfg = agent.AgentConfig()
    cfg.node_name = "corsair"
    cfg.comfyui_port = 8188

    fake_png = b"\x89PNG" + b"\x00" * 50
    fake_b64 = base64.b64encode(fake_png).decode()

    async def fake_comfyui_generate(*a, **kw):
        return fake_b64

    monkeypatch.setattr(agent, "_comfyui_generate", fake_comfyui_generate)

    channel = FakeChannel()
    system_ex = FakeExchange("jc.system")
    channel.exchanges["jc.system"] = system_ex

    asyncio.run(agent.handle_image_generate(
        cfg, channel, (FakeExchange(), system_ex),
        FakeMsg({
            "request_id": "req-789",
            "prompt": "a castle",
            "negative_prompt": "",
            "width": 1024,
            "height": 1024,
            "steps": 20,
            "seed": 42,
            "model": "",
        }),
    ))

    assert len(system_ex.published) == 1
    msg, rk = system_ex.published[0]
    assert rk == "node.corsair.image_generated"
    payload = json.loads(msg.body)
    assert payload["type"] == "image_generated"
    assert payload["request_id"] == "req-789"
    assert payload["image_base64"] == fake_b64


# ── 19. node_agent.handle_image_generate — failure ─────────────────────


def test_node_agent_handle_image_generate_failure(monkeypatch):
    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
    monkeypatch.setattr(agent, "HAS_HTTPX", True)

    cfg = agent.AgentConfig()
    cfg.node_name = "corsair"

    async def fake_comfyui_generate(*a, **kw):
        raise RuntimeError("ComfyUI crashed")

    monkeypatch.setattr(agent, "_comfyui_generate", fake_comfyui_generate)

    channel = FakeChannel()
    system_ex = FakeExchange("jc.system")
    channel.exchanges["jc.system"] = system_ex

    asyncio.run(agent.handle_image_generate(
        cfg, channel, (FakeExchange(), system_ex),
        FakeMsg({"request_id": "req-fail", "prompt": "test"}),
    ))

    assert len(system_ex.published) == 1
    msg, rk = system_ex.published[0]
    assert rk == "node.corsair.image_failed"
    payload = json.loads(msg.body)
    assert payload["type"] == "image_failed"
    assert "ComfyUI crashed" in payload["error"]


# ── 20. node_agent.handle_image_generate — empty prompt ────────────────


def test_node_agent_handle_image_generate_empty_prompt(monkeypatch):
    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)

    cfg = agent.AgentConfig()
    cfg.node_name = "corsair"

    channel = FakeChannel()
    system_ex = FakeExchange("jc.system")

    asyncio.run(agent.handle_image_generate(
        cfg, channel, (FakeExchange(), system_ex),
        FakeMsg({"request_id": "req-x", "prompt": ""}),
    ))

    assert len(system_ex.published) == 0


# ── 21. hardware.py — ComfyUI reachable ────────────────────────────────


def test_assess_hardware_comfyui_reachable(tmp_path, monkeypatch):
    hardware.HARDWARE_STATE_PATH = tmp_path / "hw.json"
    monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})())
    monkeypatch.setattr(psutil, "cpu_count", lambda: 8)

    class MockProc:
        returncode = 1
        stdout = ""

    monkeypatch.setattr(hardware.subprocess, "run", lambda cmd, **kw: MockProc())

    async def mock_get(self, url, *args, **kwargs):
        class R:
            status_code = 200
            def json(self):
                if "CheckpointLoaderSimple" in url:
                    return {"CheckpointLoaderSimple": {"input": {"required": {"ckpt_name": [["model.safetensors", "other.ckpt"]]}}}}
                return {}
            def raise_for_status(self):
                pass
        if "8188" in url:
            return R()
        if "v1/models" in url:
            class R2:
                status_code = 200
                def json(self):
                    return {"data": []}
            return R2()
        if "6333" in url:
            class R3:
                status_code = 200
                def json(self):
                    return {"result": {"collections": []}}
            return R3()
        if "8888" in url:
            class R4:
                status_code = 200
            return R4()
        class R5:
            status_code = 200
            def json(self):
                return {}
        return R5()

    monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)

    state = asyncio.run(hardware.assess_hardware())
    assert state["comfyui_reachable"] is True
    assert "model.safetensors" in state["comfyui_models"]


# ── 22. hardware.py — ComfyUI unreachable ──────────────────────────────


def test_assess_hardware_comfyui_unreachable(tmp_path, monkeypatch):
    hardware.HARDWARE_STATE_PATH = tmp_path / "hw.json"
    monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})())
    monkeypatch.setattr(psutil, "cpu_count", lambda: 8)

    class MockProc:
        returncode = 1
        stdout = ""

    monkeypatch.setattr(hardware.subprocess, "run", lambda cmd, **kw: MockProc())

    async def mock_get(self, url, *args, **kwargs):
        if "8188" in url:
            raise httpx.ConnectError("refused")
        if "v1/models" in url:
            class R2:
                status_code = 200
                def json(self):
                    return {"data": []}
            return R2()
        if "6333" in url:
            class R3:
                status_code = 200
                def json(self):
                    return {"result": {"collections": []}}
            return R3()
        if "8888" in url:
            class R4:
                status_code = 200
            return R4()
        class R5:
            status_code = 200
            def json(self):
                return {}
        return R5()

    monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)

    state = asyncio.run(hardware.assess_hardware())
    assert state["comfyui_reachable"] is False
    assert state["comfyui_models"] == []


# ── 23. node_agent config reads comfyui_port ───────────────────────────


def test_config_from_ini_comfyui_port(tmp_path):
    ini = tmp_path / "caic-node-agent.conf"
    ini.write_text(
        "[agent]\n"
        "node_name = corsair\n"
        "capabilities = llm,image_gen\n"
        "comfyui_port = 8188\n"
    )
    cfg = agent.AgentConfig.from_ini(str(ini))
    assert cfg.comfyui_port == 8188
    assert "image_gen" in cfg.capabilities


def test_config_from_ini_comfyui_port_default():
    cfg = agent.AgentConfig()
    assert cfg.comfyui_port == 8188


# ── 24. SUBSCRIBE_TABLE includes image gen handlers ────────────────────


def test_subscribe_table_includes_image_handlers():
    routing_keys = [rks for _, rks, _ in cluster.SUBSCRIBE_TABLE]
    all_keys = [rk for rks in routing_keys for rk in rks]
    assert "node.*.image_generated" in all_keys
    assert "node.*.image_failed" in all_keys