"""comfyui_generate_asset_pack_oneshot — set COHERENTE de assets 2D de un solo tiro.

One-shot del grupo `gamedev-2d`: recibe una *spec de pack* (lista de assets, cada
uno con su `kind` y `subject`) y genera el set entero compartiendo el MISMO
checkpoint, el MISMO LoRA de estilo, el MISMO estilo común y una semilla derivada
(`seed = base_seed + indice`). Eso es lo que hace que un sprite, sus iconos, sus
tiles y su UI "peguen" entre sí: misma mano, misma paleta, mismo modelo.

Promueve a UNA sola llamada la secuencia que hoy se repite a mano (issue 0087):
elegir N builders atómicos de `gamedev-2d`, pasarles el mismo estilo/checkpoint/LoRA
y encolar/esperar/descargar cada uno. El registry no crece inflando builders —
crece promoviendo esta composición repetida a un pipeline.

Compone funciones del registry:

    comfyui_build_<kind>_workflow_py_ml   (un builder atómico por `kind`, PUROS)
    comfyui_submit_workflow_py_ml          (POST /prompt)
    comfyui_wait_result_py_ml              (poll /history)
    comfyui_fetch_output_image_py_ml       (GET /view -> disco)
    comfyui_export_asset_to_godot_py_pipelines  (export opcional a un proyecto Godot)

Coherencia (lo que comparten TODOS los assets del pack):
    - checkpoint  -> mismo modelo base.
    - lora        -> mismo LoRA de estilo (solo en builders con un param `lora` genérico).
    - style       -> texto de estilo común, concatenado al `subject` de cada asset,
                     de modo que el `style` categórico propio de cada builder
                     ("game icon, clean, centered", etc.) se conserva.
    - base_seed   -> seed = base_seed + indice (reproducible y a la vez variado).

Pipeline impuro: red (HTTP a ComfyUI) + escritura en disco + (export) subprocess.
"""
from __future__ import annotations

import inspect
import os
import sys
import tempfile

_FUNCTIONS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _FUNCTIONS_ROOT not in sys.path:
    sys.path.insert(0, _FUNCTIONS_ROOT)

# --- builders atómicos del grupo gamedev-2d (puros) ---
from ml.comfyui_build_achievement_badge_workflow import comfyui_build_achievement_badge_workflow
from ml.comfyui_build_card_art_workflow import comfyui_build_card_art_workflow
from ml.comfyui_build_decal_overlay_workflow import comfyui_build_decal_overlay_workflow
from ml.comfyui_build_dialogue_box_workflow import comfyui_build_dialogue_box_workflow
from ml.comfyui_build_emote_workflow import comfyui_build_emote_workflow
from ml.comfyui_build_enemy_creature_workflow import comfyui_build_enemy_creature_workflow
from ml.comfyui_build_foliage_set_workflow import comfyui_build_foliage_set_workflow
from ml.comfyui_build_item_icon_workflow import comfyui_build_item_icon_workflow
from ml.comfyui_build_parallax_background_workflow import comfyui_build_parallax_background_workflow
from ml.comfyui_build_particle_texture_workflow import comfyui_build_particle_texture_workflow
from ml.comfyui_build_portrait_avatar_workflow import comfyui_build_portrait_avatar_workflow
from ml.comfyui_build_projectile_workflow import comfyui_build_projectile_workflow
from ml.comfyui_build_prop_object_workflow import comfyui_build_prop_object_workflow
from ml.comfyui_build_rune_glyph_workflow import comfyui_build_rune_glyph_workflow
from ml.comfyui_build_seamless_tile_workflow import comfyui_build_seamless_tile_workflow
from ml.comfyui_build_skill_tree_node_workflow import comfyui_build_skill_tree_node_workflow
from ml.comfyui_build_splash_art_workflow import comfyui_build_splash_art_workflow
from ml.comfyui_build_status_effect_icon_workflow import comfyui_build_status_effect_icon_workflow
from ml.comfyui_build_structure_workflow import comfyui_build_structure_workflow
from ml.comfyui_build_title_lettering_workflow import comfyui_build_title_lettering_workflow
from ml.comfyui_build_topdown_sprite_workflow import comfyui_build_topdown_sprite_workflow
from ml.comfyui_build_trap_hazard_workflow import comfyui_build_trap_hazard_workflow
from ml.comfyui_build_ui_hud_workflow import comfyui_build_ui_hud_workflow
from ml.comfyui_build_vehicle_mount_workflow import comfyui_build_vehicle_mount_workflow
from ml.comfyui_build_weather_overlay_workflow import comfyui_build_weather_overlay_workflow
from ml.comfyui_build_world_map_workflow import comfyui_build_world_map_workflow

