"""Pure functions over LTX 2.3 mode API-format workflow templates.

Templates in `workflows/<mode>.json` are saved from ComfyUI's editor via
"Save (API Format)". Shape: `{node_id_str: {"class_type": str, "inputs": dict}}`.
This is what ComfyUI's `PromptExecutor.execute(prompt=...)` expects directly.
"""

from __future__ import annotations

import copy
import json
import pathlib
from typing import Any

WORKFLOWS_DIR = pathlib.Path(__file__).parent / "workflows"

VALID_MODES: tuple[str, ...] = ("t2v", "a2v", "i2v", "lipsync", "keyframe", "style")

# Model files the upstream HF repos have deleted, mapped to their replacements.
# The templates are machine-generated by tools/extract_modes.py from the user's
# master ComfyUI export, so patching the JSON would be undone on the next
# re-export — the remap has to live in code. load_template is the only reader of
# workflows/*.json, which puts it upstream of both consumers: the dict handed to
# PromptExecutor and models.walk_workflow_for_models' download pass.
RETIRED_MODEL_RENAMES: dict[str, str] = {
    # Lightricks deleted x2 spatial upscaler v1.0 on 2026-08-10 (v1.1 is the
    # long-video hotfix). Loading the retired name 404s and aborts the graph
    # right after sampling, leaving the user with an empty output.
    "ltx-2.3-spatial-upscaler-x2-1.0.safetensors": "ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
}


def _rename_retired_models(value: Any) -> Any:
    """Return *value* with every retired model filename swapped for its successor."""
    if isinstance(value, str):
        return RETIRED_MODEL_RENAMES.get(value, value)
    if isinstance(value, dict):
        return {k: _rename_retired_models(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_rename_retired_models(v) for v in value]
    return value


def load_template(mode: str) -> dict[str, Any]:
    """Load a fresh, independent copy of the named mode's API workflow template."""
    if mode not in VALID_MODES:
        raise ValueError(f"unknown mode {mode!r}; expected one of {VALID_MODES}")
    path = WORKFLOWS_DIR / f"{mode}.json"
    template = copy.deepcopy(json.loads(path.read_text()))
    return _rename_retired_models(template)


def set_input(workflow: dict[str, Any], node_id: int | str, field: str, value: Any) -> None:
    """Patch a node's input field in place.

    For API-format workflows, each node has an `inputs` dict keyed by field name.
    `node_id` is the dict key (string for top-level, "<inst>:<inner>" for
    subgraph-internal). `field` is an entry name in `inputs`.

    Args:
        workflow: API-format workflow dict (mapping id → {class_type, inputs}).
        node_id: Dict key of the target node.
        field: Name of the input field to set.
        value: New value (literal, or `[src_id, src_slot]` link form).

    Raises:
        KeyError: If the node doesn't exist in the workflow.
    """
    nid = str(node_id)
    if nid not in workflow:
        raise KeyError(f"node id {nid!r} not found in workflow")
    inputs = workflow[nid].setdefault("inputs", {})
    inputs[field] = value