#!/usr/bin/env python3
"""SOTAKA Vocal Remover — Demucs on Hugging Face ZeroGPU."""

from __future__ import annotations

import os
import time
from pathlib import Path
from typing import Optional, Tuple

os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")

import numpy as np
import soundfile as sf
import torch
import torchaudio

try:
    import spaces
except ImportError:

    class spaces:  # type: ignore
        @staticmethod
        def GPU(*_a, **_k):
            def deco(fn):
                return fn

            return deco


ROOT = Path(__file__).resolve().parent
OUT_DIR = ROOT / "outputs"
OUT_DIR.mkdir(exist_ok=True)

# 4 min song + Demucs load does not fit in a single 120s ZeroGPU slot.
MAX_SOURCE_SEC = int(os.environ.get("DEMUCS_MAX_SOURCE_SEC", "300"))
CHUNK_SEC = int(os.environ.get("DEMUCS_CHUNK_SEC", "70"))
CHUNK_OVERLAP_SEC = float(os.environ.get("DEMUCS_CHUNK_OVERLAP", "1.0"))

_model = None
_model_name = None


def _device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def _get_model(name: str = "htdemucs"):
    global _model, _model_name
    name = (name or "htdemucs").strip() or "htdemucs"
    if _model is not None and _model_name == name:
        return _model
    from demucs.pretrained import get_model

    m = get_model(name)
    m.to(_device())
    m.eval()
    _model = m
    _model_name = name
    return m


def _gpu_err(exc: BaseException) -> str:
    msg = str(exc)
    low = msg.lower()
    if (
        "zerogpu" in low
        or "expired" in low
        or "proxy token" in low
        or "gpu task" in low
        or "cuda" in low
        and "out of memory" in low
    ):
        return (
            "GPU timed out or ran out of memory on this clip. "
            "Keep the tab open and click Separate again "
            f"(songs are split into ~{CHUNK_SEC}s chunks). Detail: {exc}"
        )
    return f"Separation error: {exc}"


def _crossfade(chunk1: np.ndarray, chunk2: np.ndarray, overlap: int) -> np.ndarray:
    fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
    fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
    chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
    return chunk2


def _stitch(wavs: list, sr: int, overlap_sec: float) -> np.ndarray:
    if len(wavs) == 1:
        return wavs[0]
    overlap = int(sr * overlap_sec)
    out = wavs[0]
    for nxt in wavs[1:]:
        if overlap > 0 and len(out) > overlap and len(nxt) > overlap:
            merged = _crossfade(out[-overlap:].copy(), nxt[:overlap].copy(), overlap)
            out = np.concatenate([out[:-overlap], merged, nxt[overlap:]])
        else:
            out = np.concatenate([out, nxt])
    return out.astype(np.float32)


@spaces.GPU(duration=120)
def _separate_chunk(
    audio_path: str,
    model_name: str = "htdemucs",
) -> Tuple[Optional[str], Optional[str], str]:
    """One ZeroGPU session: Demucs on a short clip, write wavs, return paths."""
    global _model, _model_name
    _model = None
    _model_name = None

    try:
        model = _get_model(model_name)
    except Exception as e:
        return None, None, f"Failed to load Demucs: {e}"

    device = _device()
    try:
        wav, sr = torchaudio.load(audio_path)
        if wav.dim() == 1:
            wav = wav.unsqueeze(0)
        if wav.size(0) == 1:
            wav = wav.repeat(2, 1)
        elif wav.size(0) > 2:
            wav = wav[:2]

        from demucs.apply import apply_model
        from demucs.audio import convert_audio

        wav = convert_audio(wav, sr, model.samplerate, model.audio_channels)
        wav = wav.unsqueeze(0).to(device)
        with torch.no_grad():
            sources = apply_model(
                model,
                wav,
                device=device,
                split=True,
                overlap=0.25,
                shifts=0,
            )[0]
        names = list(model.sources)
        if "vocals" not in names:
            return None, None, f"Model has no vocals stem: {names}"

        v_idx = names.index("vocals")
        vocals = sources[v_idx].mean(0).cpu().numpy().astype(np.float32)
        other = None
        for i, n in enumerate(names):
            if i == v_idx:
                continue
            stem = sources[i].mean(0).cpu().numpy().astype(np.float32)
            other = stem if other is None else other + stem

        stamp = f"{os.getpid()}_{int(time.time() * 1000)}"
        v_path = OUT_DIR / f"vr_v_{stamp}.wav"
        i_path = OUT_DIR / f"vr_i_{stamp}.wav"
        out_sr = int(model.samplerate)
        sf.write(v_path, vocals, out_sr)
        if other is not None:
            sf.write(i_path, other, out_sr)
            inst = str(i_path)
        else:
            inst = None
        dur = len(vocals) / out_sr
        log = f"ok device={device} dur={dur:.1f}s sr={out_sr}"
        return str(v_path), inst, log
    except Exception as e:
        return None, None, f"Separation error: {e}"


