# Copyright (c) 2025-2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
from pathlib import Path
import os, tempfile
import numpy as np
import soundfile as sf
import librosa
import torch
import gc

from audio_separator.separator import Separator
from shared.utils import files_locator as fl
from shared.utils.download import process_files_def_if_needed
from shared.utils.default_device import keep_default_device
from preprocessing.roformer.assets import query_download_def
from shared.utils.media_control import controlled_model_loading, checkpoint_modules


@controlled_model_loading
def _load_separator(separator):
    separator.load_model()

def _resolve_separator_output(out_file, output_dir: Path) -> Path:
    out = Path(out_file)
    return out if out.is_absolute() else output_dir / out


def _separator_stem_key(path_or_stem) -> str:
    return "_".join(part for part in Path(path_or_stem).stem.lower().split("_") if part)


def _prepare_separator_input(src_path: str, min_seconds: float):
    duration = librosa.get_duration(path=src_path)
    use_path = src_path
    temp_path = None
    original_duration = None

    if duration < min_seconds:
        # Load (resample) and pad in memory
        y, sr = librosa.load(src_path, sr=None, mono=False)
        if y.ndim == 1:  # ensure shape (channels, samples)
            y = y[np.newaxis, :]
        original_duration = y.shape[1] / sr
        target_len = int(min_seconds * sr)
        pad = max(0, target_len - y.shape[1])
        if pad:
            y = np.pad(y, ((0, 0), (0, pad)), mode="constant")

        # Write a temp WAV for the separator
        fd, temp_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)
        sf.write(temp_path, y.T, sr)  # soundfile expects (frames, channels)
        use_path = temp_path

    return use_path, temp_path, original_duration


@keep_default_device()
def extract_vocal_and_background_stems(src_path: str, vocals_dst_path: str, background_dst_path: str, min_seconds: float = 8) -> tuple[str, str]:
    process_files_def_if_needed(query_download_def())
    torch.set_default_device('cpu')

    vocals_dst = Path(vocals_dst_path)
    background_dst = Path(background_dst_path)
    vocals_dst.parent.mkdir(parents=True, exist_ok=True)
    background_dst.parent.mkdir(parents=True, exist_ok=True)

    sep = None
    temp_path = None
    try:
        use_path, temp_path, _ = _prepare_separator_input(src_path, min_seconds)
        sep = Separator(
            output_dir=str(vocals_dst.parent),
            output_format=(vocals_dst.suffix.lstrip(".") or "wav"),
            model_file_dir=fl.locate_folder("roformer")
        )
        _load_separator(sep)
        with checkpoint_modules(module for module in sep.model_instance.model_run.modules() if not module._modules):
            out_files = sep.separate(use_path, {"Vocals": vocals_dst.stem, "Instrumental": background_dst.stem})
        out_paths = [_resolve_separator_output(out_file, vocals_dst.parent) for out_file in out_files]
        by_stem = {_separator_stem_key(out_path): str(out_path) for out_path in out_paths}
        return by_stem[_separator_stem_key(vocals_dst)], by_stem[_separator_stem_key(background_dst)]
    finally:
        if sep is not None:
            del sep
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)

        torch.cuda.empty_cache()
        gc.collect()


@keep_default_device()
def _get_single_stem(src_path: str, dst_path: str, stem: str, min_seconds: float = 8) -> str:
    """Separate audio and save only the requested stem to ``dst_path``."""

    process_files_def_if_needed(query_download_def())
    torch.set_default_device('cpu')

    dst = Path(dst_path)
    dst.parent.mkdir(parents=True, exist_ok=True)

    sep = None
    temp_path = None
    try:
        use_path, temp_path, original_duration = _prepare_separator_input(src_path, min_seconds)
        # Run separation: emit only the selected stem with the requested filename.
        sep = Separator(
            output_dir=str(dst.parent),
            output_format=(dst.suffix.lstrip(".") or "wav"),
            output_single_stem=stem,
            model_file_dir=fl.locate_folder("roformer")
        )
        _load_separator(sep)
        with checkpoint_modules(module for module in sep.model_instance.model_run.modules() if not module._modules):
            out_files = sep.separate(use_path, {stem: dst.stem})

        output_path = _resolve_separator_output(out_files[0], dst.parent)
        if original_duration is not None:
            output_info = sf.info(output_path)
            output_frames = min(output_info.frames, round(original_duration * output_info.samplerate))
            samples, _ = sf.read(output_path, frames=output_frames, always_2d=True)
            sf.write(output_path, samples, output_info.samplerate, subtype=output_info.subtype)
        return str(output_path)
    finally:
        if sep is not None:
            del sep
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)

        torch.cuda.empty_cache()
        gc.collect()


def get_vocals(src_path: str, dst_path: str, min_seconds: float = 8) -> str:
    """Save the vocal stem, padding short inputs for the separator if needed."""
    return _get_single_stem(src_path, dst_path, "Vocals", min_seconds)


def get_instrumental(src_path: str, dst_path: str, min_seconds: float = 8) -> str:
    """Save the instrumental stem, removing the separated vocals."""
    return _get_single_stem(src_path, dst_path, "Instrumental", min_seconds)

# Example:
# final = extract_vocals("in/clip.mp3", "out/vocals.wav")
# print(final)