"""Handler smoke tests for backend/handlers/cinema_scene.py.

The base-model call is MOCKED — no network, no real generation. We assert:
- each shot's image maps to that shot's dynamic output port (id derived from
  shot.id);
- a shot whose base call raises is isolated: it gets status 'error' while the
  other shots still complete (scene completes partially);
- the License guard substitutes the commercial-OK default base for FLUX.1-dev;
- color + look stages run on the base output (real deterministic pillars).
"""

from __future__ import annotations

from pathlib import Path
from typing import Any
from uuid import uuid4

import numpy as np
import pytest
from PIL import Image

import handlers.cinema_scene as cinema_scene
from handlers.cinema_scene import (
    _guard_base_model,
    _output_port_id,
    handle_cinema_scene,
)
from models.graph import GraphNode, PortValueDict
from services.output import OUTPUT_ROOT


# ---------- helpers ----------


def _write_base_image(rgb: tuple[int, int, int] = (90, 120, 200)) -> str:
    """Write a small solid image under OUTPUT_ROOT and return its served URL —
    standing in for a base model's generated output."""
    run_dir = OUTPUT_ROOT / "cinema-scene-test"
    run_dir.mkdir(parents=True, exist_ok=True)
    arr = np.zeros((24, 24, 3), dtype=np.uint8)
    arr[:, :] = rgb
    out_path = run_dir / f"{uuid4().hex[:12]}.png"
    Image.fromarray(arr, mode="RGB").save(out_path, format="PNG")
    rel = out_path.resolve().relative_to(OUTPUT_ROOT.resolve())
    return f"/api/outputs/{rel}"


def _scene_node(scene: dict[str, Any]) -> GraphNode:
    return GraphNode(id="scene1", definitionId="cinema-scene", params={"scene": scene})


def _patch_base(monkeypatch, behavior) -> list[dict[str, Any]]:
    """Patch get_handler_registry so the base model handler is `behavior`.

    Returns a list that records each base-handler invocation's (node, inputs)
    for assertions. `behavior(node, inputs, api_keys)` is the async stand-in.
    """
    calls: list[dict[str, Any]] = []

    async def fake_base_handler(node, inputs, api_keys):
        calls.append({"node": node, "inputs": inputs})
        return await behavior(node, inputs, api_keys)

    def fake_registry(emit=None):
        # Every base model id resolves to our stand-in.
        class _AnyKeyRegistry(dict):
            def get(self, _key, _default=None):
                return fake_base_handler

        return _AnyKeyRegistry()

    monkeypatch.setattr(cinema_scene, "get_handler_registry", fake_registry, raising=False)
    # handle_cinema_scene does `from execution.sync_runner import get_handler_registry`
    # at call time, so patch the source module too.
    import execution.sync_runner as sync_runner
    monkeypatch.setattr(sync_runner, "get_handler_registry", fake_registry)
    return calls


# ---------- pure helpers ----------


def test_output_port_id_derives_from_shot_id() -> None:
    assert _output_port_id("abc") == "shot_abc"


def test_license_guard_substitutes_flux_dev() -> None:
    assert _guard_base_model({"model": "flux-1-dev"}) == "seedream-4-5"
    assert _guard_base_model({"model": "FLUX.1-dev"}) == "seedream-4-5"
    assert _guard_base_model({}) == "seedream-4-5"
    # A commercial-OK model is left alone.
    assert _guard_base_model({"model": "nano-banana"}) == "nano-banana"


# ---------- per-shot port mapping ----------


@pytest.mark.asyncio
async def test_shots_map_to_dynamic_output_ports(monkeypatch) -> None:
    async def ok(node, inputs, api_keys):
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, ok)

    scene = {
        "version": 1,
        "base": {"model": "nano-banana"},
        "aspectRatio": "16:9",
        "shots": [
            {"id": "s1", "prompt": "a stepling at dawn"},
            {"id": "s2", "prompt": "a stepling at dusk"},
        ],
    }
    result = await handle_cinema_scene(
        _scene_node(scene), inputs={}, api_keys={"FAL_KEY": "x"}, emit=None
    )

    assert set(result.keys()) == {"shot_s1", "shot_s2"}
    for port_id in ("shot_s1", "shot_s2"):
        assert result[port_id]["type"] == "Image"
        assert isinstance(result[port_id]["value"], str)
        assert result[port_id]["value"].startswith("/api/outputs/")

    # Per-shot status is recorded back on the spec.
    shots = scene["shots"]
    assert shots[0]["output"]["status"] == "done"
    assert shots[1]["output"]["status"] == "done"
    assert shots[0]["output"]["hash"]


