"""FikoYT SongGen v2 (LeVo) Loader — PT (12.9G) veya GGUF Q8 (7.6G) auto-download."""
import os

try:
    import folder_paths
except ImportError:
    folder_paths = None

from ..core import ensure_model, get_models_dir


SONGGEN_REPO = "tencent/SongGeneration"

SONGGEN_MODES = {
    "pt_full": {
        "label": "PT Full (~12.9 GB, en yüksek kalite)",
        "subdir": "SongGeneration",
        "filename": "songgeneration_base/model.pt",
        "use_gguf": False,
    },
    "gguf_q8": {
        "label": "GGUF Q8 (~7.6 GB, düşük VRAM)",
        "subdir": "gguf",
        "filename": "songgeneration_base-Q8_0.gguf",
        "use_gguf": True,
    },
}


def _try_loader_call(infer_model_path, mode_info):
    """smthemex/ComfyUI_SongGeneration build_model API'sini çağır.
    Başarısız olursa lazy-handle dön — generate node kendi try/except'iyle ele alır.
    """
    try:
        from custom_nodes.ComfyUI_SongGeneration.generate import build_model  # type: ignore
    except Exception:
        try:
            import importlib
            mod = importlib.import_module("ComfyUI_SongGeneration.generate")
            build_model = getattr(mod, "build_model")
        except Exception as e:
            return {"error": f"ComfyUI_SongGeneration node yüklenemedi: {e}",
                    "infer_model_path": infer_model_path, "use_gguf": mode_info["use_gguf"]}

    try:
        ckpt_root = os.path.join(get_models_dir("SongGeneration"), "ckpt")
        os.makedirs(ckpt_root, exist_ok=True)
        model, cfg = build_model(ckpt_root, infer_model_path, "v2", True, True)
        return {"ok": True, "model": model, "cfg": cfg,
                "infer_model_path": infer_model_path, "use_gguf": mode_info["use_gguf"]}
    except Exception as e:
        return {"error": f"build_model hata: {e}",
                "infer_model_path": infer_model_path, "use_gguf": mode_info["use_gguf"]}


class FikoYTSongGenLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mode": (list(SONGGEN_MODES.keys()), {"default": "gguf_q8"}),
                "auto_download": ("BOOLEAN", {"default": True}),
                "use_flash_attn": ("BOOLEAN", {"default": True}),
                "offload_audiolm": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("SONGGEN_MODEL",)
    FUNCTION = "load"
    CATEGORY = "FikoYT/audio"

    def load(self, mode, auto_download, use_flash_attn, offload_audiolm):
        info = SONGGEN_MODES[mode]
        target_dir = get_models_dir(info["subdir"])
        target_path = os.path.join(target_dir, os.path.basename(info["filename"]))

        if not os.path.isfile(target_path):
            if auto_download:
                try:
                    target_path = ensure_model(SONGGEN_REPO, info["filename"], info["subdir"])
                except Exception as e:
                    raise FileNotFoundError(
                        f"SongGen indirilemedi: {e}\nManuel: huggingface.co/{SONGGEN_REPO}"
                    )
            else:
                raise FileNotFoundError(
                    f"SongGen modeli yok: {target_path} (auto_download=False)"
                )

        handle = _try_loader_call(target_path, info)
        handle["mode"] = mode
        handle["use_flash_attn"] = use_flash_attn
        handle["offload_audiolm"] = offload_audiolm
        return (handle,)


NODE_CLASS_MAPPINGS = {"FikoYTSongGenLoader": FikoYTSongGenLoader}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTSongGenLoader": "FikoYT • SongGen Loader (v2 / LeVo)"}