"""
🎧 Magenta DJ — real-time AI music station built on Google Magenta RT2.

Blends genres live in real time, morphs smoothly between styles, records
to WAV. Built for the Build Small Hackathon 2026.

Model: Google Magenta RT2 Small — ~600M params, 48 kHz stereo.
Transport: gapless Web Audio scheduler (zero seams, no <audio> element gaps).
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")

import json, time, uuid, math, base64, tempfile, warnings
import numpy as np
import soundfile as sf
import torch
import spaces
import gradio as gr
from huggingface_hub import hf_hub_download

warnings.filterwarnings("ignore")
from magenta_rt import paths

# ── Model loading ─────────────────────────────────────────────────────────────
MAGENTA_HOME = "/data" if os.path.isdir("/data") else "/tmp/magenta"
HOME = os.path.join(MAGENTA_HOME, "magenta-rt-v2")
os.makedirs(HOME, exist_ok=True)

hf_hub_download("google/magenta-realtime-2",
                "checkpoints/mrt2_small.safetensors", local_dir=HOME)
paths.set_magenta_home(HOME)

from magenta_rt.torch import MagentaRT2
from magenta_rt.torch.musiccoca import MusicCoCa

print("[boot] loading Magenta RT2 Small…", flush=True)
style_model = MusicCoCa(device="cpu")
model_main  = MagentaRT2(size="mrt2_small", device="cuda",
                          dtype=torch.bfloat16, style_model=style_model)
try:
    model_main.load_compiled(repo_id="magenta-torch/magenta-rt-aoti-small")
    print("[boot] AOTI fast path active.", flush=True)
except Exception as e:
    print("[boot] AOTI unavailable → eager:", repr(e)[:80], flush=True)

SR            = 48000
EMB_DIM       = style_model.embedding_dim
FRAME_SAMPLES = 1920
SESSION_DIR   = "/tmp/mrt_dj"
os.makedirs(SESSION_DIR, exist_ok=True)

CHUNK     = 32    # frames per SSE event  (~1.28 s)
LEAD      = 6.0   # generator lead time   (s)
MAX_SLEEP = 0.30  # max throttle sleep    (s)

MAX_VARS = 8
DEFAULT_VARS = [
    ("upbeat pop",          1.4),
    ("chill lo-fi hip hop", 0.5),
    ("deep house",          0.0),
    ("jazz café piano",     0.0),
    ("80s synthwave",       0.0),
    ("cinematic strings",   0.0),
]
PRESETS = {
    "🎵 Pop":   {"upbeat pop": 2.0},
    "🏠 House": {"deep house": 2.0},
    "☕ Lo-fi": {"chill lo-fi hip hop": 2.0},
    "🎷 Jazz":  {"jazz café piano": 2.0},
}

# ── Session slot helpers ───────────────────────────────────────────────────────
def _slot(sid):
    return os.path.join(SESSION_DIR, f"{os.path.basename(sid)}.json")

def write_slot(sid, d):
    if not sid: return
    tmp = _slot(sid) + "." + uuid.uuid4().hex
    with open(tmp, "w") as f: json.dump(d, f)
    os.replace(tmp, _slot(sid))

def read_slot(sid):
    try:
        with open(_slot(sid)) as f: return json.load(f)
    except Exception: return None

def build_ctrl(glide, cfg, drums, temp, *tw):
    prompts, weights = [], []
    for i in range(MAX_VARS):
        t = tw[2*i]     if 2*i     < len(tw) else ""
        w = tw[2*i + 1] if 2*i + 1 < len(tw) else 0.0
        prompts.append((t or "").strip())
        try:    weights.append(float(w))
        except: weights.append(0.0)
    return {
        "prompts": prompts, "weights": weights,
        "transition": float(glide), "cfg": float(cfg),
        "cfg_drums":  float(drums),  "temperature": float(temp), "top_k": 50,
    }

def push_ctrl(sid, glide, cfg, drums, temp, *tw):
    prev = read_slot(sid) or {}
    c = build_ctrl(glide, cfg, drums, temp, *tw)
    c["playing"] = prev.get("playing", True)
    write_slot(sid, c)

# ── AI generator ──────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def dj_stream(sid, glide, cfg, drums, temp, *tw):
    from magenta_rt.torch.system import make_sampler, discretize_cfg, _float_to_int16
    if not sid: return
    c0 = build_ctrl(glide, cfg, drums, temp, *tw)
    c0["playing"] = True
    write_slot(sid, c0)

    if style_model.device != "cuda":
        style_model.to("cuda")
    dev, dt = "cuda", torch.bfloat16
    gen = torch.Generator(device=dev).manual_seed(0)
    txt_cache = {}

    def embed(lbl):
        e = txt_cache.get(lbl)
        if e is None:
            e = style_model.embed(lbl).float()
            txt_cache[lbl] = e
        return e

    model        = model_main
    dstate       = model.model.decoder.init_streaming_f(1, dev, dt)
    decode_state = model.init_decode_state()
    source = source_target = cur_target_tokens = cur_cfg = cur_emb = None
    emitted = seq = 0
    t0 = time.time()
    pcm_path = os.path.join(SESSION_DIR, f"{os.path.basename(sid)}.pcm")
    pcm = open(pcm_path, "wb")

    try:
        while time.time() - t0 < 116.0:
            c = read_slot(sid) or c0
            if not c.get("playing", True): break

            toks = []
            for _ in range(CHUNK):
                c       = read_slot(sid) or c
                prompts = c.get("prompts") or ["instrumental music"]
                weights = c.get("weights") or [1.0] * len(prompts)
                g       = float(c.get("transition", 4.0))

                # ── target embedding (weighted blend of active labels) ──────
                target = torch.zeros(EMB_DIM, device=dev, dtype=torch.float32)
                tot    = float(sum(w for w in weights if w > 0)) or 1.0
                any_on = False
                for i, p in enumerate(prompts):
                    w   = float(weights[i]) if i < len(weights) else 0.0
                    lbl = (p or "").strip()
                    if w <= 0 or not lbl: continue
                    target += (w / tot) * embed(lbl)
                    any_on  = True
                if not any_on:
                    target = embed("ambient instrumental music")

                # ── cur_emb glides toward target ───────────────────────────
                dtf = FRAME_SAMPLES / SR
                if cur_emb is None or g <= 0.05:
                    cur_emb = target
                else:
                    a       = 1.0 - math.exp(-dtf / g)
                    cur_emb = cur_emb + a * (target - cur_emb)

                # ── source_target: encode TARGET tokens (not cur_emb) ──────
                # Recomputed only when user actually changes controls.
                # source then glides toward source_target at the same rate,
                # giving seamless cross-fades between genres.
                target_tokens = style_model.tokenize(target).tolist()
                cfgsig = (round(float(c.get("cfg", 1.6)), 3),
                          round(float(c.get("cfg_drums", 4.0)), 3))
                if source_target is None or target_tokens != cur_target_tokens or cfgsig != cur_cfg:
                    cur_target_tokens = target_tokens
                    cur_cfg = cfgsig
                    cfgs = [discretize_cfg(c.get("cfg", 1.6),       0.2, 40),
                            discretize_cfg(0.0,                      0.2, 40),
                            discretize_cfg(c.get("cfg_drums", 4.0), 1.0,  8)]
                    cond = model._conditioning(
                        (target_tokens + [-1]*model.num_musiccoca)[:model.num_musiccoca],
                        [-1]*128, [-1], cfgs)
                    source_target = model.model.encode(cond).to(dt)

                if source is None or g <= 0.05:
                    source = source_target
                else:
                    a_s    = 1.0 - math.exp(-dtf / g)
                    source = (source.float() + a_s*(source_target.float()-source.float())).to(dt)

                sampler = make_sampler(c.get("temperature", 1.1), c.get("top_k", 50), gen)
                toks.append(model.model.decoder.step_f(
                    dstate, source, sampler=sampler,
                    temporal_step=model._temporal_step,
                    depth_step=model._depth_step))

            new_codes = torch.cat(toks, dim=1)
            audio     = model.decode_stream(new_codes, decode_state)
            emitted  += audio.shape[1]

            if audio.shape[1] > 0:
                i16 = _float_to_int16(audio[0].float().cpu().numpy()).astype(np.int16)
                raw = i16.tobytes()
                pcm.write(raw); pcm.flush()
                seq += 1
                yield f"{seq}|" + base64.b64encode(raw).decode("ascii")

            ahead = (emitted / SR) - (time.time() - t0)
            if ahead > LEAD:
                time.sleep(min(ahead - LEAD, MAX_SLEEP))
    finally:
        try: pcm.close()
        except: pass

def stop_stream(sid):
    c = read_slot(sid) or {}
    c["playing"] = False
    write_slot(sid, c)
    return ""   # JS handles the status label

def render_recording(sid):
    p = os.path.join(SESSION_DIR, f"{os.path.basename(sid)}.pcm")
    try:    raw = open(p, "rb").read()
    except: raw = b""
    if len(raw) < SR:
        return None, "Nothing recorded yet — press ▶ Play first."
    wav = np.frombuffer(raw, dtype=np.int16).reshape(-1, 2)
    out = os.path.join(tempfile.gettempdir(),
                       f"magenta_dj_{os.path.basename(sid)[:8]}.wav")
    sf.write(out, wav, SR, subtype="PCM_16")
    return out, f"✅ {wav.shape[0]/SR:.1f}s recorded. Download below."

def clear_recording(sid):
    try: os.remove(os.path.join(SESSION_DIR, f"{os.path.basename(sid)}.pcm"))
    except: pass
    return None, "Cleared."

def apply_preset(name, *texts):
    mp = PRESETS[name]
    return [float(next((v for k, v in mp.items() if k.lower()==(t or "").strip().lower()), 0.0))
            for t in texts]

# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
:root {
  --brand: #7c5cff;
  --teal:  #19d3da;
  --dark:  #080910;
  --card:  rgba(255,255,255,.055);
  --edge:  rgba(255,255,255,.11);
  --muted: rgba(255,255,255,.42);
}
*, *::before, *::after { box-sizing: border-box; }
.gradio-container { max-width: 640px !important; margin: 0 auto !important; padding: 0 12px 40px !important; }
.gradio-container, body {
  background: radial-gradient(ellipse 1000px 500px at 50% -60px, #180d3a 0%, #080910 62%) fixed !important;
}
footer, .svelte-1ipelgc { display: none !important; }

/* ── hero ── */
#dj-hero { text-align: center; padding: 22px 4px 6px; }
#dj-hero h1 {
  margin: 0 0 6px;
  font-size: clamp(1.7rem, 7vw, 2.3rem);
  font-weight: 900;
  letter-spacing: -.03em;
  background: linear-gradient(100deg, var(--brand) 0%, var(--teal) 100%);
  -webkit-background-clip: text; background-clip: text; color: transparent;
}
#dj-hero p { margin: 0; font-size: .8rem; color: var(--muted); }
.model-badge {
  display: inline-flex; align-items: center; gap: 7px;
  margin-top: 10px;
  padding: 5px 14px;
  background: var(--card);
  border: 1px solid var(--edge);
  border-radius: 999px;
  font-size: .72rem;
  color: var(--muted);
  letter-spacing: .02em;
}
.model-badge b { color: rgba(255,255,255,.72); font-weight: 600; }

/* ── status bar ── */
#dj-status-bar {
  display: flex; align-items: center; gap: 10px;
  padding: 9px 14px;
  margin: 10px 0 6px;
  background: var(--card);
  border: 1px solid var(--edge);
  border-radius: 10px;
  font-size: .82rem;
  color: rgba(255,255,255,.65);
  min-height: 38px;
}
#dj-led {
  width: 9px; height: 9px; border-radius: 50%;
  background: #333; flex-shrink: 0;
  transition: background .4s, box-shadow .4s;
}
#dj-led.warming  { background: #f59e0b; }
#dj-led.live     { background: var(--teal); box-shadow: 0 0 7px var(--teal); animation: ledpulse 1.6s ease-in-out infinite; }
#dj-led.stopped  { background: var(--brand); }
@keyframes ledpulse { 0%,100%{opacity:1} 50%{opacity:.55} }

/* ── transport ── */
#transport { gap: 8px !important; margin-bottom: 0 !important; }
#transport button { height: 50px; border-radius: 12px; font-weight: 800 !important; font-size: .97rem !important; }
#play-btn {
  background: linear-gradient(130deg, var(--brand) 0%, var(--teal) 100%) !important;
  color: #fff !important; border: none !important;
}
#stop-btn {
  background: var(--card) !important;
  border: 1px solid var(--edge) !important;
  color: rgba(255,255,255,.7) !important;
}

/* ── canvas ── */
#dj-viz-wrap { position: relative; margin: 8px 0 4px; border-radius: 12px; overflow: hidden; }
#dj_viz { display: block; width: 100%; height: 64px; background: #05060d; }

/* ── section headers ── */
.sect {
  font-size: .72rem; font-weight: 700; letter-spacing: .1em;
  text-transform: uppercase; color: var(--muted);
  margin: 20px 0 8px; display: flex; align-items: center; gap: 8px;
}
.sect::after { content:""; flex:1; height:1px; background: var(--edge); }

/* ── genre cards ── */
.varcard {
  display: flex; align-items: center; gap: 10px;
  background: var(--card);
  border: 1px solid var(--edge);
  border-radius: 12px;
  padding: 8px 12px;
  margin-bottom: 6px;
  transition: border-color .2s;
}
.varcard:focus-within { border-color: rgba(124,92,255,.45); }
.varcard textarea, .varcard input[type=text] {
  font-weight: 600 !important; font-size: .88rem !important;
  background: transparent !important; border: none !important;
}

/* ── preset pills ── */
#presetrow { gap: 8px !important; flex-wrap: wrap; }
#presetrow button {
  border-radius: 999px !important; font-weight: 700 !important;
  min-height: 38px !important; font-size: .84rem !important;
  background: var(--card) !important; border: 1px solid var(--edge) !important;
  color: rgba(255,255,255,.8) !important;
  transition: border-color .15s, background .15s !important;
}
#presetrow button:hover {
  border-color: var(--brand) !important;
  background: rgba(124,92,255,.12) !important;
}

/* ── accordions ── */
.gr-accordion { background: var(--card) !important; border: 1px solid var(--edge) !important; border-radius: 12px !important; margin-bottom: 6px !important; }
.gr-accordion summary { font-size: .85rem !important; font-weight: 600 !important; color: rgba(255,255,255,.65) !important; }

/* ── add-var btn ── */
#add-var-btn button { border-radius: 10px !important; font-size: .82rem !important; background: var(--card) !important; border: 1px solid var(--edge) !important; color: var(--muted) !important; }

/* ── record section ── */
#rec-btns { gap: 8px !important; }

/* ── footer ── */
#dj-footer {
  text-align: center; margin-top: 24px; padding-top: 14px;
  border-top: 1px solid var(--edge);
  font-size: .72rem; color: var(--muted); line-height: 1.8;
}
#dj-footer a { color: rgba(124,92,255,.8); text-decoration: none; }
#dj-footer a:hover { color: var(--teal); }
"""

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="Magenta DJ",
               theme=gr.themes.Soft(primary_hue="purple", neutral_hue="slate"),
               css=CSS) as demo:
    sid = gr.State("")

    # Hero
    gr.HTML("""
    <div id="dj-hero">
      <h1>🎧 Magenta DJ</h1>
      <p>Blend any two genres. Watch them morph in real time.</p>
      <div class="model-badge">
        <span>🧠</span>
        <span><b>Google Magenta RT2</b> · ~600M params · 48 kHz stereo · open-source</span>
      </div>
    </div>
    """)

    # Hidden PCM pipe (yields base64 PCM → Web Audio scheduler)
    pcm_pipe = gr.Textbox(visible=False, elem_id="pcm_pipe")

    # Transport
    with gr.Row(elem_id="transport"):
        play_btn = gr.Button("▶  Play", variant="primary", elem_id="play-btn", scale=4)
        stop_btn = gr.Button("■  Stop", elem_id="stop-btn", scale=1)

    # JS-managed status bar
    gr.HTML("""
    <div id="dj-status-bar">
      <span id="dj-led"></span>
      <span id="dj-status-text">Press ▶ Play to start the AI music stream</span>
    </div>
    """)

    # VU canvas
    gr.HTML("""
    <div id="dj-viz-wrap">
      <canvas id="dj_viz" width="640" height="64"></canvas>
    </div>
    """)

    # ── Genre mix ─────────────────────────────────────────────────────────────
    gr.HTML("<div class='sect'>🎚 Blend Genres</div>")
    gr.HTML("<p style='font-size:.78rem;color:rgba(255,255,255,.4);margin:-4px 0 10px'>Mix up to 8 styles. Drag the slider to set how much of each genre to add. Changes blend smoothly in real time.</p>")

    texts, sliders, cards = [], [], []
    for i in range(MAX_VARS):
        on = i < len(DEFAULT_VARS)
        t0 = DEFAULT_VARS[i][0] if on else ""
        w0 = DEFAULT_VARS[i][1] if on else 0.0
        with gr.Group(visible=on, elem_classes="varcard") as card:
            with gr.Row():
                tb = gr.Textbox(value=t0, show_label=False, container=False,
                                placeholder="genre or mood (e.g. jazz café)", max_lines=1, scale=3)
                sl = gr.Slider(0.0, 2.0, value=w0, step=0.01,
                               show_label=False, container=False, scale=2)
        texts.append(tb); sliders.append(sl); cards.append(card)

    with gr.Row(elem_id="add-var-btn"):
        add_btn = gr.Button("＋  Add a genre", size="sm")

    # ── Quick vibes ───────────────────────────────────────────────────────────
    gr.HTML("<div class='sect'>⚡ Quick Vibes</div>")
    with gr.Row(elem_id="presetrow"):
        preset_btns = [gr.Button(n, size="sm") for n in PRESETS]

    # ── Fine-tune ─────────────────────────────────────────────────────────────
    with gr.Accordion("🎛  Fine-tune the sound", open=False, elem_classes="gr-accordion"):
        glide = gr.Slider(0, 20, value=4,   step=0.5,  label="Morph speed (s) — how long genre blends take to settle")
        cfg   = gr.Slider(0.5, 7, value=1.6, step=0.1, label="Style strength — how closely the AI follows the genre labels")
        drums = gr.Slider(1,   6, value=4,   step=0.1, label="Drum intensity")
        temp  = gr.Slider(0.5, 2, value=1.1, step=0.05,label="Variation — higher = more surprising, lower = more structured")

    # ── Latency / buffer ──────────────────────────────────────────────────────
    with gr.Accordion("⚙  Playback buffer", open=False, elem_classes="gr-accordion"):
        gr.Markdown("Time the browser pre-loads before starting. Raise if you ever hear a hiccup on a slow connection.")
        buffer_ms = gr.Slider(200, 3000, value=1500, step=50, label="Startup buffer (ms)")

    # ── Record ────────────────────────────────────────────────────────────────
    with gr.Accordion("⏺  Record & download", open=False, elem_classes="gr-accordion"):
        gr.Markdown("While playing, everything is captured automatically. Hit **Render** to save it as a WAV file.")
        with gr.Row(elem_id="rec-btns"):
            render_btn = gr.Button("⏺  Render recording", variant="primary")
            clear_btn  = gr.Button("🗑  Clear")
        rec_status = gr.Markdown("")
        rec_file   = gr.File(label="Your mix (.wav)", interactive=False)

    # Footer
    gr.HTML("""
    <div id="dj-footer">
      Built on <a href="https://github.com/magenta/magenta-realtime" target="_blank">Google Magenta RT2</a>
      · open-source · Apache 2.0
      <br>
      Runs in real time on a single GPU · no audio gaps via Web Audio API scheduling
    </div>
    """)

    # ── Wire everything ───────────────────────────────────────────────────────
    ctrl = [glide, cfg, drums, temp]
    for i in range(MAX_VARS):
        ctrl += [texts[i], sliders[i]]

    demo.load(lambda: uuid.uuid4().hex, outputs=[sid])

    play_btn.click(None, None, None,
                   js="() => { window.djStart && window.djStart(); }")
    play_btn.click(dj_stream, inputs=[sid]+ctrl, outputs=[pcm_pipe])

    pcm_pipe.change(None, [pcm_pipe], None,
                    js="(b) => { window.djPlay && window.djPlay(b); }")
    buffer_ms.change(None, [buffer_ms], None,
                     js="(v) => { window.djSetBuffer && window.djSetBuffer(v); }")

    stop_btn.click(stop_stream, [sid], [gr.Markdown(visible=False)])
    stop_btn.click(None, None, None,
                   js="() => { window.djStop && window.djStop(); }")

    for comp in ctrl:
        ev = comp.release if isinstance(comp, gr.Slider) else comp.change
        ev(push_ctrl, [sid]+ctrl, None)

    vis_state = gr.State([i < len(DEFAULT_VARS) for i in range(MAX_VARS)])
    def add_var(vis):
        vis = list(vis)
        for i in range(MAX_VARS):
            if not vis[i]: vis[i] = True; break
        return [vis] + [gr.update(visible=v) for v in vis]
    add_btn.click(add_var, [vis_state], [vis_state]+cards)

    for b, name in zip(preset_btns, PRESETS):
        b.click(lambda *t, n=name: apply_preset(n, *t), inputs=texts, outputs=sliders)

    render_btn.click(render_recording, [sid], [rec_file, rec_status])
    clear_btn.click(clear_recording,   [sid], [rec_file, rec_status])


