# app.py — Hugging Face Space for SonicMaster AI Audio Mastering
# Optimized for HF Spaces ZeroGPU with pre-downloaded weights

import os
import sys
import subprocess
import shutil
import zipfile
import urllib.request
import threading
from pathlib import Path

import gradio as gr
import numpy as np
import soundfile as sf

os.environ.setdefault("GRADIO_USE_CDN", "true")

SPACE_ROOT = Path(__file__).parent.resolve()
REPO_DIR = SPACE_ROOT / "SonicMasterRepo"
REPO_URL = "https://github.com/AMAAI-Lab/SonicMaster"
CACHE_DIR = SPACE_ROOT / "weights"
CACHE_DIR.mkdir(parents=True, exist_ok=True)

MASTERING_PRESETS = {
    "Auto Enhance": "Enhance the input audio",
    "Dereverb": "Please, dereverb this audio.",
    "Brighten / Treble Boost": "Increase the clarity of this song by emphasizing treble frequencies.",
    "Bass Boost": "Make this song sound more boomy by amplifying the low end bass frequencies.",
    "Louder / Maximize": "Can you make this sound louder, please?",
    "Reduce Distortion": "Make the audio smoother and less distorted.",
    "Balance Mix": "Improve the balance in this song.",
    "Widen Stereo": "Disentangle the left and right channels to give this song a stereo feeling.",
    "Reduce Echo / Roominess": "Correct the unnatural frequency emphasis. Reduce the roominess or echo.",
    "Boost Vocals": "Raise the level of the vocals, please.",
    "Open Up / Less Squashed": "Make the sound less squashed and more open.",
    "Fix Frequency Issues": "Correct the unnatural frequency emphasis.",
}

# Global state
_repo_ready = False
_weights_path = None
_weights_loaded = threading.Event()
_load_status = "Starting..."


def _rmtree(p: Path):
    try:
        if p.exists():
            shutil.rmtree(p)
    except Exception:
        pass


def _safe_unlink(p: Path):
    try:
        if p.exists():
            p.unlink()
    except Exception:
        pass


def ensure_repo(progress=None) -> Path:
    global _repo_ready
    script0 = REPO_DIR / "infer_single.py"
    if _repo_ready and script0.exists():
        if REPO_DIR.as_posix() not in sys.path:
            sys.path.append(REPO_DIR.as_posix())
        return REPO_DIR

    if REPO_DIR.exists() and not script0.exists():
        _rmtree(REPO_DIR)

    if not REPO_DIR.exists():
        if progress:
            progress(0.02, desc="Fetching SonicMaster code")
        git_bin = shutil.which("git")
        cloned_ok = False
        if git_bin:
            try:
                subprocess.run(
                    [git_bin, "clone", "--depth", "1", REPO_URL, REPO_DIR.as_posix()],
                    check=True, capture_output=True, text=True,
                )
                cloned_ok = True
            except Exception:
                cloned_ok = False
                if REPO_DIR.exists():
                    _rmtree(REPO_DIR)

        if not cloned_ok:
            if progress:
                progress(0.05, desc="Downloading SonicMaster ZIP")
            zip_url = "https://codeload.github.com/AMAAI-Lab/SonicMaster/zip/refs/heads/main"
            zip_path = SPACE_ROOT / "SonicMaster.zip"
            urllib.request.urlretrieve(zip_url, zip_path.as_posix())
            if progress:
                progress(0.08, desc="Extracting SonicMaster ZIP")
            with zipfile.ZipFile(zip_path, "r") as zf:
                zf.extractall(SPACE_ROOT)
            extracted = SPACE_ROOT / "SonicMaster-main"
            if extracted.exists() and (extracted / "infer_single.py").exists():
                if REPO_DIR.exists():
                    _rmtree(REPO_DIR)
                extracted.rename(REPO_DIR)
            else:
                for cand in SPACE_ROOT.glob("SonicMaster-*"):
                    if cand.is_dir() and (cand / "infer_single.py").exists():
                        if REPO_DIR.exists():
                            _rmtree(REPO_DIR)
                        cand.rename(REPO_DIR)
                        break
            _safe_unlink(zip_path)

    if not (REPO_DIR / "infer_single.py").exists():
        raise RuntimeError("SonicMaster code fetch finished but infer_single.py is missing.")

    if REPO_DIR.as_posix() not in sys.path:
        sys.path.append(REPO_DIR.as_posix())
    _repo_ready = True
    return REPO_DIR


