"""MuScriptor Studio — Gradio app wrapping muscriptor with optional demucs vocal-isolate.

Pipeline:
  audio (mp3/wav/m4a/flac)
    -> [optional] demucs htdemucs -> vocals.wav  (isolate vocals)
    -> muscriptor transcribe_and_postprocess -> .mid bytes
"""
from __future__ import annotations

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

import gradio as gr

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("muscriptor-studio")

# -------------------------------------------------------------------- muscriptor

_MUSCRIPTOR_OK = False
try:
    from muscriptor.transcription_model import TranscriptionModel  # type: ignore

    _MUSCRIPTOR_OK = True
    log.info("muscriptor imported OK")
except Exception as e:  # pragma: no cover
    log.warning("muscriptor unavailable: %s", e)
    TranscriptionModel = None  # type: ignore

_MODEL_CACHE: dict[str, "TranscriptionModel"] = {}
# Map UI size label -> muscriptor model id (small / medium / large)
_MODEL_SIZES = {"small": "small", "medium": "medium", "large": "large"}


def _device() -> str:
    """Pick the best device available in this build."""
    try:
        import torch  # type: ignore

        if torch.cuda.is_available():
            return "cuda"
        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
            return "mps"
    except Exception:
        pass
    return "cpu"


def _dtype_for(device: str) -> str:
    return "float16" if device in {"cuda", "mps"} else "float32"


def _get_model(size: str) -> "TranscriptionModel":
    """Lazy-load muscriptor (downloads weights on first call)."""
    if not _MUSCRIPTOR_OK or TranscriptionModel is None:
        raise RuntimeError("muscriptor not installed in this build")
    if size not in _MODEL_CACHE:
        device = _device()
        log.info("loading muscriptor size=%s device=%s dtype=%s", size, device, _dtype_for(device))
        _MODEL_CACHE[size] = TranscriptionModel.load_model(
            weights_path=_MODEL_SIZES[size],
            device=device,
            dtype=_dtype_for(device),
        )
    return _MODEL_CACHE[size]


def _isolate_vocals(src: Path, workdir: Path) -> Path:
    """Run demucs --two-stems vocals; return path to vocals.wav."""
    log.info("demucs htdemucs on %s", src.name)
    proc = subprocess.run(
        [
            "python", "-m", "demucs",
            "--two-stems", "vocals",
            "-n", "htdemucs",
            "--device", _device(),
            "-o", str(workdir / "demucs"),
            str(src),
        ],
        capture_output=True,
        text=True,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"demucs failed:\n{proc.stderr[-1000:]}")
    sep = workdir / "demucs" / "separated" / "htdemucs" / src.stem / "vocals.wav"
    if not sep.exists():
        raise RuntimeError(f"demucs produced no vocals.wav at {sep}")
    return sep


def _transcribe(
    audio_path: str,
    model_size: str,
    instruments_csv: str,
    isolate_vocals: bool,
) -> str:
    """Run demucs (optional) + muscriptor; return path to .mid file."""
    if not audio_path:
        raise gr.Error("Upload an audio file first.")
    if not _MUSCRIPTOR_OK:
        raise gr.Error("muscriptor not available in this build.")

    workdir = Path(tempfile.mkdtemp(prefix="muscriptor-"))
    src = Path(audio_path)
    feed = workdir / src.name
    shutil.copy2(src, feed)

    if isolate_vocals:
        feed = _isolate_vocals(feed, workdir)

    instruments: list[str] | None = None
    if instruments_csv.strip():
        instruments = [s.strip() for s in instruments_csv.split(",") if s.strip()]
        log.info("constrained instruments=%s", instruments)

    model = _get_model(model_size)
    log.info("transcribe %s", feed)
    midi_bytes, beat_grid = model.transcribe_and_postprocess(
        str(feed),
        instruments=instruments,
        detect_tempo="best-effort",
        quantize=False,
    )
    if beat_grid is not None:
        log.info("detected tempo: bpm=%s", getattr(beat_grid, "bpm", "?"))

    final = Path(tempfile.gettempdir()) / (src.stem + ".muscriptor.mid")
    final.write_bytes(midi_bytes)
    log.info("midi ready: %s (%d bytes)", final, len(midi_bytes))
    return str(final)


# -------------------------------------------------------------------- UI

MODEL_CHOICES = ["small", "medium", "large"]

with gr.Blocks(title="MuScriptor Studio", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
        # MuScriptor Studio

        Audio → MIDI via [muscriptor](https://github.com/muscriptor/muscriptor)
        (transformer LM, Kyutai × Mirelo). Optional **vocal isolation**
        with [demucs](https://github.com/facebookresearch/demucs) before
        transcription — useful when the input has a heavy mix and you only
        want the lead voice.

        muscriptor available: **{_MUSCRIPTOR_OK}** ·
        device: **{_device()}**
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(
                label="Audio (mp3 / wav / m4a / flac)",
                type="filepath",
                sources=["upload"],
            )
            model_dd = gr.Dropdown(MODEL_CHOICES, value="medium", label="Model size")
            isolate_cb = gr.Checkbox(False, label="Isolate vocals first (demucs htdemucs)")
            instr_tb = gr.Textbox(
                "",
                label="Constrain instruments (comma-separated, blank = all)",
                placeholder="voice,acoustic_piano,accordion,drums",
            )
            run_btn = gr.Button("Transcribe → MIDI", variant="primary")

        with gr.Column(scale=1):
            midi_out = gr.File(label="MIDI download (.mid)")

    if _MUSCRIPTOR_OK:
        run_btn.click(
            _transcribe,
            inputs=[audio_in, model_dd, instr_tb, isolate_cb],
            outputs=[midi_out],
        )

    gr.Markdown(
        """
        **Tips**
        - `small` is fastest, `large` is highest quality. On CPU, default
          to `small`; on GPU, `medium` is the sweet spot.
        - Tick *Isolate vocals* if the mix is busy — voice is decoded
          best when it stands alone.
        - Leave the instrument constraint blank to let muscriptor detect
          whatever's there. Constrain it if you know the instrumentation
          (e.g. Balkan folk = `voice,acoustic_piano,accordion,drums`).
        - First run downloads model weights (~hundreds of MB).
        """
    )

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