"""comfyui_build_skill_workflow — compila una receta de *skill* a un workflow ComfyUI.

Una **skill** es una receta versionada (el dict de `recipe.json` del grupo `comfyui-skill`) que
fija checkpoint, LoRAs, parametros de sampling, scaffold de prompt y bloques de post-proceso.
Esta funcion PURA la compila a un dict de workflow en "API format" listo para
`comfyui_submit_workflow`, componiendo los builders del registry segun `recipe['base_workflow']`.

Despacho de `base_workflow` (los que se construyen solo a partir de un `subject` de texto):

    txt2img       -> comfyui_build_txt2img_workflow
    flux          -> comfyui_build_flux_workflow
    sdxl_refiner  -> comfyui_build_sdxl_refiner_workflow

Bloques de post-proceso (`recipe['blocks']`, aplicados en orden sobre el dict resultante):

    facedetailer  -> comfyui_build_facedetailer_workflow   (encadena sobre el dict)
    hires_fix     -> comfyui_inject_hires_fix              (encadena sobre el dict)

Pasos: scaffold de prompt ({subject} + trigger_words) -> builder base -> LoRAs (inject_lora) ->
bloques (en orden). Los `base_workflow` que necesitan una imagen de entrada (img2img, inpaint,
controlnet) NO se soportan aqui — `build_skill_workflow` arranca de texto.

Funcion pura: sin red, sin I/O. Solo compone builders puros del registry.
"""

from __future__ import annotations

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))


class SkillWorkflowError(ValueError):
    """Error tipado al compilar una receta de skill (base_workflow desconocido,
    recipe invalida, bloque no soportado, dependencia ausente)."""


_IMAGE_INPUT_BASES = {"img2img", "inpaint", "controlnet"}


def _pick(params: dict, keys) -> dict:
    """Subconjunto de `params` con solo las claves presentes en `keys`."""
    return {k: params[k] for k in keys if k in params}


def _scaffold_prompts(recipe: dict, subject: str) -> tuple[str, str]:
    """Sustituye `{subject}` y antepone los trigger_words en el scaffold de prompt.

    Devuelve (positive, negative). Si no hay `prompt_scaffold`, usa el subject como
    positivo y "" como negativo.
    """
    scaffold = recipe.get("prompt_scaffold") or {}
    positive = str(scaffold.get("positive", "") or "")
    negative = str(scaffold.get("negative", "") or "")
    if "{subject}" in positive:
        positive = positive.replace("{subject}", subject)
    elif not positive:
        positive = subject
    else:
        # scaffold sin placeholder: el subject se antepone para no perderlo.
        positive = f"{subject}, {positive}"
    negative = negative.replace("{subject}", subject)

    triggers = scaffold.get("trigger_words") or []
    if triggers:
        positive = ", ".join(list(triggers) + [positive]) if positive else ", ".join(triggers)
    return positive, negative


def _build_base(recipe: dict, positive: str, negative: str, seed: int) -> dict:
    """Despacha al builder base del registry segun `recipe['base_workflow']`."""
    base = recipe.get("base_workflow")
    params = dict(recipe.get("params") or {})
    ckpt = recipe.get("checkpoint", "")

    if base == "txt2img":
        from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
        if not ckpt:
            raise SkillWorkflowError("base_workflow txt2img requiere recipe['checkpoint']")
        kw = _pick(params, ("steps", "cfg", "width", "height", "sampler_name", "scheduler"))
        return comfyui_build_txt2img_workflow(ckpt, positive, negative, seed=seed, **kw)

    if base == "flux":
        from ml.comfyui_build_flux_workflow import comfyui_build_flux_workflow
        kw = _pick(params, ("clip_l", "t5xxl", "vae", "width", "height", "steps",
                            "guidance", "weight_dtype", "sampler_name", "scheduler"))
        # En Flux el "checkpoint" de la receta es el modelo de difusion (unet).
        if ckpt:
            kw.setdefault("unet", ckpt)
        return comfyui_build_flux_workflow(positive, seed=seed, **kw)

    if base == "sdxl_refiner":
        from ml.comfyui_build_sdxl_refiner_workflow import comfyui_build_sdxl_refiner_workflow
        refiner = params.get("refiner_ckpt")
        if not ckpt or not refiner:
            raise SkillWorkflowError(
                "base_workflow sdxl_refiner requiere recipe['checkpoint'] (base) y "
                "recipe['params']['refiner_ckpt']")
        kw = _pick(params, ("base_steps", "refiner_steps", "cfg", "width", "height"))
        return comfyui_build_sdxl_refiner_workflow(ckpt, refiner, positive, negative, seed=seed, **kw)

    if base in _IMAGE_INPUT_BASES:
        raise SkillWorkflowError(
            f"base_workflow {base!r} requiere una imagen de entrada; build_skill_workflow "
            "compila a partir de un subject de texto. Construye ese workflow aparte.")

    raise SkillWorkflowError(
        f"base_workflow desconocido: {base!r}. Soportados: txt2img, flux, sdxl_refiner.")


