#!/usr/bin/env python3
import argparse
import contextlib
import json
import sys
from pathlib import Path


VALID_MODELS = {"htdemucs", "htdemucs_6s", "mdx_extra"}


@contextlib.contextmanager
def _stdout_to_stderr():
    """Redirect sys.stdout to stderr so Demucs progress doesn't pollute our JSON channel."""
    original = sys.stdout
    sys.stdout = sys.stderr
    try:
        yield
    finally:
        sys.stdout = original


def _progress(message: str):
    print(json.dumps({"type": "progress", "message": message}), file=sys.stderr, flush=True)


def separate(input_file: str, model: str, output_dir: str) -> dict:
    input_path = Path(input_file)
    if not input_path.exists():
        return {"success": False, "error": f"Input file not found: {input_file}", "stems": None}
    if not input_path.is_file():
        return {"success": False, "error": f"Input path is not a file: {input_file}", "stems": None}

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    try:
        import demucs.separate

        _progress("Separating stems...")
        args = ["--name", model, "--out", output_dir, input_file]
        with _stdout_to_stderr():
            demucs.separate.main(args)

        stem_dir = Path(output_dir) / model / input_path.stem
        stems = {p.stem: str(p) for p in sorted(stem_dir.glob("*.wav"))}

        if not stems:
            return {"success": False, "error": f"No output stems found in {stem_dir}", "stems": None}

        _progress("Done.")
        return {"success": True, "error": None, "stems": stems}

    except Exception as e:
        return {"success": False, "error": str(e), "stems": None}


def main():
    parser = argparse.ArgumentParser(description="Stem separation via Demucs")
    parser.add_argument("--input-file", required=True, help="Path to the audio file to separate")
    parser.add_argument("--model", required=True, choices=sorted(VALID_MODELS), help="Demucs model name")
    parser.add_argument("--output-dir", required=True, help="Directory to write separated stems into")
    args = parser.parse_args()

    result = separate(args.input_file, args.model, args.output_dir)
    print(json.dumps(result))
    sys.exit(0 if result["success"] else 1)


if __name__ == "__main__":
    main()
