"""FikoYT Demucs Split — vocal/drums/bass/other 4 stem ayrımı."""
import os
import sys
import tempfile
import subprocess

import torch


DEMUCS_MODELS = ["htdemucs", "htdemucs_ft", "mdx_extra", "mdx_extra_q"]


def _audio_to_wav(audio: dict, dst: str):
    import torchaudio
    wav = audio["waveform"]
    sr = int(audio["sample_rate"])
    if wav.dim() == 3:
        wav = wav.squeeze(0)
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    torchaudio.save(dst, wav.cpu().float(), sr)


def _load_wav(path: str) -> dict:
    import torchaudio
    wav, sr = torchaudio.load(path)
    if wav.dim() == 2:
        wav = wav.unsqueeze(0)
    return {"waveform": wav.float(), "sample_rate": int(sr)}


def _empty_audio(sr: int = 44100) -> dict:
    return {"waveform": torch.zeros(1, 2, sr), "sample_rate": sr}


class FikoYTDemucsSplit:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO",),
                "model": (DEMUCS_MODELS, {"default": "htdemucs"}),
            }
        }

    RETURN_TYPES = ("AUDIO", "AUDIO", "AUDIO", "AUDIO")
    RETURN_NAMES = ("vocals", "drums", "bass", "other")
    FUNCTION = "split"
    CATEGORY = "FikoYT/audio"

    def split(self, audio, model):
        try:
            import importlib.util
            if importlib.util.find_spec("demucs") is None:
                raise ImportError
        except ImportError:
            raise RuntimeError(
                f"Demucs kurulu değil.\nKurulum: \"{sys.executable}\" -m pip install demucs"
            )

        tmpdir = tempfile.mkdtemp(prefix="fikoyt_demucs_")
        in_path = os.path.join(tmpdir, "input.wav")
        _audio_to_wav(audio, in_path)

        cmd = [sys.executable, "-m", "demucs", "-n", model, "-o", tmpdir, in_path]
        print(f"[FikoYT Demucs] {model} çalışıyor…")
        env = os.environ.copy()
        env.pop("PYTHONHASHSEED", None)
        try:
            proc = subprocess.run(cmd, capture_output=True, timeout=1800, env=env)
            if proc.returncode != 0:
                err = (proc.stderr or b"").decode("utf-8", errors="replace")[-600:]
                raise RuntimeError(f"demucs hata: {err}")
        except subprocess.TimeoutExpired:
            raise RuntimeError("demucs 30 dk içinde tamamlanamadı")

        stem_dir = os.path.join(tmpdir, model, "input")
        if not os.path.isdir(stem_dir):
            raise RuntimeError(f"demucs çıktı klasörü yok: {stem_dir}")

        outputs = {}
        for fname in os.listdir(stem_dir):
            if not fname.lower().endswith(".wav"):
                continue
            stem_id = os.path.splitext(fname)[0].lower()
            outputs[stem_id] = _load_wav(os.path.join(stem_dir, fname))

        sr_fallback = next(iter(outputs.values()))["sample_rate"] if outputs else 44100
        vocals = outputs.get("vocals", _empty_audio(sr_fallback))
        drums = outputs.get("drums", _empty_audio(sr_fallback))
        bass = outputs.get("bass", _empty_audio(sr_fallback))
        other = outputs.get("other", _empty_audio(sr_fallback))

        print(f"[FikoYT Demucs] tamam — {list(outputs.keys())}")
        return (vocals, drums, bass, other)


NODE_CLASS_MAPPINGS = {"FikoYTDemucsSplit": FikoYTDemucsSplit}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTDemucsSplit": "FikoYT • Demucs Stem Split"}