def get_weights_path(progress=None) -> Path:
    global _weights_path
    if _weights_path is None:
        if progress:
            progress(0.10, desc="Locating model weights")
        from huggingface_hub import hf_hub_download
        wp = hf_hub_download(
            repo_id="amaai-lab/SonicMaster",
            filename="model.safetensors",
            local_dir=str(CACHE_DIR),
            local_dir_use_symlinks=False,
            force_download=False,
            resume_download=True,
        )
        _weights_path = Path(wp)
    return _weights_path


def preload_weights():
    """Pre-download weights at startup so they're cached for inference."""
    global _load_status, _weights_path
    try:
        _load_status = "Downloading model weights (~3 GB)... This happens once."
        print(_load_status)
        ensure_repo()
        wp = get_weights_path()
        _weights_path = wp
        _load_status = f"Model weights ready: {wp}"
        print(_load_status)
    except Exception as e:
        _load_status = f"Weight download failed: {e}"
        print(_load_status)
    finally:
        _weights_loaded.set()


def run_inference_with_extra(input_wav: Path, prompt: str, output_wav: Path, extra_args: list, progress=None) -> tuple[bool, str]:
    ensure_repo(progress=progress)
    prompt = (prompt or "").strip() or "Enhance the input audio"
    if progress:
        progress(0.14, desc="Preparing inference")
    ckpt = get_weights_path(progress=progress)

    script = REPO_DIR / "infer_single.py"
    if not script.exists():
        return False, "infer_single.py not found."

    py = sys.executable or "python3"
    env = os.environ.copy()
    env["PYTHONDONTWRITEBYTECODE"] = "1"
    cwd = REPO_DIR.as_posix()

    cmd = [
        py, script.as_posix(),
        "--ckpt", ckpt.as_posix(),
        "--input", input_wav.as_posix(),
        "--prompt", prompt,
        "--output", output_wav.as_posix(),
    ] + extra_args

    try:
        if progress:
            progress(0.20, desc="Running SonicMaster inference on GPU...")
        res = subprocess.run(cmd, capture_output=True, text=True, check=True, env=env, cwd=cwd, timeout=600)
        if output_wav.exists() and output_wav.stat().st_size > 0:
            if progress:
                progress(0.95, desc="Done!")
            stdout = (res.stdout or "").strip()
            return True, stdout or "Mastering completed."
        return False, "Inference finished but produced no output file."
    except subprocess.TimeoutExpired:
        return False, "Inference timed out (600s). Try shorter audio or fewer steps."
    except subprocess.CalledProcessError as e:
        snippet = "\n".join(filter(None, [e.stdout or "", e.stderr or ""])).strip()
        return False, snippet or f"Failed with return code {e.returncode}."
    except Exception as e:
        import traceback
        return False, f"Unexpected error: {e}\n{traceback.format_exc()}"


def read_audio(path: str) -> tuple[np.ndarray, int]:
    wav, sr = sf.read(path, always_2d=False)
    if wav.dtype == np.float64:
        wav = wav.astype(np.float32)
    return wav, sr


def save_wav(wav: np.ndarray, sr: int, path: Path):
    if wav.ndim == 2 and wav.shape[0] < wav.shape[1]:
        wav = wav.T
    if wav.dtype == np.float64:
        wav = wav.astype(np.float32)
    sf.write(path.as_posix(), wav, sr)


def get_load_status():
    return _load_status


