# app.py — end-to-end Hugging Face Spaces app for SonicMaster
# - Works even if `git` is NOT available (falls back to GitHub ZIP download)
# - Lazily fetches code + weights only when user clicks "Enhance"
# - Runs inference ONLY via infer_single.py

import os
import sys
import subprocess
import shutil
import zipfile
import urllib.request
from pathlib import Path
from typing import Tuple, Optional, List, Any

# Make Gradio assets reliable on Spaces
os.environ.setdefault("GRADIO_USE_CDN", "true")

# --- HF Spaces SDK (optional) ---
try:
    import spaces  # HF Spaces SDK
except Exception:
    class _DummySpaces:
        def GPU(self, *_, **__):
            def deco(fn):
                return fn
            return deco
    spaces = _DummySpaces()

import gradio as gr
import numpy as np
import soundfile as sf
from huggingface_hub import hf_hub_download

# ================= Runtime hints (safe on CPU) =================
USE_ZEROGPU = os.getenv("SPACE_RUNTIME", "").lower() == "zerogpu"

SPACE_ROOT = Path(__file__).parent.resolve()
REPO_DIR = SPACE_ROOT / "SonicMasterRepo"
REPO_URL = "https://github.com/AMAAI-Lab/SonicMaster"
WEIGHTS_REPO = "amaai-lab/SonicMaster"
WEIGHTS_FILE = "model.safetensors"

CACHE_DIR = SPACE_ROOT / "weights"
CACHE_DIR.mkdir(parents=True, exist_ok=True)

# ================== SAFE repo handling (NO network at import) ==================
_repo_ready: bool = False


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


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


def ensure_repo(progress: Optional[gr.Progress] = None) -> Path:
    """
    Ensure SonicMaster repo is available.
    Tries:
      1) git clone (if git exists)
      2) GitHub ZIP download + extract (if git missing or clone fails)
    Called lazily (on button click), not at import time.
    """
    global _repo_ready

    # If already ready and infer exists, done
    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 directory exists but script missing, treat as corrupted and reset
    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 (first run)")

        git_bin = shutil.which("git")
        cloned_ok = False

        # Try git clone if available
        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 partial directory created, remove and retry via ZIP
                if REPO_DIR.exists():
                    _rmtree(REPO_DIR)

        # Fallback: download ZIP
        if not cloned_ok:
            if progress:
                progress(0.05, desc="Downloading SonicMaster ZIP (no git / clone failed)")

            zip_url = "https://codeload.github.com/AMAAI-Lab/SonicMaster/zip/refs/heads/main"
            zip_path = SPACE_ROOT / "SonicMaster.zip"

            # Download ZIP
            urllib.request.urlretrieve(zip_url, zip_path.as_posix())

            if progress:
                progress(0.08, desc="Extracting SonicMaster ZIP")

            # Extract ZIP to SPACE_ROOT
            with zipfile.ZipFile(zip_path, "r") as zf:
                zf.extractall(SPACE_ROOT)

            # GitHub zip usually extracts to SonicMaster-main/
            extracted = SPACE_ROOT / "SonicMaster-main"
            if extracted.exists() and (extracted / "infer_single.py").exists():
                # If destination already exists, remove it
                if REPO_DIR.exists():
                    _rmtree(REPO_DIR)
                extracted.rename(REPO_DIR)
            else:
                # If structure differs, try to locate the folder containing infer_single.py
                found = None
                for cand in SPACE_ROOT.glob("SonicMaster-*"):
                    if cand.is_dir() and (cand / "infer_single.py").exists():
                        found = cand
                        break
                if found:
                    if REPO_DIR.exists():
                        _rmtree(REPO_DIR)
                    found.rename(REPO_DIR)

            _safe_unlink(zip_path)

    # Final sanity check
    if not (REPO_DIR / "infer_single.py").exists():
        existing = sorted([p.name for p in REPO_DIR.glob("*")]) if REPO_DIR.exists() else []
        raise RuntimeError(
            "SonicMaster code fetch finished, but infer_single.py is still missing.\n"
            f"REPO_DIR: {REPO_DIR}\n"
            f"Contents: {existing[:80]}"
        )

    if REPO_DIR.as_posix() not in sys.path:
        sys.path.append(REPO_DIR.as_posix())

    _repo_ready = True
    return REPO_DIR


