// MuScriptor front-end. Plain ES modules, no build step.
// Loads @gradio/client (for the HF Space) and @tonejs/midi + Tone (for playback) from esm.sh.

import { Client, handle_file } from "https://esm.sh/@gradio/client@1.6.0";
import { Midi } from "https://esm.sh/@tonejs/midi@2.0.28?bundle";
import * as Tone from "https://esm.sh/tone@15.1.22?bundle";

// ---- DOM helpers ----
const $ = (id) => document.getElementById(id);
const log = (msg, level = "info") => {
  const el = $("log");
  const line = `[${new Date().toISOString().slice(11, 19)}] ${level.toUpperCase()}  ${msg}\n`;
  el.textContent = (el.textContent === "(no events yet)" ? "" : el.textContent) + line;
  el.scrollTop = el.scrollHeight;
  if (level === "error") console.error(msg);
  else console.log(msg);
};
const setStatus = (id, text, cls) => {
  const el = $(id);
  el.textContent = text;
  el.className = "status " + (cls || "muted");
};

// ---- State ----
let audioFile = null;     // { file: File, url: string, source: 'upload' | 'sample' }
let audioEl = null;       // for sample playback
let midiBuffer = null;    // ArrayBuffer from Space
let midiBlobUrl = null;
let musicxmlBlobUrl = null;
let midiData = null;      // parsed @tonejs/midi structure
let toneReady = false;
let client = null;
let clientConnecting = null;

// ---- Wire up the form once the DOM is parsed ----
document.addEventListener("DOMContentLoaded", () => {
  // 1. audio file
  $("audio-file").addEventListener("change", (e) => {
    const f = e.target.files[0];
    if (f) setAudio(f, "upload");
  });

  // 2. sample buttons
  document.querySelectorAll(".sample").forEach((btn) => {
    btn.addEventListener("click", async () => {
      const url = btn.dataset.sample;
      const fname = url.split("/").pop();
      setStatus("audio-name", `loading sample ${fname}…`, "muted");
      try {
        const r = await fetch(url);
        if (!r.ok) throw new Error(`fetch ${url} -> ${r.status}`);
        const blob = await r.blob();
        const file = new File([blob], fname, { type: blob.type || "audio/wav" });
        setAudio(file, "sample", url);
      } catch (err) {
        log(`failed to fetch sample ${url}: ${err.message}`, "error");
        setStatus("audio-name", `failed to load ${fname}`, "err");
      }
    });
  });

  // 3. play audio
  $("play-audio").addEventListener("click", async () => {
    if (!audioFile) return;
    if (!audioEl) {
      // Attach to the DOM so the element appears in querySelectorAll and is
      // inspectable by tests / dev tools.
      audioEl = new Audio(audioFile.url);
      audioEl.id = "sample-audio-el";
      audioEl.preload = "auto";
      audioEl.crossOrigin = "anonymous";
      document.body.appendChild(audioEl);
    }
    try { await audioEl.play(); } catch (e) { log(`play audio failed: ${e.message}`, "error"); }
  });

  // 4. transcribe
  $("transcribe").addEventListener("click", transcribe);

  // 5. backend check
  $("check-backend").addEventListener("click", checkBackend);

  // 6. model + instrument change -> reset result card status
  ["model-size", "instruments"].forEach((id) => $(id).addEventListener("change", () => {
    if (midiBuffer || musicxmlBlobUrl) {
      setStatus("status-line", "options changed — re-transcribe to refresh results", "warn");
    }
  }));

  log("UI ready", "info");
});

function setAudio(file, source, sampleUrl) {
  if (audioFile?.url) URL.revokeObjectURL(audioFile.url);
  audioFile = {
    file,
    url: sampleUrl || URL.createObjectURL(file),
    source,
  };
  $("audio-name").textContent = `${file.name} (${(file.size / 1024).toFixed(1)} KB)`;
  $("audio-name").className = "status";
  $("play-audio").disabled = false;
  $("transcribe").disabled = false;
  audioEl = null;
  log(`audio set: ${file.name} (${source})`, "info");
}

async function checkBackend() {
  const url = $("space-url").value.trim().replace(/\/+$/, "");
  setStatus("backend-status", "checking…", "muted");
  try {
    // Gradio Spaces serve /config with the full app config.
    const r = await fetch(url + "/config", { method: "GET" });
    if (!r.ok) throw new Error(`GET /config -> ${r.status}`);
    const cfg = await r.json();
    setStatus("backend-status", `online · ${cfg.title || "Gradio app"} · v${cfg.version || "?"}`, "ok");
    log(`backend ${url} online`, "info");
  } catch (e) {
    setStatus("backend-status", `offline (${e.message})`, "err");
    log(`backend check failed for ${url}: ${e.message}`, "error");
  }
}

async function ensureClient() {
  if (client) return client;
  if (clientConnecting) return clientConnecting;
  const url = $("space-url").value.trim().replace(/\/+$/, "");
  setStatus("status-line", `connecting to ${url}…`, "muted");
  log(`connecting to ${url}`, "info");
  clientConnecting = (async () => {
    try {
      client = await Client.connect(url);
      log(`client connected`, "info");
      return client;
    } catch (e) {
      log(`client connect failed: ${e.message}`, "error");
      clientConnecting = null;
      throw e;
    } finally {
      clientConnecting = null;
    }
  })();
  return clientConnecting;
}

