const MIDI_OUTPUT_BATCH_SIZE = 4; // Automatically set by Python backend, do not change.

// Utility to query elements across shadow DOM
function deepQuerySelector(selector) {
  function deepSearch(root, selector) {
    let element = root.querySelector(selector);
    if (element) return element;
    const shadowHosts = root.querySelectorAll('*');
    for (let host of shadowHosts) {
      if (host.shadowRoot) {
        element = deepSearch(host.shadowRoot, selector);
        if (element) return element;
      }
    }
    return null;
  }
  return deepSearch(this, selector);
}
Element.prototype.deepQuerySelector = deepQuerySelector;
Document.prototype.deepQuerySelector = deepQuerySelector;

// Get Gradio app root
function gradioApp() {
  const elems = document.getElementsByTagName('gradio-app');
  const gradioShadowRoot = elems.length ? elems[0].shadowRoot : null;
  return gradioShadowRoot || document;
}

// Callback registries
const uiUpdateCallbacks = [];
const msgReceiveCallbacks = [];

function onUiUpdate(callback) {
  uiUpdateCallbacks.push(callback);
}

function onMsgReceive(callback) {
  msgReceiveCallbacks.push(callback);
}

function runCallback(callback, data) {
  try {
    callback(data);
  } catch (e) {
    console.error(e);
  }
}

function executeCallbacks(queue, data) {
  queue.forEach((callback) => runCallback(callback, data));
}

// Observe DOM changes
document.addEventListener("DOMContentLoaded", () => {
  const observer = new MutationObserver((mutations) => {
    executeCallbacks(uiUpdateCallbacks, mutations);
  });
  observer.observe(gradioApp(), { childList: true, subtree: true });
});

// HSV to RGB conversion for note coloring
function HSVtoRGB(h, s, v) {
  let r, g, b, i = Math.floor(h * 6), f = h * 6 - i, p = v * (1 - s), q = v * (1 - f * s), t = v * (1 - (1 - f) * s);
  switch (i % 6) {
    case 0: r = v; g = t; b = p; break;
    case 1: r = q; g = v; b = p; break;
    case 2: r = p; g = v; b = t; break;
    case 3: r = p; g = q; b = v; break;
    case 4: r = t; g = p; b = v; break;
    case 5: r = v; g = p; b = q; break;
  }
  return { r: Math.round(r * 255), g: Math.round(g * 255), b: Math.round(b * 255) };
}

// Detect mobile devices
function isMobile() {
  return /iPhone|iPad|iPod|Android|Windows Phone/i.test(navigator.userAgent);
}

// MIDI Visualizer Custom Element
class MidiVisualizer extends HTMLElement {
  constructor() {
    super();
    this.midiEvents = [];
    this.activeNotes = [];
    this.midiTimes = [];
    this.trackMap = new Map();
    this.patches = Array(16).fill().map(() => [[0, 0]]);
    this.config = { noteHeight: isMobile() ? 1 : 4, beatWidth: isMobile() ? 16 : 32 };
    this.timePreBeat = 16;
    this.svgWidth = 0;
    this.t1 = 0;
    this.totalTimeMs = 0;
    this.playTime = 0;
    this.playTimeMs = 0;
    this.lastUpdateTime = 0;
    this.colorMap = new Map();
    this.playing = false;
    this.version = "v2";
    this.init();
  }

  init() {
    this.innerHTML = '';
    const shadow = this.attachShadow({ mode: 'open' });
    const style = document.createElement("style");
    style.textContent = ".note.active {stroke: black; stroke-width: 0.75; stroke-opacity: 0.75;}";
    const container = document.createElement('div');
    container.style.display = "flex";
    container.style.height = `${this.config.noteHeight * 128 + 25}px`;
    const pianoRoll = document.createElement('div');
    pianoRoll.style.overflowX = "scroll";
    pianoRoll.style.flexGrow = "1";
    const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg');
    svg.style.height = `${this.config.noteHeight * 128}px`;
    const timeLine = document.createElementNS('http://www.w3.org/2000/svg', 'line');
    timeLine.style.stroke = "green";
    timeLine.style.strokeWidth = isMobile() ? "1" : "2";

    shadow.appendChild(style);
    shadow.appendChild(container);
    container.appendChild(pianoRoll);
    pianoRoll.appendChild(svg);
    svg.appendChild(timeLine);

    this.container = container;
    this.pianoRoll = pianoRoll;
    this.svg = svg;
    this.timeLine = timeLine;
    for (let i = 0; i < 128; i++) {
      this.colorMap.set(i, HSVtoRGB(i / 128, 1, 1));
    }
    this.setPlayTime(0);
  }