def separate_vocals(
    audio_path: Optional[str],
    model_name: str = "htdemucs",
) -> Tuple[Optional[str], Optional[str], str]:
    """Return (vocals_wav, instrumental_wav, log). Long mixes are chunked for ZeroGPU."""
    if not audio_path:
        return None, None, "Upload a song / mix."

    model_name = (model_name or "htdemucs").strip() or "htdemucs"

    try:
        wav, sr = torchaudio.load(audio_path)
    except Exception as e:
        return None, None, f"Could not read audio: {e}"

    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    if wav.size(0) > 2:
        wav = wav[:2]

    dur = wav.size(-1) / float(sr)
    notes = []
    if dur > MAX_SOURCE_SEC:
        wav = wav[..., : int(sr * MAX_SOURCE_SEC)]
        notes.append(f"Trimmed {dur:.0f}s → {MAX_SOURCE_SEC}s.")
        dur = MAX_SOURCE_SEC

    chunk_samples = int(sr * CHUNK_SEC)
    hop = max(1, int(sr * (CHUNK_SEC - CHUNK_OVERLAP_SEC)))
    starts = [0]
    if wav.size(-1) > chunk_samples:
        starts = list(range(0, wav.size(-1), hop))
        if len(starts) > 1 and wav.size(-1) - starts[-1] < sr * 8:
            starts = starts[:-1]

    notes.append(
        f"Song {dur:.0f}s → {len(starts)} Demucs chunk(s) (~{CHUNK_SEC}s). Keep this tab open."
    )

    v_parts: list[np.ndarray] = []
    i_parts: list[np.ndarray] = []
    out_sr = 44100
    logs = list(notes)

    for n, start in enumerate(starts):
        end = min(wav.size(-1), start + chunk_samples)
        part = wav[:, start:end]
        src = OUT_DIR / f"vr_src_{n}.wav"
        torchaudio.save(str(src), part.cpu(), sr)
        try:
            vp, ip, log = _separate_chunk(str(src), model_name)
        except Exception as e:
            return None, None, "\n".join(logs) + "\n" + _gpu_err(e)
        logs.append(f"[{n + 1}/{len(starts)}] {log}")
        if not vp or not Path(vp).is_file():
            return None, None, "\n".join(logs) + "\nGPU finished but vocals file missing."
        yv, out_sr = sf.read(vp, dtype="float32", always_2d=False)
        v_parts.append(np.asarray(yv, dtype=np.float32).reshape(-1))
        if ip and Path(ip).is_file():
            yi, _ = sf.read(ip, dtype="float32", always_2d=False)
            i_parts.append(np.asarray(yi, dtype=np.float32).reshape(-1))

    vocals = _stitch(v_parts, int(out_sr), CHUNK_OVERLAP_SEC)
    stamp = int(time.time())
    v_out = OUT_DIR / f"SOTAKA_Vocals_{stamp}.wav"
    sf.write(v_out, vocals, int(out_sr))
    inst = None
    if i_parts and len(i_parts) == len(v_parts):
        instrumental = _stitch(i_parts, int(out_sr), CHUNK_OVERLAP_SEC)
        i_out = OUT_DIR / f"SOTAKA_Instrumental_{stamp}.wav"
        sf.write(i_out, instrumental, int(out_sr))
        inst = str(i_out)

    log = (
        f"Engine: Demucs ({model_name}, chunked ZeroGPU)\n"
        + "\n".join(logs)
        + f"\nduration={len(vocals) / out_sr:.1f}s | sr={out_sr}\n"
        + "Tip: use Vocals as Song Convert source."
    )
    return str(v_out), inst, log