# --- transporte ComfyUI (impuro) ---
from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result

# --- export opcional a Godot ---
from pipelines.comfyui_export_asset_to_godot import comfyui_export_asset_to_godot

# Dispatch declarativo: kind del asset -> (kind Godot para export, builder atómico).
# El "kind Godot" es el bucket de carpeta/import que usa comfyui_export_asset_to_godot
# (sprite -> assets/sprites, tileset -> assets/tilesets, vfx -> assets/vfx).
_SUPPORTED: dict[str, tuple[str, object]] = {
    "item_icon":         ("sprite",  comfyui_build_item_icon_workflow),
    "enemy_creature":    ("sprite",  comfyui_build_enemy_creature_workflow),
    "prop_object":       ("sprite",  comfyui_build_prop_object_workflow),
    "ui_hud":            ("sprite",  comfyui_build_ui_hud_workflow),
    "status_effect_icon":("sprite",  comfyui_build_status_effect_icon_workflow),
    "structure":         ("sprite",  comfyui_build_structure_workflow),
    "topdown_sprite":    ("sprite",  comfyui_build_topdown_sprite_workflow),
    "projectile":        ("sprite",  comfyui_build_projectile_workflow),
    "vehicle_mount":     ("sprite",  comfyui_build_vehicle_mount_workflow),
    "trap_hazard":       ("sprite",  comfyui_build_trap_hazard_workflow),
    "foliage_set":       ("sprite",  comfyui_build_foliage_set_workflow),
    "rune_glyph":        ("sprite",  comfyui_build_rune_glyph_workflow),
    "achievement_badge": ("sprite",  comfyui_build_achievement_badge_workflow),
    "skill_tree_node":   ("sprite",  comfyui_build_skill_tree_node_workflow),
    "portrait_avatar":   ("sprite",  comfyui_build_portrait_avatar_workflow),
    "emote":             ("sprite",  comfyui_build_emote_workflow),
    "card_art":          ("sprite",  comfyui_build_card_art_workflow),
    "splash_art":        ("sprite",  comfyui_build_splash_art_workflow),
    "world_map":         ("sprite",  comfyui_build_world_map_workflow),
    "title_lettering":   ("sprite",  comfyui_build_title_lettering_workflow),
    "parallax_background":("sprite", comfyui_build_parallax_background_workflow),
    "dialogue_box":      ("sprite",  comfyui_build_dialogue_box_workflow),
    "particle_texture":  ("vfx",     comfyui_build_particle_texture_workflow),
    "decal_overlay":     ("vfx",     comfyui_build_decal_overlay_workflow),
    "weather_overlay":   ("vfx",     comfyui_build_weather_overlay_workflow),
    "seamless_tile":     ("tileset", comfyui_build_seamless_tile_workflow),
}

# Nombres de parámetro candidatos para el checkpoint, por orden de preferencia.
_CKPT_PARAMS = ("checkpoint", "ckpt_name", "ckpt")
# Claves del item del pack que NO se reenvían como kwargs al builder.
_RESERVED_ITEM_KEYS = {"kind", "subject"}


def supported_kinds() -> list[str]:
    """Lista ordenada de los `kind` de asset que el pack sabe despachar."""
    return sorted(_SUPPORTED)


def _inject_style(subject: str, style: str) -> str:
    """Estilo común al final del subject; conserva el style categórico del builder."""
    style = (style or "").strip()
    subject = (subject or "").strip()
    if not style:
        return subject
    if not subject:
        return style
    return f"{subject}, {style}"


