import os

# ZeroGPU and library caches must be configured before any third-party import.
os.environ.setdefault("HF_HOME", "/tmp/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  # Must precede torch on ZeroGPU.

import base64
import hashlib
import json
import math
import re
import sys
import tempfile
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterator

import gradio as gr
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download
from muscriptor.events import NoteEndEvent, NoteStartEvent, ProgressEvent
from muscriptor.tokenizer.mt3 import MT3Tokenizer, MT3_FULL_PLUS_GROUP_NAMES
from muscriptor.transcription_model import (
    TranscriptionModel,
    _build_model,
    _remap_single_codebook_keys,
    _resolve_config,
    _resolve_source,
)
from safetensors.torch import load_file


APP_ROOT = Path(__file__).resolve().parent
ASSET_ROOT = APP_ROOT / "assets"
EXAMPLE_ROOT = APP_ROOT / "examples"
MODEL_ID = "MuScriptor/muscriptor-medium"
MODEL_VARIANT = "medium"
MAX_AUDIO_SECONDS = 60.0

# Calibrated on the deployed ZeroGPU Space with 12s / 3-window audio:
# 27.90s cold, 20.17s warm, 17.02s conditioned. round(27.90 * 1.4) = 39s.
GPU_BASE_SECONDS = 21
GPU_SECONDS_PER_CHUNK = 6
GPU_DURATION_CAP = 120

COLORS = (
    "#8b7cff",
    "#35c8ff",
    "#67e8a5",
    "#f6c667",
    "#fb7185",
    "#c084fc",
    "#2dd4bf",
    "#60a5fa",
    "#f472b6",
    "#a3e635",
)


@spaces.GPU(duration=1)
def _zerogpu_probe() -> str:
    """Required lightweight ZeroGPU allocation probe."""
    return "ready"


def _load_model_zerogpu() -> TranscriptionModel:
    """Build MuScriptor with CUDA attributes while loading weights on CPU.

    ZeroGPU intercepts ``.to('cuda')`` at module scope and packs the singleton
    for restoration inside decorated calls. Safetensors cannot target the fake
    CUDA device at startup, so its state dictionary must first be read on CPU.
    """

    token = os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError(
            "HF_TOKEN is required. Add a read token whose owner has accepted "
            f"the gated license at https://huggingface.co/{MODEL_ID}."
        )

    device = torch.device("cuda")
    source = _resolve_source(MODEL_VARIANT)
    weights_path = Path(
        hf_hub_download(
            repo_id=MODEL_ID,
            filename="model.safetensors",
            token=token,
        )
    )
    config = _resolve_config(source, weights_path)

    model = _build_model(device, config)
    model.eval()
    state_dict = load_file(str(weights_path), device="cpu")
    state_dict = _remap_single_codebook_keys(state_dict)
    model.load_state_dict(state_dict)
    model.to("cuda")

    tokenizer = MT3Tokenizer(
        instrument_vocabulary="MT3_FULL_PLUS",
        max_shift_steps=1001,
    )
    return TranscriptionModel(model=model, tokenizer=tokenizer, device=device)


print(f"[MuScriptor Studio] Loading {MODEL_VARIANT} model once at startup…", flush=True)
_model_load_started = time.perf_counter()
MODEL = _load_model_zerogpu()
print(
    "[MuScriptor Studio] Model ready in "
    f"{time.perf_counter() - _model_load_started:.2f}s.",
    flush=True,
)


def _audio_path(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, (str, Path)):
        return str(value)
    if isinstance(value, dict):
        path = value.get("path") or value.get("name")
        return str(path) if path else None
    path = getattr(value, "path", None) or getattr(value, "name", None)
    return str(path) if path else None


def _audio_duration(path: str | None) -> float | None:
    if not path:
        return None
    try:
        return float(sf.info(path).duration)
    except Exception:
        return None


def _estimate_gpu_duration(
    audio: Any,
    instruments: list[str] | None = None,
    use_sampling: bool = False,
    temperature: float = 1.0,
    beam_size: int = 1,
    *args: Any,
    **kwargs: Any,
) -> int:
    """Estimate allocation from 5-second chunks; accepts Gradio extras."""

    duration = _audio_duration(_audio_path(audio))
    chunks = max(1, math.ceil((duration or 15.0) / 5.0))
    beam = max(1, min(4, int(beam_size or 1)))
    return min(
        GPU_DURATION_CAP,
        GPU_BASE_SECONDS + chunks * GPU_SECONDS_PER_CHUNK * beam,
    )


def _display_name(name: str) -> str:
    if name.startswith("program_"):
        return f"MIDI program {name.removeprefix('program_')}"
    special = {
        "drums": "Drums",
        "voice": "Voice",
        "flutes": "Flutes",
        "soprano_and_alto_sax": "Soprano & alto sax",
    }
    return special.get(name, name.replace("_", " ").title())


def _color_for(name: str) -> str:
    digest = hashlib.sha1(name.encode("utf-8")).digest()
    return COLORS[int.from_bytes(digest[:2], "big") % len(COLORS)]


def _slug(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-") or "track"


def _data_uri(data: bytes) -> str:
    encoded = base64.b64encode(data).decode("ascii")
    return f"data:audio/midi;base64,{encoded}"


def _program_for(name: str) -> int:
    if name == "drums":
        return 0
    try:
        return int(MODEL._program_for_instrument(name))
    except Exception:
        if name.startswith("program_"):
            try:
                return int(name.removeprefix("program_"))
            except ValueError:
                pass
        return 0


def _notes_from_events(
    events: list[NoteStartEvent | NoteEndEvent | ProgressEvent],
) -> dict[str, list[dict[str, float | int]]]:
    notes: dict[str, list[dict[str, float | int]]] = defaultdict(list)
    for event in events:
        if not isinstance(event, NoteEndEvent):
            continue
        start = event.start_event
        notes[start.instrument].append(
            {
                "pitch": int(start.pitch),
                "start": round(float(start.start_time), 4),
                "end": round(max(float(event.end_time), float(start.start_time) + 0.01), 4),
                "velocity": 100,
            }
        )
    for instrument_notes in notes.values():
        instrument_notes.sort(key=lambda note: (note["start"], note["pitch"]))
    return dict(notes)


def _track_payloads(
    notes: dict[str, list[dict[str, float | int]]],
    requested: list[str] | None = None,
    midi_by_instrument: dict[str, bytes] | None = None,
) -> list[dict[str, Any]]:
    names = set(notes)
    names.update(requested or [])
    ordered = sorted(
        names,
        key=lambda name: (
            notes.get(name, [{}])[0].get("start", float("inf")) if notes.get(name) else float("inf"),
            name,
        ),
    )
    tracks: list[dict[str, Any]] = []
    for index, name in enumerate(ordered):
        track_notes = notes.get(name, [])
        midi = (midi_by_instrument or {}).get(name)
        tracks.append(
            {
                "id": index,
                "key": name,
                "name": _display_name(name),
                "color": _color_for(name),
                "note_count": len(track_notes),
                "program": _program_for(name),
                "is_drum": name == "drums",
                "midi": _data_uri(midi) if midi else "",
                "notes": track_notes,
            }
        )
    return tracks


def _payload(
    *,
    state: str,
    status: str,
    progress: float,
    audio_name: str,
    elapsed: float,
    duration: float,
    tracks: list[dict[str, Any]],
    full_midi: bytes | None = None,
) -> str:
    return json.dumps(
        {
            "state": state,
            "status": status,
            "progress": round(max(0.0, min(1.0, progress)), 4),
            "audio_name": audio_name,
            "elapsed": round(elapsed, 2),
            "duration": round(duration, 2),
            "note_count": sum(track["note_count"] for track in tracks),
            "full_midi": _data_uri(full_midi) if full_midi else "",
            "tracks": tracks,
        },
        separators=(",", ":"),
    )


def _event_is_for_instrument(
    event: NoteStartEvent | NoteEndEvent | ProgressEvent,
    instrument: str,
) -> bool:
    if isinstance(event, NoteStartEvent):
        return event.instrument == instrument
    if isinstance(event, NoteEndEvent):
        return event.start_event.instrument == instrument
    return False


def _write_outputs(
    events: list[NoteStartEvent | NoteEndEvent | ProgressEvent],
    track_names: list[str],
) -> tuple[Path, dict[str, Path], bytes, dict[str, bytes]]:
    request_dir = Path(tempfile.mkdtemp(prefix="muscriptor-output-"))
    full_bytes = MODEL.events_to_midi_bytes(iter(events))
    full_path = request_dir / "muscriptor-transcription.mid"
    full_path.write_bytes(full_bytes)

    track_paths: dict[str, Path] = {}
    midi_by_instrument: dict[str, bytes] = {}
    for name in track_names:
        filtered = [event for event in events if _event_is_for_instrument(event, name)]
        track_bytes = MODEL.events_to_midi_bytes(iter(filtered))
        track_path = request_dir / f"{_slug(name)}.mid"
        track_path.write_bytes(track_bytes)
        track_paths[name] = track_path
        midi_by_instrument[name] = track_bytes
    return full_path, track_paths, full_bytes, midi_by_instrument


def _manifest(
    *,
    state: str,
    audio_name: str,
    elapsed: float,
    duration: float,
    tracks: list[dict[str, Any]],
    full_path: Path | None = None,
    track_paths: dict[str, Path] | None = None,
    message: str | None = None,
) -> dict[str, Any]:
    result: dict[str, Any] = {
        "state": state,
        "model": MODEL_VARIANT,
        "model_id": MODEL_ID,
        "audio_name": audio_name,
        "audio_duration_seconds": round(duration, 3),
        "inference_seconds": round(elapsed, 3),
        "note_count": sum(track["note_count"] for track in tracks),
        "track_count": len(tracks),
        "tracks": [
            {
                "instrument": track["key"],
                "display_name": track["name"],
                "note_count": track["note_count"],
                "program": track["program"],
                "is_drum": track["is_drum"],
                "file": str(track_paths[track["key"]])
                if track_paths and track["key"] in track_paths
                else None,
            }
            for track in tracks
        ],
        "full_midi": str(full_path) if full_path else None,
    }
    if message:
        result["message"] = message
    return result


@spaces.GPU(duration=_estimate_gpu_duration)
def transcribe_audio(
    audio: Any,
    instruments: list[str] | None,
    use_sampling: bool,
    temperature: float,
    beam_size: int,
) -> Iterator[tuple[str | None, list[str] | None, str, dict[str, Any]]]:
    """Stream chunk progress, then return combined and isolated MIDI files."""

    path = _audio_path(audio)
    if not path:
        empty_tracks: list[dict[str, Any]] = []
        message = "Drop an audio file before starting transcription."
        yield (
            None,
            None,
            _payload(
                state="error",
                status=message,
                progress=0,
                audio_name="",
                elapsed=0,
                duration=0,
                tracks=empty_tracks,
            ),
            _manifest(
                state="error",
                audio_name="",
                elapsed=0,
                duration=0,
                tracks=empty_tracks,
                message=message,
            ),
        )
        return

    audio_name = Path(path).name
    duration = _audio_duration(path) or 0.0
    if duration > MAX_AUDIO_SECONDS:
        tracks = _track_payloads({}, instruments)
        message = f"This demo accepts audio up to {MAX_AUDIO_SECONDS:.0f} seconds."
        yield (
            None,
            None,
            _payload(
                state="error",
                status=message,
                progress=0,
                audio_name=audio_name,
                elapsed=0,
                duration=duration,
                tracks=tracks,
            ),
            _manifest(
                state="error",
                audio_name=audio_name,
                elapsed=0,
                duration=duration,
                tracks=tracks,
                message=message,
            ),
        )
        return

    requested = list(dict.fromkeys(instruments or []))
    beam = max(1, min(4, int(beam_size or 1)))
    started = time.perf_counter()
    events: list[NoteStartEvent | NoteEndEvent | ProgressEvent] = []
    last_completed = -1

    try:
        stream = MODEL.transcribe(
            path,
            instruments=requested or None,
            use_sampling=bool(use_sampling),
            temperature=float(temperature),
            beam_size=beam,
            batch_size=1,
        )
        for event in stream:
            events.append(event)
            if not isinstance(event, ProgressEvent):
                continue
            if event.completed == last_completed:
                continue
            last_completed = event.completed
            elapsed = time.perf_counter() - started
            fraction = event.completed / max(1, event.total)
            notes = _notes_from_events(events)
            tracks = _track_payloads(notes, requested)
            status = (
                "GPU ready · preparing the first 5-second window"
                if event.completed == 0
                else f"Transcribed {event.completed} of {event.total} audio windows"
            )
            yield (
                None,
                None,
                _payload(
                    state="transcribing",
                    status=status,
                    progress=fraction,
                    audio_name=audio_name,
                    elapsed=elapsed,
                    duration=duration,
                    tracks=tracks,
                ),
                _manifest(
                    state="running",
                    audio_name=audio_name,
                    elapsed=elapsed,
                    duration=duration,
                    tracks=tracks,
                ),
            )

        elapsed = time.perf_counter() - started
        notes = _notes_from_events(events)
        # Instrument hints guide decoding, but a hinted instrument with no
        # decoded notes is not a real output track. Keep hinted placeholders
        # during progress, then export only non-empty detections.
        preview_tracks = _track_payloads(notes)
        track_names = [track["key"] for track in preview_tracks]
        full_path, track_paths, full_bytes, midi_by_instrument = _write_outputs(
            events, track_names
        )
        tracks = _track_payloads(notes, midi_by_instrument=midi_by_instrument)
        print(
            "[MuScriptor Studio] "
            f"{audio_name}: {duration:.2f}s audio, {len(tracks)} tracks, "
            f"{sum(track['note_count'] for track in tracks)} notes, "
            f"{elapsed:.2f}s inference.",
            flush=True,
        )
        yield (
            str(full_path),
            [str(track_paths[track["key"]]) for track in tracks],
            _payload(
                state="complete",
                status=f"Transcription complete in {elapsed:.1f} seconds",
                progress=1,
                audio_name=audio_name,
                elapsed=elapsed,
                duration=duration,
                tracks=tracks,
                full_midi=full_bytes,
            ),
            _manifest(
                state="complete",
                audio_name=audio_name,
                elapsed=elapsed,
                duration=duration,
                tracks=tracks,
                full_path=full_path,
                track_paths=track_paths,
            ),
        )
    except Exception as exc:
        elapsed = time.perf_counter() - started
        notes = _notes_from_events(events)
        tracks = _track_payloads(notes, requested)
        message = f"Transcription stopped: {type(exc).__name__}: {exc}"
        print(f"[MuScriptor Studio] {message}", file=sys.stderr, flush=True)
        yield (
            None,
            None,
            _payload(
                state="error",
                status=message,
                progress=0,
                audio_name=audio_name,
                elapsed=elapsed,
                duration=duration,
                tracks=tracks,
            ),
            _manifest(
                state="error",
                audio_name=audio_name,
                elapsed=elapsed,
                duration=duration,
                tracks=tracks,
                message=message,
            ),
        )


def _read_asset(name: str) -> str:
    return (ASSET_ROOT / name).read_text(encoding="utf-8")


APP_CSS = _read_asset("app.css")
VISUALIZER_TEMPLATE = _read_asset("visualizer.html")
VISUALIZER_JS = _read_asset("visualizer.js")

INSTRUMENT_CHOICES = [
    (_display_name(name), name)
    for name in sorted(MT3_FULL_PLUS_GROUP_NAMES, key=MT3_FULL_PLUS_GROUP_NAMES.get)
]

INITIAL_VISUALIZER = _payload(
    state="idle",
    status="Ready for audio",
    progress=0,
    audio_name="",
    elapsed=0,
    duration=0,
    tracks=[],
)

HEADER = """
<header class="masthead-inner">
  <a class="brand" href="https://github.com/muscriptor/muscriptor" target="_blank" rel="noreferrer">
    <span class="brand-mark" aria-hidden="true"><i></i><i></i><i></i><i></i></span>
    <span>MuScriptor</span>
  </a>
  <div class="header-meta">
    <span id="model-badge"><b></b> Medium · 307M</span>
    <span class="zerogpu-badge">ZeroGPU</span>
  </div>
</header>
"""

INTRO = """
<div class="intro-copy">
  <span class="eyebrow">Multitrack transcription</span>
  <h1>Turn a recording into<br><em>editable music.</em></h1>
  <p>MuScriptor separates notes by instrument and shapes them into MIDI—one 5-second window at a time.</p>
</div>
"""

UPLOAD_HEADING = """
<div class="section-heading">
  <span class="step-number">01</span>
  <div><h2>Choose audio</h2><p>WAV, MP3, FLAC or OGG · up to 60 seconds</p></div>
</div>
"""

RESULT_HEADING = """
<div class="section-heading result-heading">
  <span class="step-number">02</span>
  <div><h2>Transcription</h2><p>Live notes, isolated tracks and MIDI export</p></div>
</div>
"""

EXAMPLES_HEADING = """
<div class="examples-heading">
  <div>
    <span class="examples-kicker">Try a session</span>
    <h2>Start with an example</h2>
    <p>Real performances spanning symphony orchestra, modern jazz and layered pop-rock.</p>
  </div>
  <span class="examples-hint">One click · 10–14 seconds</span>
</div>
"""

EXAMPLE_ROWS = [
    [
        str(EXAMPLE_ROOT / "mozart-magic-flute-overture.mp3"),
        ["string_ensemble", "flutes", "oboe", "clarinet", "bassoon"],
        False,
        1.0,
        1,
    ],
    [
        str(EXAMPLE_ROOT / "harry-mitchell-jazz-quartet.mp3"),
        ["acoustic_piano", "tenor_sax", "acoustic_bass", "drums"],
        False,
        1.0,
        1,
    ],
    [
        str(EXAMPLE_ROOT / "double-tracked-pop-rock.mp3"),
        [
            "distorted_electric_guitar",
            "acoustic_guitar",
            "electric_bass",
            "drums",
        ],
        False,
        1.0,
        1,
    ],
]

EXAMPLE_LABELS = [
    "Mozart — Magic Flute Overture · orchestra · 12s",
    "Harry Mitchell Quartet · live jazz · 14s",
    "Layered Pop-Rock · real guitars & bass · 10s",
]

FOOTNOTE = """
<div class="footnote">
  <span>Your audio is processed ephemerally and is not retained.</span>
  <span>Kyutai × Mirelo · CC BY-NC 4.0 weights</span>
</div>
"""

with gr.Blocks(title="MuScriptor — Audio to MIDI") as demo:
    with gr.Column(elem_id="app-shell"):
        gr.HTML(HEADER, elem_id="masthead")
        gr.HTML(INTRO, elem_id="intro")

        with gr.Row(elem_id="studio-grid", equal_height=False):
            with gr.Column(scale=4, min_width=320, elem_id="upload-card"):
                gr.HTML(UPLOAD_HEADING)
                audio_input = gr.Audio(
                    label="Drop audio here",
                    sources=["upload"],
                    type="filepath",
                    elem_id="audio-input",
                )
                with gr.Accordion(
                    "Guide the transcription",
                    open=False,
                    elem_id="advanced-settings",
                ):
                    instrument_input = gr.CheckboxGroup(
                        choices=INSTRUMENT_CHOICES,
                        value=[],
                        label="Known instruments",
                        info="Optional. Leave blank for automatic discovery.",
                    )
                    use_sampling_input = gr.Checkbox(
                        label="Creative decoding",
                        value=False,
                        info="Uses stochastic sampling instead of deterministic decoding.",
                    )
                    temperature_input = gr.Slider(
                        minimum=0.2,
                        maximum=1.4,
                        value=1.0,
                        step=0.1,
                        label="Temperature",
                    )
                    beam_size_input = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Beam width",
                        info=(
                            "1 is faster. Widths 2–4 use slower deterministic "
                            "beam search; sampling and temperature apply only at 1."
                        ),
                    )
                transcribe_button = gr.Button(
                    "Transcribe audio",
                    variant="primary",
                    elem_id="transcribe-button",
                )
                gr.HTML(
                    '<p class="cold-note"><span></span> Cold starts can take a moment while ZeroGPU allocates the model.</p>'
                )

            with gr.Column(scale=7, min_width=420, elem_id="result-panel"):
                gr.HTML(RESULT_HEADING)
                visualizer_output = gr.HTML(
                    value=INITIAL_VISUALIZER,
                    html_template=VISUALIZER_TEMPLATE,
                    js_on_load=VISUALIZER_JS,
                    elem_id="visualizer-host",
                )
                full_midi_output = gr.File(
                    label="Full arrangement MIDI",
                    elem_id="download-output",
                )
                track_files_output = gr.File(
                    label="Instrument MIDI files",
                    file_count="multiple",
                    visible=False,
                )
                manifest_output = gr.JSON(label="Transcription manifest", visible=False)

        with gr.Column(elem_id="examples-section"):
            gr.HTML(EXAMPLES_HEADING)
            gr.Examples(
                examples=EXAMPLE_ROWS,
                inputs=[
                    audio_input,
                    instrument_input,
                    use_sampling_input,
                    temperature_input,
                    beam_size_input,
                ],
                outputs=[
                    full_midi_output,
                    track_files_output,
                    visualizer_output,
                    manifest_output,
                ],
                fn=transcribe_audio,
                cache_examples=True,
                cache_mode="lazy",
                run_on_click=True,
                preload=False,
                examples_per_page=3,
                example_labels=EXAMPLE_LABELS,
                label=None,
                elem_id="examples-gallery-v2",
                api_visibility="private",
                api_name="load_example",
            )

        gr.HTML(FOOTNOTE)

    transcribe_button.click(
        fn=transcribe_audio,
        inputs=[
            audio_input,
            instrument_input,
            use_sampling_input,
            temperature_input,
            beam_size_input,
        ],
        outputs=[
            full_midi_output,
            track_files_output,
            visualizer_output,
            manifest_output,
        ],
        api_name="transcribe",
        concurrency_limit=1,
        concurrency_id="muscriptor-medium",
        show_progress="full",
    )

demo.queue(default_concurrency_limit=1, max_size=8)


if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Base(
            primary_hue="indigo",
            neutral_hue="slate",
        ),
        css=APP_CSS,
        js="""() => {
          document.documentElement.classList.add('dark');
          document.body.classList.add('dark');
          document.body.style.background = '#090b0f';
          const host = document.querySelector('gradio-app');
          if (host) host.style.background = '#090b0f';
          try { localStorage.setItem('theme', 'dark'); } catch (_) {}
        }""",
        show_error=True,
    )