"""ComfyUIAdapter — local Krea 2 Turbo image generation.

No network/server. The three HTTP seams (_post_json, _get_json, _get_bytes) are
monkeypatched; the workflow builder and direct_submit_image control flow run for
real. Also asserts the model routes to the comfyui provider via resolve_adapter.
"""

from __future__ import annotations

import pytest

from recoil.execution.providers.base import UnifiedVideoPayload
from recoil.execution.providers.comfyui import ComfyUIAdapter


def test_workflow_has_krea2_graph_with_rebalance():
    a = ComfyUIAdapter()
    g = a._build_workflow(
        "a bloodied prisoner", seed=42, width=768, height=1344,
        rebalance_multiplier=4.0,
        rebalance_weights="1,1,1,1,1,1,1,2.5,5,1.1,4,1", filename_prefix="EP1_SH1",
    )
    # Loaders point at the bf16 Krea 2 assets + krea2 CLIP type.
    assert g["10"]["inputs"]["unet_name"] == "krea2_turbo_bf16.safetensors"
    assert g["11"]["inputs"]["type"] == "krea2"
    assert g["12"]["inputs"]["vae_name"] == "qwen_image_vae.safetensors"
    # Rebalance node present and feeds the KSampler positive.
    assert g["13"]["class_type"] == "ConditioningKrea2Rebalance"
    assert g["13"]["inputs"]["multiplier"] == 4.0
    assert g["3"]["inputs"]["positive"] == ["13", 0]
    # Turbo sampler settings.
    assert g["3"]["inputs"]["steps"] == 8
    assert g["3"]["inputs"]["cfg"] == 1.0
    assert g["3"]["inputs"]["sampler_name"] == "euler"
    assert g["5"]["inputs"]["width"] == 768 and g["5"]["inputs"]["height"] == 1344


def test_workflow_omits_rebalance_at_multiplier_one():
    g = ComfyUIAdapter()._build_workflow(
        "x", seed=1, width=1024, height=1024,
        rebalance_multiplier=1.0, rebalance_weights="", filename_prefix="p",
    )
    assert "13" not in g
    assert g["3"]["inputs"]["positive"] == ["6", 0]


def test_negative_is_always_zeroed_out():
    # cfg is fixed at 1.0 → negative is inert; always ConditioningZeroOut.
    g = ComfyUIAdapter()._build_workflow(
        "x", seed=1, width=1024, height=1024,
        rebalance_multiplier=4.0, rebalance_weights="w", filename_prefix="p",
    )
    assert g["7"]["class_type"] == "ConditioningZeroOut"
    assert g["3"]["inputs"]["cfg"] == 1.0


def test_does_not_advertise_unhonored_capabilities():
    caps = ComfyUIAdapter().capabilities
    assert caps["t2v"] and caps["resolution_720p"]
    assert not caps["negative_prompt"]
    assert not caps["resolution_480p"] and not caps["resolution_1080p"]


def _patch_http(monkeypatch, *, history_seq, image_bytes=b"PNGDATA"):
    posted = {}

    def _post(self, path, body):  # noqa: ANN001
        posted["workflow"] = body["prompt"]
        return {"prompt_id": "pid_123"}

    seq = iter(history_seq)

    def _get_json(self, path):  # noqa: ANN001
        return next(seq)

    def _get_bytes(self, path):  # noqa: ANN001
        posted["view_path"] = path
        return image_bytes

    monkeypatch.setattr(ComfyUIAdapter, "_post_json", _post)
    monkeypatch.setattr(ComfyUIAdapter, "_get_json", _get_json)
    monkeypatch.setattr(ComfyUIAdapter, "_get_bytes", _get_bytes)
    monkeypatch.setattr("time.sleep", lambda *_a, **_k: None)
    return posted


def test_direct_submit_image_happy_path(monkeypatch):
    posted = _patch_http(
        monkeypatch,
        history_seq=[
            {},  # poll #1: prompt_id not in history yet
            {"pid_123": {  # poll #2: done
                "status": {"status_str": "success"},
                "outputs": {"9": {"images": [
                    {"filename": "EP1_SH1_00001_.png", "subfolder": "", "type": "output"}]}},
            }},
        ],
    )
    a = ComfyUIAdapter()
    out = a.direct_submit_image(UnifiedVideoPayload(
        prompt="bloodied prisoner", aspect_ratio="9:16",
        shot_id="EP1_SH1", model_id="krea-2-turbo"))
    assert out["image_bytes"] == b"PNGDATA"
    assert out["cost_usd"] == 0.0
    assert out["native_id"] == "pid_123"
    # 9:16 → 768x1344, rebalance default present.
    wf = posted["workflow"]
    assert wf["5"]["inputs"]["width"] == 768
    assert "13" in wf
    assert "filename=EP1_SH1_00001_.png" in posted["view_path"]


def test_direct_submit_image_raises_on_execution_error(monkeypatch):
    _patch_http(
        monkeypatch,
        history_seq=[
            {"pid_123": {
                "status": {"status_str": "error", "messages": [
                    ["execution_error", {"node_type": "KSampler",
                                          "exception_message": "boom"}]]},
                "outputs": {},
            }},
        ],
    )
    with pytest.raises(RuntimeError, match="KSampler: boom"):
        ComfyUIAdapter().direct_submit_image(
            UnifiedVideoPayload(prompt="x", model_id="krea-2-turbo"))


def test_routes_to_comfyui_provider():
    from recoil.execution.providers import resolve_adapter
    from recoil.execution.providers.registry import reset_caches_for_tests
    reset_caches_for_tests()
    p = UnifiedVideoPayload(prompt="x", resolution="720p", model_id="krea-2-turbo")
    adapter, tier = resolve_adapter("krea-2-turbo", p)
    assert adapter.provider_id == "comfyui"