def master_audio(
    audio_path: str,
    preset: str,
    custom_prompt: str,
    steps: int,
    guidance: float,
    chunk_sec: int,
    overlap_sec: int,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        if not audio_path:
            raise gr.Error("Please upload an audio file to master.")

        # Wait for weights to be ready
        if not _weights_loaded.is_set():
            return None, "Model weights are still downloading. Please wait and try again in a moment."

        prompt = custom_prompt.strip() if custom_prompt.strip() else (MASTERING_PRESETS.get(preset, "") or "Enhance the input audio")

        if progress:
            progress(0.01, desc="Reading input audio")
        wav, sr = read_audio(audio_path)

        tmp_in = SPACE_ROOT / "tmp_master_in.wav"
        tmp_out = SPACE_ROOT / "tmp_master_out.wav"
        _safe_unlink(tmp_out)
        save_wav(wav, sr, tmp_in)

        extra_args = [
            "--num_inference_steps", str(steps),
            "--guidance_scale", str(guidance),
            "--chunk_duration", str(chunk_sec),
            "--overlap_duration", str(overlap_sec),
        ]

        ok, msg = run_inference_with_extra(tmp_in, prompt, tmp_out, extra_args, progress=progress)

        if ok and tmp_out.exists() and tmp_out.stat().st_size > 0:
            out_wav, out_sr = read_audio(tmp_out.as_posix())
            return (out_sr, out_wav), f"Mastering complete!\n\nPrompt: {prompt}\n\n{msg}"
        else:
            return None, f"Mastering failed:\n{msg}"

    except gr.Error as e:
        return None, str(e)
    except Exception as e:
        import traceback
        return None, f"Unexpected error: {e}\n{traceback.format_exc()}"


# ================== Gradio UI ==================
with gr.Blocks(title="SonicMaster — AI Audio Mastering Studio") as demo:
    gr.Markdown(
        "## 🎧 SonicMaster — AI Audio Mastering Studio\n"
        "*Text-guided music restoration & mastering (ICML 2026)*\n\n"
        "Upload a track, pick a preset or write your own prompt, then hit **Master**.\n"
        "Runs on Hugging Face ZeroGPU — weights are pre-cached on startup."
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Input")
            in_audio = gr.Audio(label="Upload Track", type="filepath")

            gr.Markdown("### 🎛️ Preset")
            preset = gr.Dropdown(
                choices=list(MASTERING_PRESETS.keys()),
                value="Auto Enhance",
                label="Mastering Preset",
            )

            gr.Markdown("### ✏️ Custom Prompt")
            custom_prompt = gr.Textbox(
                label="Text Prompt (overrides preset if filled)",
                placeholder="e.g., Reduce reverb and brighten vocals...",
                lines=2,
            )

            gr.Markdown("### ⚙️ Advanced")
            with gr.Accordion("Inference Settings", open=False):
                steps = gr.Slider(5, 30, value=10, step=1, label="Inference Steps")
                guidance = gr.Slider(0.5, 3.0, value=1.0, step=0.1, label="Guidance Scale")
                chunk_sec = gr.Slider(10, 60, value=30, step=5, label="Chunk Duration (sec)")
                overlap_sec = gr.Slider(2, 20, value=10, step=1, label="Overlap Duration (sec)")

            run_btn = gr.Button("🚀 Master Track", variant="primary", size="lg")

        with gr.Column(scale=1):
            gr.Markdown("### 📤 Output")
            out_audio = gr.Audio(label="Mastered Audio", interactive=False)
            status = gr.Textbox(label="Status", interactive=False, lines=6)

    gr.Markdown(
        "### 💡 Prompt Examples\n"
        + "\n".join(f"- `{p}`" for p in list(MASTERING_PRESETS.values())[:6])
    )

    run_btn.click(
        fn=master_audio,
        inputs=[in_audio, preset, custom_prompt, steps, guidance, chunk_sec, overlap_sec],
        outputs=[out_audio, status],
        concurrency_limit=1,
    )

if __name__ == "__main__":
    # Start weight download in background thread
    t = threading.Thread(target=preload_weights, daemon=True)
    t.start()

    demo.queue(max_size=4).launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(primary_hue="violet", secondary_hue="indigo"),
    )