import logging
from pathlib import Path

from audio_separator.separator import Separator

logger = logging.getLogger(__name__)

STEM_NAMES = ("vocals", "drums", "bass", "other")
_FALLBACK_MODEL = "htdemucs_ft.yaml"


def separate_stems(audio_path: Path, output_dir: Path, model: str) -> dict[str, Path]:
    """Run audio-separator on the given file.

    Returns a dict mapping stem name to output file path,
    e.g. {"vocals": Path(...), "drums": Path(...), "bass": Path(...), "other": Path(...)}.

    If the requested model produces only 2 stems (vocals + instrumental),
    automatically falls back to htdemucs_ft for 4-stem separation.
    """
    stems = _run_separator(audio_path, output_dir, model)

    if len(stems) != 4 and model != _FALLBACK_MODEL:
        logger.warning(
            "Model '%s' produced %d stems %s instead of 4. "
            "Falling back to '%s' for 4-stem separation.",
            model, len(stems), list(stems.keys()), _FALLBACK_MODEL,
        )
        # Clean up partial output
        for p in stems.values():
            if p.exists():
                p.unlink()
        stems = _run_separator(audio_path, output_dir, _FALLBACK_MODEL)

    if len(stems) != 4:
        raise RuntimeError(
            f"Expected 4 stems but got {len(stems)}: {list(stems.keys())}. "
            f"Neither '{model}' nor fallback '{_FALLBACK_MODEL}' produced 4 stems."
        )

    for name, path in stems.items():
        logger.info("  %s: %s", name, path.name)

    return stems


def _run_separator(audio_path: Path, output_dir: Path, model: str) -> dict[str, Path]:
    """Run audio-separator with the given model and return parsed stems."""
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info("Separating stems for %s using model %s...", audio_path.name, model)

    separator = Separator(
        output_dir=str(output_dir),
        output_format="WAV",
        output_bitrate=None,
        normalization_threshold=1.0,
        sample_rate=44100,
    )
    separator.load_model(model_filename=model)
    output_files = separator.separate(str(audio_path))

    return _parse_stem_files(output_files, output_dir)


def _parse_stem_files(output_files: list[str], output_dir: Path) -> dict[str, Path]:
    """Parse audio-separator output file list into a stem name -> path dict."""
    stems: dict[str, Path] = {}
    for fpath in output_files:
        p = Path(fpath)
        # audio-separator may return bare filenames; ensure full path
        if not p.is_absolute():
            p = output_dir / p
        name_lower = p.stem.lower()
        if "vocal" in name_lower:
            stems["vocals"] = p
        elif "drum" in name_lower:
            stems["drums"] = p
        elif "bass" in name_lower:
            stems["bass"] = p
        else:
            stems["other"] = p
    return stems