  getTrack(tr, cl) {
    const id = tr * 16 + cl;
    let track = this.trackMap.get(id);
    if (!track) {
      const color = this.colorMap.get((this.trackMap.size * 53) % 128);
      track = { id, tr, cl, svg: document.createElementNS('http://www.w3.org/2000/svg', 'g'), color };
      this.svg.appendChild(track.svg);
      this.trackMap.set(id, track);
    }
    return track;
  }

  clearMidiEvents() {
    this.midiEvents = [];
    this.activeNotes = [];
    this.midiTimes = [];
    this.trackMap.clear();
    this.patches = Array(16).fill().map(() => [[0, 0]]);
    this.t1 = 0;
    this.setPlayTime(0);
    this.totalTimeMs = 0;
    this.playTimeMs = 0;
    this.lastUpdateTime = 0;
    this.svgWidth = 0;
    this.svg.innerHTML = '';
    this.svg.appendChild(this.timeLine);
    this.svg.style.width = `${this.svgWidth}px`;
  }

  appendMidiEvent(midiEvent) {
    if (Array.isArray(midiEvent) && midiEvent.length > 0) {
      this.t1 += midiEvent[1];
      const t = this.t1 * this.timePreBeat + midiEvent[2];
      midiEvent = [midiEvent[0], t].concat(midiEvent.slice(3));
      if (midiEvent[0] === "note") {
        const track = midiEvent[2];
        const channel = midiEvent[3];
        const pitch = midiEvent[4];
        const velocity = midiEvent[5];
        const duration = midiEvent[6];
        const visTrack = this.getTrack(track, channel);
        const x = (t / this.timePreBeat) * this.config.beatWidth;
        const y = (127 - pitch) * this.config.noteHeight;
        const w = (duration / this.timePreBeat) * this.config.beatWidth;
        const h = this.config.noteHeight;
        this.svgWidth = Math.ceil(Math.max(x + w, this.svgWidth));
        const opacity = Math.min(1, velocity / 127 + 0.1).toFixed(2);
        const rect = this.drawNote(visTrack, x, y, w, h, opacity);
        midiEvent.push(rect);
        this.setPlayTime(t);
        this.pianoRoll.scrollTo(this.svgWidth - this.pianoRoll.offsetWidth, 0);
      }
      this.midiEvents.push(midiEvent);
      this.svg.style.width = `${this.svgWidth}px`;
    }
  }

  drawNote(track, x, y, w, h, opacity) {
    const rect = document.createElementNS('http://www.w3.org/2000/svg', 'rect');
    rect.classList.add('note');
    const color = track.color;
    rect.setAttribute('fill', `rgba(${color.r}, ${color.g}, ${color.b}, ${opacity})`);
    rect.setAttribute('x', `${Math.round(x)}`);
    rect.setAttribute('y', `${Math.round(y)}`);
    rect.setAttribute('width', `${Math.round(w)}`);
    rect.setAttribute('height', `${Math.round(h)}`);
    track.svg.appendChild(rect);
    return rect;
  }

  finishAppendMidiEvent() {
    this.midiEvents.sort((a, b) => a[1] - b[1]);
    let tempo = (60 / 120) * 1000;
    let ms = 0;
    let lastT = 0;
    this.midiTimes.push({ ms, t: 0, tempo });
    this.midiEvents.forEach((midiEvent) => {
      const t = midiEvent[1];
      ms += ((t - lastT) / this.timePreBeat) * tempo;
      if (midiEvent[0] === "note") {
        this.totalTimeMs = Math.max(this.totalTimeMs, ms + (midiEvent[6] / this.timePreBeat) * tempo);
      } else {
        this.totalTimeMs = Math.max(this.totalTimeMs, ms);
      }
      lastT = t;
    });
  }