def _apply_loras(workflow: dict, loras) -> dict:
    """Encadena los LoRAs de la receta via comfyui_inject_lora (en orden)."""
    if not loras:
        return workflow
    from ml.comfyui_inject_lora import comfyui_inject_lora
    wf = workflow
    for lora in loras:
        name = lora.get("name")
        if not name:
            raise SkillWorkflowError("cada lora de la receta necesita 'name'")
        wf = comfyui_inject_lora(
            wf,
            name,
            strength_model=lora.get("strength_model", 1.0),
            strength_clip=lora.get("strength_clip", 1.0),
        )
    return wf


def _apply_block(workflow: dict, block: dict, recipe: dict, positive: str,
                 negative: str, seed: int) -> dict:
    """Aplica un bloque de post-proceso sobre el workflow (facedetailer | hires_fix)."""
    btype = block.get("type")
    bparams = dict(block.get("params") or {})

    if btype == "facedetailer":
        from ml.comfyui_build_facedetailer_workflow import comfyui_build_facedetailer_workflow
        bparams.setdefault("seed", seed)
        return comfyui_build_facedetailer_workflow(
            workflow, recipe.get("checkpoint", ""), positive, negative, **bparams)

    if btype == "hires_fix":
        try:
            from ml.comfyui_inject_hires_fix import comfyui_inject_hires_fix
        except ImportError as exc:
            raise SkillWorkflowError(
                "bloque hires_fix requiere comfyui_inject_hires_fix_py_ml (no disponible)") from exc
        bparams.setdefault("seed", seed)
        return comfyui_inject_hires_fix(workflow, **bparams)

    raise SkillWorkflowError(
        f"bloque de tipo desconocido: {btype!r}. Soportados: facedetailer, hires_fix.")


def build_skill_workflow(recipe: dict, subject: str, *, seed: int = 0) -> dict:
    """Compila una receta de skill a un workflow ComfyUI en API format.

    Args:
        recipe: dict de la receta (schema `comfyui-skill`). Campos usados:
            `base_workflow` (txt2img|flux|sdxl_refiner), `checkpoint`, `loras`,
            `params`, `prompt_scaffold` (`positive`/`negative` con `{subject}` +
            `trigger_words`), `blocks` (lista de `{type, params}`).
        subject: el sujeto concreto que sustituye `{subject}` en el scaffold del
            prompt (p.ej. "a woman with red hair").
        seed: semilla de generacion; se pasa al builder base y por defecto a cada
            bloque que la acepte. keyword-only.

    Returns:
        dict en API format listo para `comfyui_submit_workflow`.

    Raises:
        SkillWorkflowError: si `base_workflow` es desconocido o necesita imagen, si
            la receta no es un dict valido, si falta un checkpoint requerido, o si un
            bloque es de tipo no soportado / su dependencia no esta disponible.
    """
    if not isinstance(recipe, dict):
        raise SkillWorkflowError(f"recipe debe ser dict, no {type(recipe).__name__}")

    positive, negative = _scaffold_prompts(recipe, subject)
    workflow = _build_base(recipe, positive, negative, seed)
    workflow = _apply_loras(workflow, recipe.get("loras"))

    for block in (recipe.get("blocks") or []):
        if not isinstance(block, dict):
            raise SkillWorkflowError(f"cada block debe ser dict, no {type(block).__name__}")
        workflow = _apply_block(workflow, block, recipe, positive, negative, seed)

    return workflow


# Alias con el nombre completo del ID para descubrimiento por convencion.
comfyui_build_skill_workflow = build_skill_workflow


if __name__ == "__main__":
    import json

    demo = {
        "schema_version": 1,
        "slug": "portrait_demo",
        "version": "1.0.0",
        "base_workflow": "txt2img",
        "checkpoint": "IMG_juggernaut_xl_v11.safetensors",
        "loras": [{"name": "add_detail.safetensors", "strength_model": 0.6, "strength_clip": 0.6}],
        "params": {"steps": 30, "cfg": 5.5, "sampler_name": "dpmpp_2m",
                   "scheduler": "karras", "width": 832, "height": 1216},
        "prompt_scaffold": {"positive": "cinematic portrait of {subject}, sharp focus",
                            "negative": "blurry, lowres", "trigger_words": []},
        "blocks": [{"type": "facedetailer", "params": {"denoise": 0.45}}],
    }
    wf = build_skill_workflow(demo, "a woman with red hair", seed=42)
    print(json.dumps(wf, indent=2))
