"""FikoYT SongGen v2 — Tencent SongGeneration (LeVo) sampler.

smthemex/ComfyUI_SongGeneration node'unun düşük seviyeli API'sini çağırır:
generate.inference_lowram_final / song_infer_lowram. API bulunamazsa try/except ile
anlamlı hata fırlatır.
"""
import os
import time
import torch
import torchaudio

try:
    import folder_paths
except ImportError:
    folder_paths = None

from ..core import normalize_seed, get_models_dir


class FikoYTSongGenV2:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "songgen_model": ("MODEL",),
                "lyrics": ("STRING", {"multiline": True,
                                       "default": "[verse]\nHelal sana güzel insanlık\n[chorus]\nUmut kalır en sonunda"}),
                "style_prompt": ("STRING", {"default": "pop, female vocal, uplifting, 120 bpm"}),
                "duration_sec": ("FLOAT", {"default": 60.0, "min": 5.0, "max": 180.0, "step": 1.0}),
                "temperature": ("FLOAT", {"default": 0.95, "min": 0.5, "max": 1.5, "step": 0.01}),
                "top_k": ("INT", {"default": 250, "min": 10, "max": 500, "step": 1}),
                "cfg": ("FLOAT", {"default": 7.5, "min": 1.0, "max": 15.0, "step": 0.1}),
                "seed": ("INT", {"default": -1, "min": -1, "max": 2**31 - 1}),
            },
        }

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("audio",)
    FUNCTION = "generate"
    CATEGORY = "FikoYT/audio"

    def generate(self, songgen_model, lyrics, style_prompt, duration_sec,
                 temperature, top_k, cfg, seed):
        if isinstance(songgen_model, dict) and songgen_model.get("error"):
            raise RuntimeError(
                f"SongGen loader hatası: {songgen_model['error']}\n"
                "ComfyUI_SongGeneration node'unun kurulu olduğundan emin ol."
            )
        if not isinstance(songgen_model, dict) or not songgen_model.get("ok"):
            raise RuntimeError("Geçersiz SongGen handle — Loader node'unu yeniden çalıştır.")

        seed = normalize_seed(seed) if seed >= 0 else int(time.time() * 1000) % (2**31)

        try:
            from custom_nodes.ComfyUI_SongGeneration.generate import (  # type: ignore
                inference_lowram_final, song_infer_lowram,
            )
        except Exception:
            import importlib
            mod = importlib.import_module("ComfyUI_SongGeneration.generate")
            inference_lowram_final = getattr(mod, "inference_lowram_final", None)
            song_infer_lowram = getattr(mod, "song_infer_lowram", None)

        model = songgen_model["model"]
        cfg_obj = songgen_model["cfg"]

        infer_fn = song_infer_lowram if songgen_model.get("use_gguf") else inference_lowram_final
        if infer_fn is None:
            raise RuntimeError("ComfyUI_SongGeneration inference API'si bulunamadı.")

        # Tencent SongGeneration sampler kwargs — sürüm farklılıklarına karşı esnek dispatch
        kw = dict(
            lyric=lyrics,
            tags=style_prompt,
            duration=float(duration_sec),
            seed=seed,
            temperature=float(temperature),
            top_k=int(top_k),
            cfg_scale=float(cfg),
        )

        try:
            result = infer_fn(model, cfg_obj, **kw)
        except TypeError:
            # Bazı sürümlerde positional order farklı — sade fallback
            result = infer_fn(model, cfg_obj, lyrics, style_prompt, float(duration_sec))

        # Result: (waveform, sample_rate) ya da {"waveform":..., "sample_rate":...}
        if isinstance(result, dict) and "waveform" in result:
            wav = result["waveform"]
            sr = int(result.get("sample_rate", 48000))
        elif isinstance(result, (tuple, list)) and len(result) >= 2:
            wav, sr = result[0], int(result[1])
        else:
            raise RuntimeError(f"Beklenmedik SongGen çıktı tipi: {type(result)}")

        if not isinstance(wav, torch.Tensor):
            wav = torch.as_tensor(wav, dtype=torch.float32)
        if wav.dim() == 1:
            wav = wav.unsqueeze(0)  # [samples] → [1, samples]
        if wav.dim() == 2:
            wav = wav.unsqueeze(0)  # [channels, samples] → [1, channels, samples]

        print(f"[FikoYT SongGen] {duration_sec:.0f}s @ {sr}Hz tamamlandı (seed={seed})")
        return ({"waveform": wav, "sample_rate": sr},)


NODE_CLASS_MAPPINGS = {"FikoYTSongGenV2": FikoYTSongGenV2}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTSongGenV2": "FikoYT • SongGen v2 (LeVo)"}