async function transcribe() {
  if (!audioFile) {
    setStatus("status-line", "load an audio file first", "warn");
    return;
  }
  const modelSize = $("model-size").value;
  const instrumentSelection = Array.from($("instruments").selectedOptions).map((o) => o.value);
  if (instrumentSelection.length === 0) instrumentSelection.push("all");

  const btn = $("transcribe");
  btn.disabled = true;
  setStatus("status-line", "transcribing… this can take 10-60s for medium/large", "warn");

  try {
    const c = await ensureClient();

    // /transcribe is the endpoint name (the Python function name).
    const result = await c.predict("/transcribe", [
      handle_file(audioFile.file),  // audio
      modelSize,                    // model_size
      instrumentSelection,          // instruments
    ]);

    log(`predict done · data keys: ${Object.keys(result.data).join(", ")}`, "info");

    // Output order matches the Python signature: roll, midi_dl, musicxml_dl, summary
    const [roll, midiDl, musicxmlDl, summary] = result.data;

    // 1. piano roll image
    if (roll && roll.url) {
      $("roll").src = roll.url;
      $("roll").alt = `Piano roll for ${audioFile.file.name}`;
    } else if (roll && roll.path) {
      $("roll").src = `${$("space-url").value.trim().replace(/\/+$/, "")}/file=${roll.path}`;
    }

    // 2. MIDI file
    if (midiDl) {
      const midiUrl = midiDl.url || `${$("space-url").value.trim().replace(/\/+$/, "")}/file=${midiDl.path}`;
      const mr = await fetch(midiUrl);
      midiBuffer = await mr.arrayBuffer();
      midiData = new Midi(midiBuffer);
      if (midiBlobUrl) URL.revokeObjectURL(midiBlobUrl);
      midiBlobUrl = URL.createObjectURL(new Blob([midiBuffer], { type: "audio/midi" }));
      $("midi-audio").src = midiBlobUrl;
      const dl = $("midi-download");
      dl.href = midiBlobUrl;
      dl.hidden = false;
      $("midi-info").textContent = `${midiData.tracks.length} tracks · ${midiData.duration.toFixed(1)}s · ${midiData.header?.tempos?.[0]?.bpm?.toFixed(0) || "?"} BPM`;
      $("midi-info").className = "status ok";
      log(`MIDI parsed: ${midiData.tracks.length} tracks, ${midiData.notes.length} notes`, "info");
    }

    // 3. MusicXML
    if (musicxmlDl && (musicxmlDl.url || musicxmlDl.path)) {
      const xurl = musicxmlDl.url || `${$("space-url").value.trim().replace(/\/+$/, "")}/file=${musicxmlDl.path}`;
      const xr = await fetch(xurl);
      const xb = await xr.arrayBuffer();
      if (musicxmlBlobUrl) URL.revokeObjectURL(musicxmlBlobUrl);
      musicxmlBlobUrl = URL.createObjectURL(new Blob([xb], { type: "application/vnd.recordare.musicxml+xml" }));
      const dl = $("musicxml-download");
      dl.href = musicxmlBlobUrl;
      dl.hidden = false;
      $("musicxml-info").textContent = `ready (${(xb.byteLength / 1024).toFixed(1)} KB)`;
      $("musicxml-info").className = "status ok";
    } else {
      $("musicxml-info").textContent = "(not generated for this file)";
      $("musicxml-info").className = "muted";
    }

    // 4. summary
    if (summary) {
      const txt = typeof summary === "string" ? summary : JSON.stringify(summary, null, 2);
      $("summary").textContent = txt;
    }

    setStatus("status-line", "done — press Play on the MIDI result to hear it", "ok");
  } catch (e) {
    log(`transcribe failed: ${e.message}`, "error");
    setStatus("status-line", `failed: ${e.message}`, "err");
  } finally {
    btn.disabled = false;
  }
}

// Auto-play the parsed MIDI when the user clicks the audio player's play button.
$("midi-audio")?.addEventListener("play", async () => {
  // The <audio> element just streams the .mid file. But <audio> doesn't actually play MIDI
  // reliably in all browsers. We intercept and render the parsed @tonejs/midi instead.
  if (midiData && toneReady) {
    $("midi-audio").pause();
    await playMidi();
  }
});

async function playMidi() {
  if (!midiData) return;
  if (!toneReady) {
    await Tone.start();
    toneReady = true;
  }
  Tone.Transport.stop();
  Tone.Transport.cancel();
  const synth = new Tone.PolySynth(Tone.Synth, {
    envelope: { attack: 0.005, decay: 0.1, sustain: 0.3, release: 1.2 },
  }).toDestination();
  Tone.Destination.volume.value = -8;

  const now = Tone.now() + 0.05;
  midiData.tracks.forEach((track) => {
    track.notes.forEach((note) => {
      synth.triggerAttackRelease(
        note.name,
        Math.max(0.05, note.duration),
        now + note.time,
        Math.min(0.9, note.velocity / 127)
      );
    });
  });

  const total = midiData.duration + 0.5;
  setStatus("status-line", `playing ${midiData.tracks.length}-track MIDI · ${total.toFixed(1)}s`, "ok");
  setTimeout(() => {
    setStatus("status-line", "playback finished", "muted");
    synth.dispose();
  }, total * 1000);
}

// First-user gesture hook for Tone (browsers require a user interaction).
document.addEventListener("click", async () => {
  if (toneReady) return;
  try {
    await Tone.start();
    toneReady = true;
    log("audio context unlocked", "info");
  } catch (e) { /* ignore — will retry on first play */ }
}, { once: false });