import os
import shutil
import subprocess
import tempfile
from pathlib import Path

import gradio as gr
from pyharp import ModelCard, build_endpoint

MODELS_DIR = Path("/app/models")
MODEL_SLUG = "roformer-model-bs-roformer-sw-by-jarredou"
MODEL_DIR = MODELS_DIR / MODEL_SLUG

STEMS = ["vocals", "drums", "bass", "guitar", "piano", "other", "instrumental"]

def get_model_paths():
    configs = list(MODEL_DIR.glob("*.yaml"))
    checkpoints = list(MODEL_DIR.glob("*.ckpt"))
    if not configs or not checkpoints:
        raise FileNotFoundError(f"Model files not found in {MODEL_DIR}")
    return str(configs[0]), str(checkpoints[0])

model_card = ModelCard(
    name="BS-RoFormer Source Separation",
    description="Separate audio into stems (vocals, drums, bass, guitar, piano, other) using the Band-Split RoPE Transformer.",
    author="Wei-Tsung Lu, Ju-Chiang Wang, Qiuqiang Kong, Yun-Ning Hung (ByteDance)",
    tags=["source-separation", "stems", "vocals", "drums", "bass"],
)

def process_fn(input_audio_path: str, stem: str) -> str:
    print(f"Separating stem: {stem}...", flush=True)

    input_dir = Path(tempfile.mkdtemp())
    output_dir = Path(tempfile.mkdtemp())

    input_path = Path(input_audio_path)
    shutil.copy(input_path, input_dir / input_path.name)

    config_path, model_path = get_model_paths()
    print(f"Config: {config_path}", flush=True)
    print(f"Model: {model_path}", flush=True)
    print(f"Input dir contents: {list(input_dir.iterdir())}", flush=True)

    result = subprocess.run([
        "bs-roformer-infer",
        "--config_path", config_path,
        "--model_path", model_path,
        "--input_folder", str(input_dir),
        "--store_dir", str(output_dir),
    ], capture_output=True, text=True)

    print(f"STDOUT:\n{result.stdout}", flush=True)
    print(f"STDERR:\n{result.stderr}", flush=True)

    if result.returncode != 0:
        raise RuntimeError(f"bs-roformer-infer failed:\n{result.stderr}")

    stem_files = list(output_dir.glob(f"*_{stem}.wav"))
    if not stem_files:
        available = [f.name for f in output_dir.glob("*.wav")]
        raise ValueError(f"Stem '{stem}' not found. Available: {available}")

    output_path = tempfile.mktemp(suffix=".wav")
    shutil.copy(stem_files[0], output_path)
    print("Done.", flush=True)
    return output_path

with gr.Blocks() as demo:
    input_components = [
        gr.Audio(
            type="filepath",
            label="Input Audio",
        ),
        gr.Dropdown(
            choices=STEMS,
            value="vocals",
            label="Stem to Extract",
        ),
    ]
    output_components = [
        gr.Audio(
            type="filepath",
            label="Separated Stem",
        ).set_info("Extracted stem audio."),
    ]
    app = build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

print("Launching Gradio...", flush=True)
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)