def _build_item_workflow(
    item: dict,
    *,
    checkpoint: str,
    style: str,
    lora: str | None,
    lora_strength: float,
    base_seed: int,
    size: int | None,
    index: int,
) -> tuple[dict, str, str, int]:
    """Construye el workflow de UN asset del pack inyectando la coherencia compartida.

    PURO (no toca la red): solo invoca el builder atómico correspondiente. Útil para
    testear el dispatch sin ComfyUI.

    Returns:
        (workflow, subject_efectivo, kind_godot, seed) del asset.

    Raises:
        ValueError: si el `kind` no está soportado o falta un campo obligatorio del
            builder (p.ej. `expression` en `emote`).
    """
    kind = str(item.get("kind", "")).strip()
    if kind not in _SUPPORTED:
        raise ValueError(
            f"kind {kind!r} no soportado. Soportados: {', '.join(supported_kinds())}"
        )
    godot_kind, fn = _SUPPORTED[kind]
    params = inspect.signature(fn).parameters
    seed = base_seed + index
    subject_eff = _inject_style(str(item.get("subject", "")), style)

    # Coherencia compartida -> kwargs que el builder admita.
    kw: dict = {}
    for cand in _CKPT_PARAMS:
        if cand in params:
            kw[cand] = checkpoint
            break
    if "seed" in params:
        kw["seed"] = seed
    if lora and "lora" in params:
        kw["lora"] = lora
        if "lora_strength" in params:
            kw["lora_strength"] = lora_strength
    if size is not None and "size" in params:
        kw["size"] = size

    # Posicionales requeridos extra (más allá del primero = subject). Ej: emote.expression.
    pos_extra = [
        p.name
        for p in params.values()
        if p.default is inspect.Parameter.empty
        and p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
    ][1:]
    pos_args = [subject_eff]
    for name in pos_extra:
        if name not in item:
            raise ValueError(f"kind {kind!r} requiere el campo {name!r} en el item del pack")
        pos_args.append(item[name])

    # Passthrough por-asset: cualquier otra clave del item que el builder admita
    # (size, width, height, tier, view, direction, glow, transparent, negative, ...).
    # Sobreescribe la coherencia común si el usuario lo pide explícitamente.
    used = _RESERVED_ITEM_KEYS | set(pos_extra)
    for key, val in item.items():
        if key in used:
            continue
        if key in params:
            kw[key] = val

    workflow = fn(*pos_args, **kw)
    return workflow, subject_eff, godot_kind, seed


def _first_image(outputs: dict) -> dict | None:
    """Primer descriptor de imagen {filename, subfolder, type} en los outputs."""
    for node_out in outputs.values():
        images = node_out.get("images") if isinstance(node_out, dict) else None
        if images:
            return images[0]
    return None


