"""Lista los modelos instalados de ComfyUI por carpeta, resolviendo rutas reales.

ComfyUI puede leer los modelos desde rutas externas declaradas en
`<comfyui_dir>/extra_model_paths.yaml` (en este equipo, /mnt/2tb/comfyui_models/),
ademas de las nativas `<comfyui_dir>/models/<folder>/`. Esta funcion escanea
AMBAS para cada carpeta de tipo (checkpoints, loras, vae, controlnet,
upscale_models), de modo que ve los modelos aunque no esten bajo `models/`.

El escaneo es del sistema de archivos (no depende del servidor ComfyUI), asi que
funciona aunque el servidor este reiniciandose.

Impura: lectura de disco (FS scan + parse de YAML). Solo stdlib + PyYAML.
"""
import os

_DEFAULT_FOLDERS = ["checkpoints", "loras", "vae", "controlnet", "upscale_models"]
_MODEL_EXTS = (
    ".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft", ".onnx",
)


def _resolve_external_roots(base: str) -> dict:
    """Lee extra_model_paths.yaml y devuelve {folder: [dir_externo, ...]}.

    Maneja valores de carpeta multilinea (varias subrutas por clave). Si el YAML
    no existe o no se puede parsear, devuelve {} (solo se usaran las rutas
    nativas).
    """
    roots: dict = {}
    yml = os.path.join(base, "extra_model_paths.yaml")
    if not os.path.isfile(yml):
        return roots
    try:
        import yaml
        with open(yml, encoding="utf-8") as fh:
            data = yaml.safe_load(fh) or {}
    except Exception:  # noqa: BLE001 — YAML ilegible: degradar a rutas nativas
        return roots

    if not isinstance(data, dict):
        return roots
    for section in data.values():
        if not isinstance(section, dict):
            continue
        bp = os.path.expanduser(str(section.get("base_path", "")))
        for key in _DEFAULT_FOLDERS:
            sub = section.get(key)
            if not sub:
                continue
            for line in str(sub).splitlines():
                line = line.strip()
                if not line:
                    continue
                roots.setdefault(key, []).append(os.path.join(bp, line))
    return roots


def comfyui_list_installed_models(
    folder: str | None = None,
    comfyui_dir: str = "~/ComfyUI",
) -> dict:
    """Lista los modelos en disco por carpeta de tipo.

    Args:
        folder: carpeta concreta a listar (ej. "checkpoints"). Si None, lista
            todas las de _DEFAULT_FOLDERS.
        comfyui_dir: raiz de la instalacion de ComfyUI (se expande ~).

    Returns:
        dict {ok, models, error}. models es {folder: [nombre, ...]} con los
        archivos de modelo (dedup por nombre) hallados tanto en la ruta nativa
        `models/<folder>/` como en las externas de extra_model_paths.yaml. ok es
        True salvo fallo inesperado.
    """
    base = os.path.expanduser(comfyui_dir)
    folders = [folder] if folder else list(_DEFAULT_FOLDERS)
    external = _resolve_external_roots(base)

    models: dict = {}
    try:
        for f in folders:
            dirs = [os.path.join(base, "models", f)] + external.get(f, [])
            names: list = []
            seen: set = set()
            for d in dirs:
                if not os.path.isdir(d):
                    continue
                for entry in sorted(os.listdir(d)):
                    if entry in seen:
                        continue
                    p = os.path.join(d, entry)
                    if os.path.isfile(p) and entry.lower().endswith(_MODEL_EXTS):
                        seen.add(entry)
                        names.append(entry)
            models[f] = names
    except OSError as exc:
        return {"ok": False, "models": models, "error": f"fallo escaneando: {exc}"}

    return {"ok": True, "models": models, "error": ""}


if __name__ == "__main__":
    import json

    print(json.dumps(comfyui_list_installed_models(), ensure_ascii=False, indent=2))