# ---------- per-shot isolation ----------


@pytest.mark.asyncio
async def test_failing_shot_is_isolated(monkeypatch) -> None:
    async def flaky(node, inputs, api_keys):
        prompt = inputs.get("prompt")
        if prompt and "boom" in str(prompt.value):
            raise RuntimeError("base model exploded")
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, flaky)

    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},
        "aspectRatio": "2.39:1",
        "shots": [
            {"id": "good1", "prompt": "calm scene"},
            {"id": "bad", "prompt": "boom scene"},
            {"id": "good2", "prompt": "another calm scene"},
        ],
    }
    result = await handle_cinema_scene(
        _scene_node(scene), inputs={}, api_keys={"FAL_KEY": "x"}, emit=None
    )

    # All three ports present — scene completed partially.
    assert set(result.keys()) == {"shot_good1", "shot_bad", "shot_good2"}
    assert result["shot_good1"]["value"].startswith("/api/outputs/")
    assert result["shot_good2"]["value"].startswith("/api/outputs/")
    assert result["shot_bad"]["value"] is None

    shots = {s["id"]: s for s in scene["shots"]}
    assert shots["good1"]["output"]["status"] == "done"
    assert shots["good2"]["output"]["status"] == "done"
    assert shots["bad"]["output"]["status"] == "error"
    assert "exploded" in shots["bad"]["output"]["error"]


# ---------- color + look stages run on the base output ----------


@pytest.mark.asyncio
async def test_color_and_look_modify_base_output(monkeypatch) -> None:
    base_url = _write_base_image(rgb=(128, 128, 128))

    async def ok(node, inputs, api_keys):
        return {"image": {"type": "Image", "value": base_url}}

    _patch_base(monkeypatch, ok)

    # A strong warm palette + vignette should change the pixels vs the base.
    scene = {
        "version": 1,
        "base": {"model": "nano-banana"},
        "aspectRatio": "16:9",
        "palette": {"swatches": ["#d98841", "#c46a28"], "strength": 0.9, "method": "reinhard"},
        "look": {"vignette": 0.6, "contrast": 0.3},
        "shots": [{"id": "only", "prompt": "warm-graded shot"}],
    }
    result = await handle_cinema_scene(
        _scene_node(scene), inputs={}, api_keys={"FAL_KEY": "x"}, emit=None
    )

    out_url = result["shot_only"]["value"]
    rel = out_url[len("/api/outputs/"):]
    out_path = (OUTPUT_ROOT / rel).resolve()
    assert out_path.exists()

    base_rel = base_url[len("/api/outputs/"):]
    base_arr = np.asarray(Image.open((OUTPUT_ROOT / base_rel).resolve()).convert("RGB"), dtype=np.int64)
    out_arr = np.asarray(Image.open(out_path).convert("RGB"), dtype=np.int64)
    # The graded output must differ from the flat base.
    assert out_arr.shape == base_arr.shape
    assert int(np.abs(out_arr - base_arr).sum()) > 0


# ---------- character refs flow into the base call ----------


@pytest.mark.asyncio
async def test_character_refs_passed_to_base(monkeypatch) -> None:
    ref_url = _write_base_image(rgb=(10, 200, 10))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["images"] = inputs.get("images")
        captured["prompt"] = inputs.get("prompt")
        captured["aspectRatio"] = node.params.get("aspectRatio")
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    scene = {
        "version": 1,
        "base": {"model": "nano-banana"},
        "character": {"refImageUrls": [ref_url], "strength": 0.8},
        "prompt": "cinematic",
        "aspectRatio": "4:5",
        "shots": [{"id": "s1", "prompt": "close-up", "refImageUrls": []}],
    }
    await handle_cinema_scene(
        _scene_node(scene), inputs={}, api_keys={"FAL_KEY": "x"}, emit=None
    )

    images_port = captured["images"]
    assert images_port is not None
    assert ref_url in images_port.value
    # Prompt combines scene prompt + shot prompt.
    assert captured["prompt"].value == "cinematic close-up"
    assert captured["aspectRatio"] == "4:5"


# ---------- typed Character bundle input expands verbatim ----------


_BUNDLE_TRAIT = (
    "freckled olive skin, deep-set hazel eyes, asymmetric undercut (left side), "
    "a faint scar above the RIGHT brow; 1.78m; wears a worn leather aviator jacket"
)


def _bundle(views: list[str], seed: int = 84) -> dict[str, Any]:
    return {
        "characterId": "char_1",
        "name": "Iris Vane",
        "referenceViews": list(views),
        "frozenTraitString": _BUNDLE_TRAIT,
        "seed": seed,
        "consistencyStrength": 0.65,
    }