def comfyui_generate_asset_pack_oneshot(
    pack: list,
    *,
    checkpoint: str = "IMG_dreamshaper_8.safetensors",
    style: str = "",
    lora: str | None = None,
    lora_strength: float = 1.0,
    base_seed: int = 0,
    size: int | None = None,
    server: str = "127.0.0.1:8188",
    export_godot: str | None = None,
    out_dir: str | None = None,
    wait_timeout: float = 600.0,
    godot_bin: str | None = None,
) -> dict:
    """Genera un set COHERENTE de assets 2D para un mismo juego, de un solo tiro.

    Args:
        pack: lista de assets. Cada item es un dict con al menos `kind` (uno de
            `supported_kinds()`) y `subject` (qué dibujar). Campos extra se reenvían
            como kwargs al builder si los admite (p.ej. {"kind":"achievement_badge",
            "subject":"first kill","tier":"gold"} o {"kind":"emote","subject":"hero",
            "expression":"angry"}).
        checkpoint: modelo base COMPARTIDO por todos los assets. keyword-only.
        style: estilo común concatenado al `subject` de cada asset (la firma visual
            del pack, p.ej. "dark fantasy, hand-painted"). keyword-only.
        lora: LoRA de estilo COMPARTIDO; se aplica solo a builders con un param
            `lora` genérico (no pisa LoRAs funcionales como el de isometric/pixelart).
            keyword-only.
        lora_strength: fuerza del LoRA común. keyword-only.
        base_seed: semilla base; el asset i usa `seed = base_seed + i` (reproducible
            y variado). keyword-only.
        size: resolución cuadrada común para los builders que aceptan `size`
            (icons/sprites); None deja el default de cada builder. Útil para bajar
            VRAM (p.ej. 512 con SD1.5). keyword-only.
        server: host:port del servidor ComfyUI; se acepta con o sin esquema
            (`http://`). keyword-only.
        export_godot: ruta raíz de un proyecto Godot 4; si se da, cada asset se
            exporta a `res://assets/...` con su `.import`. None = no exportar.
            keyword-only.
        out_dir: directorio local donde descargar los PNG; None = un dir temporal
            por pack. keyword-only.
        wait_timeout: segundos máximos esperando cada asset en ComfyUI. keyword-only.
        godot_bin: binario de Godot para el reimport headless; None autodetecta.
            keyword-only.

    Returns:
        dict {ok, pack_dir, checkpoint, style, lora, base_seed, coherence_note,
        count, generated, failed, assets, error}. `assets` es una lista, una entrada
        por item del pack, con {index, kind, subject, seed, prompt_id, path, exported,
        ok, error}. `ok` global es True solo si TODOS los assets se generaron. Un
        fallo aislado (p.ej. OOM en uno) NO aborta el resto: queda registrado en su
        entrada con `ok=False` y `error`.
    """
    server = server.replace("http://", "").replace("https://", "").strip("/")
    out = {
        "ok": False, "pack_dir": "", "checkpoint": checkpoint, "style": style,
        "lora": lora, "base_seed": base_seed, "coherence_note": "",
        "count": 0, "generated": 0, "failed": 0, "assets": [], "error": "",
    }

    if not pack or not isinstance(pack, list):
        out["error"] = "pack vacío o no es una lista"
        return out
    out["count"] = len(pack)

    # Validación previa: todos los kinds deben estar soportados ANTES de tocar la GPU.
    bad = [str(it.get("kind", "")) for it in pack if str(it.get("kind", "")).strip() not in _SUPPORTED]
    if bad:
        out["error"] = (
            f"kind(s) no soportado(s): {', '.join(sorted(set(bad)))}. "
            f"Soportados: {', '.join(supported_kinds())}"
        )
        return out

    pack_dir = out_dir or os.path.join(tempfile.gettempdir(), f"asset_pack_seed{base_seed}")
    os.makedirs(pack_dir, exist_ok=True)
    out["pack_dir"] = pack_dir
    out["coherence_note"] = (
        f"Set coherente: checkpoint={checkpoint!r}, style={style!r}, lora={lora!r} "
        f"compartidos por los {len(pack)} assets; seed = base_seed({base_seed}) + indice."
    )

    # Índice del último asset a exportar -> único reimport headless de Godot al final.
    last_export_idx = -1
    if export_godot:
        last_export_idx = max(range(len(pack)), default=-1)

    for i, item in enumerate(pack):
        entry = {
            "index": i, "kind": str(item.get("kind", "")), "subject": item.get("subject", ""),
            "seed": base_seed + i, "prompt_id": "", "path": "", "exported": None,
            "ok": False, "error": "",
        }
        try:
            workflow, subject_eff, godot_kind, seed = _build_item_workflow(
                item, checkpoint=checkpoint, style=style, lora=lora,
                lora_strength=lora_strength, base_seed=base_seed, size=size, index=i,
            )
            entry["subject"] = subject_eff
            entry["seed"] = seed
        except (ValueError, TypeError) as exc:
            entry["error"] = f"build fallo: {exc}"
            out["assets"].append(entry)
            out["failed"] += 1
            continue

        # Encolar.
        try:
            sub = comfyui_submit_workflow(workflow, server=server)
            entry["prompt_id"] = sub["prompt_id"]
        except (RuntimeError, KeyError) as exc:
            entry["error"] = f"submit fallo: {exc}"
            out["assets"].append(entry)
            out["failed"] += 1
            continue

        # Esperar + localizar el PNG.
        try:
            outputs = comfyui_wait_result(entry["prompt_id"], server=server, timeout=wait_timeout)
        except (TimeoutError, RuntimeError) as exc:
            entry["error"] = f"wait fallo: {exc}"
            out["assets"].append(entry)
            out["failed"] += 1
            continue
        img = _first_image(outputs)
        if img is None:
            entry["error"] = f"el workflow no produjo imágenes (outputs={list(outputs)})"
            out["assets"].append(entry)
            out["failed"] += 1
            continue

        # Descargar.
        fetched = comfyui_fetch_output_image(
            img["filename"], subfolder=img.get("subfolder", ""),
            type_=img.get("type", "output"), server=server, dest_dir=pack_dir,
        )
        if not fetched.get("ok"):
            entry["error"] = f"fetch fallo: {fetched.get('error')}"
            out["assets"].append(entry)
            out["failed"] += 1
            continue
        entry["path"] = fetched["path"]
        entry["ok"] = True
        out["generated"] += 1

        # Export opcional a Godot (reimport solo en el último -> un único scan).
        if export_godot:
            exp = comfyui_export_asset_to_godot(
                entry["path"], godot_kind, export_godot,
                reimport=(i == last_export_idx), godot_bin=godot_bin,
            )
            entry["exported"] = exp
            if not exp.get("ok"):
                entry["error"] = f"export fallo (asset generado igual): {exp.get('error')}"

        out["assets"].append(entry)

    out["ok"] = out["failed"] == 0 and out["generated"] == len(pack)
    if out["failed"]:
        out["error"] = f"{out['failed']}/{len(pack)} assets fallaron (ver assets[].error)"
    return out


# Alias con el nombre completo del ID para descubrimiento por convención.
generate_asset_pack_oneshot = comfyui_generate_asset_pack_oneshot


if __name__ == "__main__":
    import json

    res = comfyui_generate_asset_pack_oneshot(
        [
            {"kind": "item_icon", "subject": "magic sword"},
            {"kind": "enemy_creature", "subject": "goblin warrior"},
        ],
        checkpoint="IMG_dreamshaper_8.safetensors",
        style="dark fantasy, hand-painted",
        base_seed=42,
        size=512,
        out_dir="/tmp/asset_pack_demo",
    )
    print(json.dumps(res, indent=2, ensure_ascii=False))