# ================ Weights: lazy download (first click) ================
_weights_path: Optional[Path] = None


def get_weights_path(progress: Optional[gr.Progress] = None) -> Path:
    """
    Download/resolve weights lazily (keeps startup fast).
    """
    global _weights_path
    if _weights_path is None:
        if progress:
            progress(0.10, desc="Downloading model weights (first run)")
        wp = hf_hub_download(
            repo_id=WEIGHTS_REPO,
            filename=WEIGHTS_FILE,
            local_dir=str(CACHE_DIR),
            local_dir_use_symlinks=False,
            force_download=False,
            resume_download=True,
        )
        _weights_path = Path(wp)
    return _weights_path


# ================== Audio helpers ==================
def save_temp_wav(wav: np.ndarray, sr: int, path: Path):
    # Ensure shape (samples, channels)
    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 read_audio(path: str) -> Tuple[np.ndarray, int]:
    wav, sr = sf.read(path, always_2d=False)
    if isinstance(wav, np.ndarray) and wav.dtype == np.float64:
        wav = wav.astype(np.float32)
    return wav, sr


def _has_cuda() -> bool:
    try:
        import torch
        return torch.cuda.is_available()
    except Exception:
        return False


# ================== CLI runner ==================
def _candidate_commands(
    py: str, script: Path, ckpt: Path, inp: Path, prompt: str, out: Path
) -> List[List[str]]:
    """
    Only support infer_single.py variants.
    Expected flags: --ckpt --input --prompt --output
    """
    return [[
        py,
        script.as_posix(),
        "--ckpt", ckpt.as_posix(),
        "--input", inp.as_posix(),
        "--prompt", prompt,
        "--output", out.as_posix(),
    ]]


def run_sonicmaster_cli(
    input_wav_path: Path,
    prompt: str,
    out_path: Path,
    progress: Optional[gr.Progress] = None,
) -> Tuple[bool, str]:
    """
    Run inference via subprocess; returns (ok, message).
    Uses ONLY infer_single.py.
    """
    # Ensure repo is present when needed (NOT at startup)
    ensure_repo(progress=progress)

    # Ensure a non-empty prompt for the CLI
    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 in the SonicMaster repo (code fetch likely failed)."

    py = sys.executable or "python3"
    env = os.environ.copy()

    # Make sure subprocess runs inside repo (some projects rely on relative paths)
    cwd = REPO_DIR.as_posix()

    last_err = ""
    for cidx, cmd in enumerate(_candidate_commands(py, script, ckpt, input_wav_path, prompt, out_path), 1):
        try:
            if progress:
                progress(min(0.25 + 0.10 * cidx, 0.70), desc=f"Running inference (try {cidx})")
            res = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True,
                env=env,
                cwd=cwd,
            )

            if out_path.exists() and out_path.stat().st_size > 0:
                if progress:
                    progress(0.88, desc="Post-processing output")
                stdout = (res.stdout or "").strip()
                return True, (stdout or "Inference completed.")
            last_err = "infer_single.py finished but produced no output file."
        except subprocess.CalledProcessError as e:
            snippet = "\n".join(filter(None, [e.stdout or "", e.stderr or ""])).strip()
            last_err = snippet if snippet else f"infer_single.py failed with return code {e.returncode}."
        except Exception as e:
            import traceback
            last_err = f"Unexpected error: {e}\n{traceback.format_exc()}"

    return False, last_err or "Inference failed."


# ============ GPU path (ZeroGPU) ============
@spaces.GPU(duration=60)  # safe cap for ZeroGPU tiers
def enhance_on_gpu(input_path: str, prompt: str, output_path: str) -> Tuple[bool, str]:
    from pathlib import Path as _P
    return run_sonicmaster_cli(_P(input_path), prompt, _P(output_path), progress=None)