  setPlayTime(t) {
    this.playTime = t;
    const x = Math.round((t / this.timePreBeat) * this.config.beatWidth);
    this.timeLine.setAttribute('x1', `${x}`);
    this.timeLine.setAttribute('y1', '0');
    this.timeLine.setAttribute('x2', `${x}`);
    this.timeLine.setAttribute('y2', `${this.config.noteHeight * 128}`);
    this.pianoRoll.scrollTo(Math.max(0, x - this.pianoRoll.offsetWidth / 2), 0);

    if (this.playing && Date.now() - this.lastUpdateTime > 50) {
      const activeNotes = [];
      this.removeActiveNotes(this.activeNotes);
      this.midiEvents.forEach((midiEvent) => {
        if (midiEvent[0] === "note") {
          const time = midiEvent[1];
          const duration = midiEvent[6];
          const note = midiEvent[midiEvent.length - 1];
          if (time <= this.playTime && time + duration >= this.playTime) {
            activeNotes.push(note);
          }
        }
      });
      this.addActiveNotes(activeNotes);
      this.lastUpdateTime = Date.now();
    }
  }

  addActiveNotes(notes) {
    notes.forEach((note) => {
      this.activeNotes.push(note);
      note.classList.add('active');
    });
  }

  removeActiveNotes(notes) {
    notes.forEach((note) => {
      const idx = this.activeNotes.indexOf(note);
      if (idx > -1) this.activeNotes.splice(idx, 1);
      note.classList.remove('active');
    });
  }

  play() {
    this.playing = true;
  }

  pause() {
    this.removeActiveNotes(this.activeNotes);
    this.playing = false;
  }

  bindAudioPlayer(audio) {
    this.pause();
    audio.addEventListener("play", () => this.play());
    audio.addEventListener("pause", () => this.pause());
    audio.addEventListener("loadedmetadata", () => {
      this.totalTimeMs = audio.duration * 1000;
    });
    audio.addEventListener("timeupdate", () => {
      this.setPlayTimeMs(audio.currentTime * 1000);
    });
  }
}

customElements.define('midi-visualizer', MidiVisualizer);

// Setup visualizers and progress bar
(() => {
  const midiVisualizers = Array.from({ length: MIDI_OUTPUT_BATCH_SIZE }, () => document.createElement('midi-visualizer'));
  midiVisualizers.forEach((visualizer, idx) => {
    onUiUpdate(() => {
      const app = gradioApp();
      const container = app.querySelector(`#midi_visualizer_container_${idx}`);
      if (container && !container.contains(visualizer)) {
        container.appendChild(visualizer);
      }
      const audio = app.querySelector(`#midi_audio_${idx} audio`);
      if (audio) {
        visualizer.bindAudioPlayer(audio);
      }
    });
  });

  let hasProgressBar = false;
  let outputTabsInited = null;
  onUiUpdate(() => {
    const app = gradioApp();
    const outputTabs = app.querySelector("#output_tabs");
    if (outputTabs && outputTabsInited !== outputTabs) {
      outputTabsInited = outputTabs;
    }
  });

  function createProgressBar() {
    const parent = outputTabsInited.parentNode;
    const divProgress = document.createElement('div');
    divProgress.className = 'progressDiv';
    divProgress.style.width = "100%";
    divProgress.style.background = "#b4c0cc";
    divProgress.style.borderRadius = "8px";
    const divInner = document.createElement('div');
    divInner.className = 'progress';
    divInner.style.color = "white";
    divInner.style.background = "#0060df";
    divInner.style.textAlign = "right";
    divInner.style.fontWeight = "bold";
    divInner.style.borderRadius = "8px";
    divInner.style.height = "20px";
    divInner.style.lineHeight = "20px";
    divInner.style.paddingRight = "8px";
    divInner.style.width = "0%";
    divProgress.appendChild(divInner);
    parent.insertBefore(divProgress, outputTabsInited);
    hasProgressBar = true;
    return divInner;
  }

  function setProgressBar(progress, total) {
    if (!hasProgressBar) createProgressBar();
    const divInner = gradioApp().querySelector(".progress");
    if (divInner) {
      divInner.style.width = `${(progress / total) * 100}%`;
      divInner.textContent = `${progress}/${total}`;
    }
  }

  onMsgReceive((msgs) => {
    msgs.forEach((msg) => {
      const { name, data } = msg;
      const idx = data[0];
      const visualizer = midiVisualizers[idx];
      switch (name) {
        case "visualizer_clear":
          visualizer.clearMidiEvents();
          break;
        case "visualizer_append":
          data[1].forEach((event) => visualizer.appendMidiEvent(event));
          break;
        case "visualizer_end":
          visualizer.finishAppendMidiEvent();
          visualizer.setPlayTime(0);
          break;
        case "progress":
          setProgressBar(data[0], data[1]);
          break;
      }
    });
  });
})();