@pytest.mark.asyncio
async def test_character_bundle_expands_verbatim(monkeypatch) -> None:
    """A typed `character` input expands each shot's base call: trait string
    prepended verbatim, referenceViews first then per-use refs, seed injected."""
    v1 = _write_base_image(rgb=(10, 200, 10))
    v2 = _write_base_image(rgb=(200, 10, 10))
    v3 = _write_base_image(rgb=(10, 10, 200))
    shot_ref = _write_base_image(rgb=(50, 50, 50))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["images"] = inputs.get("images")
        captured["prompt"] = inputs.get("prompt")
        captured["seed"] = node.params.get("seed")
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    bundle = _bundle([v1, v2, v3], seed=84)
    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},  # maxRefs = 10, comfortably above 4
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot", "refImageUrls": [shot_ref]}],
    }
    await handle_cinema_scene(
        _scene_node(scene),
        inputs={"character": PortValueDict(type="Character", value=bundle)},
        api_keys={"FAL_KEY": "x"},
        emit=None,
    )

    # Trait string prepended VERBATIM, then the base (scene + shot) prompt.
    assert captured["prompt"].value == f"{_BUNDLE_TRAIT}. a forest at dawn wide shot"
    # referenceViews FIRST (stored order), then the per-use shot ref.
    assert captured["images"].value == [v1, v2, v3, shot_ref]
    # Seed injected into the base node params.
    assert captured["seed"] == 84


@pytest.mark.asyncio
async def test_character_bundle_exceeding_max_refs_raises(monkeypatch) -> None:
    """A bundle whose view count exceeds the base model's reference cap fails
    fast for the whole scene with the clear anti-FLORA capability error."""

    async def ok(node, inputs, api_keys):
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, ok)

    # 12 views > seedream-4-5's cap of 10 -> raise up-front (scene-level).
    views = [_write_base_image() for _ in range(12)]
    bundle = _bundle(views, seed=7)
    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot"}],
    }
    with pytest.raises(ValueError) as exc:
        await handle_cinema_scene(
            _scene_node(scene),
            inputs={"character": PortValueDict(type="Character", value=bundle)},
            api_keys={"FAL_KEY": "x"},
            emit=None,
        )
    msg = str(exc.value)
    assert "10" in msg  # the cap
    assert "12" in msg  # the required ref count
    assert "reference" in msg.lower()


@pytest.mark.asyncio
async def test_character_bundle_on_nano_banana_default_does_not_raise(monkeypatch) -> None:
    """The spec §5 headline default base — `nano-banana` (its model enum
    defaults to Nano Banana 2, 14 refs) — must ACCEPT a real multi-view Character
    bundle and expand it verbatim, NOT reject it with the capability error.

    Regression guard for the dead-key bug: `MODEL_MAX_REFS` used to key the 14
    cap on the unreachable `nano-banana-2` def id, so a `nano-banana` base fell
    through to DEFAULT_MAX_REFS (1) and rejected the spec's own default bundle.
    """
    v1 = _write_base_image(rgb=(10, 200, 10))
    v2 = _write_base_image(rgb=(200, 10, 10))
    v3 = _write_base_image(rgb=(10, 10, 200))
    shot_ref = _write_base_image(rgb=(50, 50, 50))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["images"] = inputs.get("images")
        captured["prompt"] = inputs.get("prompt")
        captured["seed"] = node.params.get("seed")
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    bundle = _bundle([v1, v2, v3], seed=84)
    scene = {
        "version": 1,
        "base": {"model": "nano-banana"},  # defaults to NB2 -> maxRefs = 14
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot", "refImageUrls": [shot_ref]}],
    }
    outputs = await handle_cinema_scene(
        _scene_node(scene),
        inputs={"character": PortValueDict(type="Character", value=bundle)},
        api_keys={"FAL_KEY": "x"},
        emit=None,
    )

    # Shot completed (no capability error) — output port carries an image URL.
    assert outputs["shot_s1"]["value"] is not None
    assert scene["shots"][0]["output"]["status"] == "done"
    # Trait string prepended VERBATIM, then the base (scene + shot) prompt.
    assert captured["prompt"].value == f"{_BUNDLE_TRAIT}. a forest at dawn wide shot"
    # referenceViews FIRST (stored order), then the per-use shot ref.
    assert captured["images"].value == [v1, v2, v3, shot_ref]
    # Seed injected into the base node params.
    assert captured["seed"] == 84