# ================== Optional Examples (NO FETCH AT STARTUP) ==================
PROMPTS_10 = [
    "Increase the clarity of this song by emphasizing treble frequencies.",
    "Make this song sound more boomy by amplifying the low end bass frequencies.",
    "Can you make this sound louder, please?",
    "Make the audio smoother and less distorted.",
    "Improve the balance in this song.",
    "Disentangle the left and right channels to give this song a stereo feeling.",
    "Correct the unnatural frequency emphasis. Reduce the roominess or echo.",
    "Raise the level of the vocals, please.",
    "Increase the clarity of this song by emphasizing treble frequencies.",
    "Please, dereverb this audio.",
]


def build_examples_if_repo_present() -> List[List[Any]]:
    """
    Build examples WITHOUT cloning/downloading.
    If repo isn't present yet, return [].
    """
    wav_dir = REPO_DIR / "samples" / "inputs"
    if not wav_dir.exists():
        return []
    wav_paths = sorted(p for p in wav_dir.glob("*.wav") if p.is_file())
    ex: List[List[Any]] = []
    for i, p in enumerate(wav_paths[:10]):
        pr = PROMPTS_10[i] if i < len(PROMPTS_10) else PROMPTS_10[-1]
        ex.append([p.as_posix(), pr])
    return ex


# ================== Main callback ==================
def enhance_audio_ui(
    audio_path: str,
    prompt: str,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Returns (audio, message). On failure, audio=None and message=error text.
    """
    try:
        prompt = (prompt or "").strip() or "Enhance the input audio"
        if not audio_path:
            raise gr.Error("Please upload or select an input audio file.")

        if progress:
            progress(0.03, desc="Preparing audio")
        wav, sr = read_audio(audio_path)

        tmp_in = SPACE_ROOT / "tmp_in.wav"
        tmp_out = SPACE_ROOT / "tmp_out.wav"

        # Clean previous output to avoid stale returns
        if tmp_out.exists():
            _safe_unlink(tmp_out)

        save_temp_wav(wav, sr, tmp_in)

        use_gpu_call = USE_ZEROGPU or _has_cuda()
        if progress:
            progress(0.12, desc="Starting inference")

        if use_gpu_call:
            ok, msg = enhance_on_gpu(tmp_in.as_posix(), prompt, tmp_out.as_posix())
        else:
            ok, msg = run_sonicmaster_cli(tmp_in, prompt, tmp_out, 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), (msg or "Done.")
        else:
            return None, (msg or "Inference failed without a specific error message.")

    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 – Text-Guided Restoration & Mastering", fill_height=True) as _demo:
    gr.Markdown(
        "## 🎧 SonicMaster\n"
        "Upload audio, write a prompt (or leave blank), then click **Enhance**.\n"
        "If left blank, we use: _Enhance the input audio_.\n\n"
        "- First run will fetch code + download weights.\n"
        "- Subsequent runs are much faster.\n"
    )

    with gr.Row():
        with gr.Column(scale=1):
            in_audio = gr.Audio(label="Input Audio", type="filepath")
            prompt_box = gr.Textbox(
                label="Text Prompt",
                placeholder="e.g., Reduce reverb and brighten vocals. (Optional)",
            )
            run_btn = gr.Button("🚀 Enhance", variant="primary")

            # Examples only if already present locally (no startup fetch)
            examples = build_examples_if_repo_present()
            if examples:
                gr.Examples(examples=examples, inputs=[in_audio, prompt_box], label="Sample Inputs (10)")
            else:
                gr.Markdown("> ℹ️ Samples will appear after the code is fetched (first run).")

        with gr.Column(scale=1):
            out_audio = gr.Audio(label="Enhanced Audio (output)")
            status = gr.Textbox(label="Status / Messages", interactive=False, lines=8)

    run_btn.click(
        fn=enhance_audio_ui,
        inputs=[in_audio, prompt_box],
        outputs=[out_audio, status],
        concurrency_limit=1,
    )

demo = _demo.queue(max_size=16)
iface = demo
app = demo

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)