# ── Web Audio gapless scheduler + status manager ──────────────────────────────
HEAD = r"""
<script>
(function(){
  const SR = 48000;
  var ctx=null, nextTime=0, bufferSec=1.5, lastSeq=0, running=false, started=false;
  var gainNode=null, analyser=null, freq=null;
  var chunkCount=0, bufferAhead=0;

  /* ── status helper ─────────────────────────────────────────── */
  var STATES = {
    idle:     { dot:'',         text:'Press ▶ Play to start the AI music stream' },
    warming:  { dot:'warming',  text:'⏳  Warming up Magenta AI…  first sound in ~5–8 s' },
    notes:    { dot:'warming',  text:'♫  Generating first notes…' },
    live:     { dot:'live',     text:'● LIVE — drag the genre sliders to morph the music' },
    stopped:  { dot:'stopped',  text:'■  Stopped — press ▶ Play to resume' },
  };
  function setStatus(key) {
    var s = STATES[key] || STATES.idle;
    var led  = document.getElementById('dj-led');
    var text = document.getElementById('dj-status-text');
    if (led)  { led.className = s.dot; }
    if (text) { text.textContent = s.text; }
  }

  /* ── AudioContext ─────────────────────────────────────────── */
  function ensureCtx() {
    if (!ctx) {
      ctx = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: SR });
      gainNode = ctx.createGain();
      analyser = ctx.createAnalyser(); analyser.fftSize = 256;
      freq = new Uint8Array(analyser.frequencyBinCount);
      gainNode.connect(analyser); analyser.connect(ctx.destination);
      drawViz();
    }
  }

  /* ── public API ───────────────────────────────────────────── */
  window.djStart = function() {
    ensureCtx();
    if (ctx.state === 'suspended') ctx.resume();
    running=true; started=false; chunkCount=0; nextTime=0;
    setStatus('warming');
  };
  window.djStop = function() {
    running=false; started=false; bufferAhead=0;
    if (ctx) try { ctx.suspend(); } catch(e) {}
    setStatus('stopped');
  };
  window.djSetBuffer = function(ms) {
    bufferSec = Math.max(0.1, (parseFloat(ms)||1500)/1000);
  };

  /* ── scheduler ────────────────────────────────────────────── */
  window.djPlay = function(payload) {
    if (!payload || !running) return;
    ensureCtx();
    if (ctx.state === 'suspended') ctx.resume();

    var bar = payload.indexOf('|');
    if (bar < 0) return;
    var seq = parseInt(payload.slice(0,bar), 10);
    if (seq <= lastSeq) return;
    lastSeq = seq;

    var b64 = payload.slice(bar+1);
    var bin; try { bin = atob(b64); } catch(e) { return; }
    var bytes = new Uint8Array(bin.length);
    for (var i=0; i<bin.length; i++) bytes[i] = bin.charCodeAt(i);
    var i16 = new Int16Array(bytes.buffer);
    var frames = i16.length >> 1;
    if (frames < 1) return;

    var buf = ctx.createBuffer(2, frames, SR);
    var L=buf.getChannelData(0), R=buf.getChannelData(1);
    for (var i=0; i<frames; i++) { L[i]=i16[2*i]/32768; R[i]=i16[2*i+1]/32768; }

    var src = ctx.createBufferSource();
    src.buffer = buf; src.connect(gainNode);
    var now = ctx.currentTime;
    if (!started)                  { nextTime = now + bufferSec; started = true; }
    else if (nextTime < now+0.04)  { nextTime = now + 0.04; }
    src.start(nextTime);
    nextTime += frames/SR;
    bufferAhead = nextTime - ctx.currentTime;

    /* update status based on how many chunks have arrived */
    chunkCount++;
    if (chunkCount === 1) setStatus('notes');
    if (chunkCount === 2) setStatus('live');
  };

  /* ── canvas visualiser ────────────────────────────────────── */
  function drawViz() {
    var canvas = document.getElementById('dj_viz');
    if (!canvas) { requestAnimationFrame(drawViz); return; }
    var vc = canvas.getContext('2d');
    (function loop() {
      requestAnimationFrame(loop);
      var w=canvas.width, h=canvas.height;
      vc.clearRect(0,0,w,h);
      if (analyser && running) analyser.getByteFrequencyData(freq);

      var BARS=56, BAR_H=h-10, bw=w/BARS;
      for (var i=0; i<BARS; i++) {
        var v = (analyser && running)
          ? freq[i % freq.length]/255
          : 0.05 + 0.03*Math.sin(Date.now()/350+i*0.6);
        var bh = Math.max(2, v*BAR_H);
        var g = vc.createLinearGradient(0, BAR_H, 0, BAR_H-bh);
        /* colour shifts toward teal as bar gets taller */
        var ratio = Math.min(v*1.4, 1);
        var r1 = Math.round(124 + (25-124)*ratio);
        var g1 = Math.round(92  + (211-92)*ratio);
        var b1 = Math.round(255 + (218-255)*ratio);
        g.addColorStop(0, 'rgba('+r1+','+g1+','+b1+',.9)');
        g.addColorStop(1, 'rgba('+r1+','+g1+','+b1+',.25)');
        vc.fillStyle = g;
        vc.fillRect(i*bw+1, BAR_H-bh, bw-2, bh);
      }

      /* buffer health strip */
      var MAX_BUF=6, ratio2=Math.min(bufferAhead/MAX_BUF, 1);
      vc.fillStyle='rgba(0,0,0,.5)';
      vc.fillRect(0, h-10, w, 10);
      if (ratio2 > 0) {
        var hg=vc.createLinearGradient(0,0,w,0);
        hg.addColorStop(0,'#7c5cff'); hg.addColorStop(1,'#19d3da');
        vc.fillStyle=hg;
        vc.fillRect(0, h-10, w*ratio2, 10);
      }
      if (running && bufferAhead > 0.1) {
        vc.fillStyle='rgba(255,255,255,.3)';
        vc.font='8px monospace';
        vc.textAlign='right';
        vc.fillText('buffer '+bufferAhead.toFixed(1)+'s', w-4, h-1);
      }
    })();
  }

  /* grab canvas once Gradio finishes rendering */
  var t=setInterval(function(){
    if (document.getElementById('dj_viz')) { drawViz(); clearInterval(t); }
    /* also reset status dot if it isn't set yet */
    setStatus('idle');
  }, 300);
})();
</script>
"""

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=4, max_size=32).launch(
        show_error=True, ssr_mode=False, head=HEAD)