# ---------- per-use override layer flows through to the base call ----------


@pytest.mark.asyncio
async def test_override_prompt_and_refs_flow_to_base_call(monkeypatch) -> None:
    """A bundle carrying the per-use override layer (overridePrompt /
    overrideRefs) expands each shot's base call: the override prompt trails the
    base prompt, and the node-level overrideRefs slot in after referenceViews and
    before the scene/shot refs."""
    v1 = _write_base_image(rgb=(10, 200, 10))
    v2 = _write_base_image(rgb=(200, 10, 10))
    node_ref = _write_base_image(rgb=(120, 120, 0))
    shot_ref = _write_base_image(rgb=(50, 50, 50))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["images"] = inputs.get("images")
        captured["prompt"] = inputs.get("prompt")
        captured["params"] = dict(node.params)
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    bundle = _bundle([v1, v2], seed=84)
    bundle["overridePrompt"] = "low-angle, dramatic rim light"
    bundle["overrideRefs"] = [node_ref]
    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},  # maxRefs = 10
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot", "refImageUrls": [shot_ref]}],
    }
    await handle_cinema_scene(
        _scene_node(scene),
        inputs={"character": PortValueDict(type="Character", value=bundle)},
        api_keys={"FAL_KEY": "x"},
        emit=None,
    )

    # PROMPT ORDER: "{trait}. {base}. {override}".
    assert captured["prompt"].value == (
        f"{_BUNDLE_TRAIT}. a forest at dawn wide shot. low-angle, dramatic rim light"
    )
    # referenceViews FIRST, then node overrideRefs, then the scene/shot ref.
    assert captured["images"].value == [v1, v2, node_ref, shot_ref]


# ---------- strength applied HONESTLY: only where the base supports it ----------


@pytest.mark.asyncio
async def test_strength_omitted_for_base_without_adherence_param(monkeypatch) -> None:
    """The HONEST behaviour for v1 reference-edit bases: seedream-4-5 has NO
    IP-adherence knob, so the effective strength is carried in the bundle but
    NOT injected into the base call's params (no silent no-op)."""
    v1 = _write_base_image(rgb=(10, 200, 10))
    v2 = _write_base_image(rgb=(200, 10, 10))
    v3 = _write_base_image(rgb=(10, 10, 200))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["params"] = dict(node.params)
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    bundle = _bundle([v1, v2, v3], seed=84)
    bundle["strengthOverride"] = 0.9  # would-be strength value
    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot"}],
    }
    await handle_cinema_scene(
        _scene_node(scene),
        inputs={"character": PortValueDict(type="Character", value=bundle)},
        api_keys={"FAL_KEY": "x"},
        emit=None,
    )

    params = captured["params"]
    # No fabricated strength key reached the base for a model without the knob.
    assert "strength" not in params
    assert "image_prompt_strength" not in params
    assert "consistencyStrength" not in params
    assert "strengthOverride" not in params
    # The seed (a real, supported param) still flows through.
    assert params.get("seed") == 84


@pytest.mark.asyncio
async def test_strength_injected_for_base_with_adherence_param(monkeypatch) -> None:
    """When a base model DOES expose an IP-adherence param (mapped in
    MODEL_STRENGTH_PARAM), the effective strength is injected under that key.

    v1 has no such base, so we register a hypothetical one for this test to
    prove the wiring fires when a real knob exists (e.g. v2 trained-LoRA)."""
    import cinema.identity as identity

    monkeypatch.setitem(identity.MODEL_STRENGTH_PARAM, "seedream-4-5", "ip_strength")

    v1 = _write_base_image(rgb=(10, 200, 10))
    v2 = _write_base_image(rgb=(200, 10, 10))

    captured: dict[str, Any] = {}

    async def capture(node, inputs, api_keys):
        captured["params"] = dict(node.params)
        return {"image": {"type": "Image", "value": _write_base_image()}}

    _patch_base(monkeypatch, capture)

    bundle = _bundle([v1, v2], seed=84)  # consistencyStrength = 0.65
    scene = {
        "version": 1,
        "base": {"model": "seedream-4-5"},
        "prompt": "a forest at dawn",
        "aspectRatio": "16:9",
        "shots": [{"id": "s1", "prompt": "wide shot"}],
    }
    await handle_cinema_scene(
        _scene_node(scene),
        inputs={"character": PortValueDict(type="Character", value=bundle)},
        api_keys={"FAL_KEY": "x"},
        emit=None,
    )

    # Effective strength (no override -> consistencyStrength) under the mapped key.
    assert captured["params"]["ip_strength"] == 0.65
