/**
 * 自动绕过 shadowRoot 的 querySelector
 * @param {string} selector - 要查询的 CSS 选择器
 * @returns {Element|null} - 匹配的元素或 null 如果未找到
 */
function deepQuerySelector(selector) {
  /**
   * 在指定的根元素或文档对象下深度查询元素
   * @param {Element|Document} root - 要开始搜索的根元素或文档对象
   * @param {string} selector - 要查询的 CSS 选择器
   * @returns {Element|null} - 匹配的元素或 null 如果未找到
   */
  function deepSearch(root, selector) {
    let element = root.querySelector(selector);
    if (element) {
      return element;
    }

    const shadowHosts = root.querySelectorAll('*');
    for (let i = 0; i < shadowHosts.length; i++) {
      const host = shadowHosts[i];
      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;

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

let uiUpdateCallbacks = [];
let msgReceiveCallbacks = [];

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

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

function runCallback(x, m) {
  try {
    x(m);
  } catch (e) {
    (console.error || console.log).call(console, e.message, e);
  }
}

function executeCallbacks(queue, m) {
  queue.forEach(function(x) {
    runCallback(x, m);
  });
}

document.addEventListener("DOMContentLoaded", function() {
  const mutationObserver = new MutationObserver(function(m) {
    executeCallbacks(uiUpdateCallbacks, m);
  });
  mutationObserver.observe(gradioApp(), { childList: true, subtree: true });
});

function HSVtoRGB(h, s, v) {
  let r, g, b, i, f, p, q, t;
  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)
  };
}

class MidiVisualizer extends HTMLElement {
  constructor() {
    super();
    this.midiEvents = [];
    this.activeNotes = [];
    this.midiTimes = [];
    this.wrapper = null;
    this.svg = null;
    this.timeLine = null;
    this.config = {
      noteHeight: 4,
      beatWidth: 32
    };
    this.timePreBeat = 16;
    this.svgWidth = 0;
    this.t1 = 0;
    this.totalTimeMs = 0;
    this.playTime = 0;
    this.playTimeMs = 0;
    this.colorMap = new Map();
    this.playing = false;
    this.timer = null;
    this.init();
  }

  init() {
    this.innerHTML = '';
    const shadow = this.attachShadow({ mode: 'open' });
    const style = document.createElement('style');
    const wrapper = document.createElement('div');
    style.textContent = ".note.active {stroke: black;stroke-width: 0.75;stroke-opacity: 0.75;}";
    wrapper.style.overflowX = "scroll";
    const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg');
    svg.style.height = `${this.config.noteHeight * 128}px`;
    svg.style.width = `${this.svgWidth}px`;
    const timeLine = document.createElementNS('http://www.w3.org/2000/svg', 'line');
    timeLine.style.stroke = "green";
    timeLine.style.strokeWidth = 2;
    shadow.appendChild(style);
    shadow.appendChild(wrapper);
    wrapper.appendChild(svg);
    svg.appendChild(timeLine);
    this.wrapper = wrapper;
    this.svg = svg;
    this.timeLine = timeLine;
    this.setPlayTime(0);
  }

  clearMidiEvents() {
    this.pause();
    this.midiEvents = [];
    this.activeNotes = [];
    this.midiTimes = [];
    this.t1 = 0;
    this.colorMap.clear();
    this.setPlayTime(0);
    this.totalTimeMs = 0;
    this.playTimeMs = 0;
    this.svgWidth = 0;
    this.svg.innerHTML = '';
    this.svg.style.width = `${this.svgWidth}px`;
    this.svg.appendChild(this.timeLine);
  }

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

  getColor(track, channel) {
    let key = `${track},${channel}`;
    let color = this.colorMap.get(key);
    if (color) {
      return color;
    }
    color = HSVtoRGB(Math.random(), Math.random() * 0.5 + 0.5, 1);
    this.colorMap.set(key, color);
    return color;
  }

  drawNote(x, y, w, h, fill) {
    if (!this.svg) {
      return null;
    }
    const rect = document.createElementNS('http://www.w3.org/2000/svg', 'rect');
    rect.classList.add('note');
    rect.setAttribute('fill', fill);
    rect.setAttribute('x', `${Math.round(x)}`);
    rect.setAttribute('y', `${Math.round(y)}`);
    rect.setAttribute('width', `${Math.round(w)}`);
    rect.setAttribute('height', `${Math.round(h)}`);
    this.svg.appendChild(rect);
    return rect;
  }

  finishAppendMidiEvent() {
    this.pause();
    let midiEvents = this.midiEvents.sort((a, b) => a[1] - b[1]);
    let tempo = (60 / 120) * 10 ** 3;
    let ms = 0;
    for (let i = 0; i < midiEvents.length; i++) {
      let midiEvent = midiEvents[i];
      if (i === 0) {
        ms = 0;
      } else {
        ms += (midiEvent[1] - midiEvents[i - 1][1]) * tempo / this.timePreBeat;
      }
      midiEvent.push(ms);
      this.midiTimes.push(ms);
    }
    this.totalTimeMs = ms;
  }

  getDuration() {
    return this.totalTimeMs;
  }

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

  setPlayTimeMs(ms) {
    if (ms < 0 || ms > this.totalTimeMs) {
      return;
    }
    let i = this.midiTimes.findIndex((t) => t > ms);
    if (i === -1) {
      i = this.midiTimes.length;
    }
    this.playTimeMs = ms;
    this.setPlayTime(this.midiEvents[i - 1][1]);
    let notes = this.activeNotes.filter((note) => {
      note[7].classList.remove('active');
      if (note[1] + note[3] * this.timePreBeat > this.playTime) {
        return true;
      }
      return false;
    });
    this.activeNotes = notes;
    let actives = this.midiEvents.filter((event) => event[1] <= this.playTime && this.playTime < event[1] + event[3] * this.timePreBeat);
    actives.forEach((event) => {
      event[7].classList.add('active');
    });
    this.activeNotes.push(...actives);
  }

  isPlaying() {
    return this.playing;
  }

  play() {
    if (this.playing) {
      return;
    }
    this.playing = true;
    let t0 = this.playTimeMs;
    let start = performance.now();
    const step = (now) => {
      this.timer = requestAnimationFrame(step);
      this.setPlayTimeMs(t0 + now - start);
      if (t0 + now - start >= this.totalTimeMs) {
        this.pause();
      }
    };
    this.timer = requestAnimationFrame(step);
  }

  pause() {
    if (this.timer) {
      cancelAnimationFrame(this.timer);
      this.timer = null;
    }
    this.playing = false;
  }
}

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

function getProgressBar() {
  const progressbar = gradioApp().deepQuerySelector('span#prompt_progress span');
  if (progressbar !== null) {
    return progressbar;
  }
  return { innerHTML: '' };
}

function getPromptResult() {
  const promptResult = gradioApp().deepQuerySelector('div#prompt_result');
  if (promptResult !== null) {
    return promptResult;
  }
  return { innerHTML: '' };
}

function updateProgressBarHTML(html) {
  const bar = getProgressBar();
  if (bar.innerHTML !== html) {
    bar.innerHTML = html;
  }
}

function updateProgressBar(text, value, time) {
  let percent = !text.includes('%') ? `${Math.round(value * 100)}%` : '';
  updateProgressBarHTML(`${text} ${percent}`);
  const result = getPromptResult();
  result.innerHTML += `<br>${text} ${percent}`;
}

onUiUpdate(function() {
  const bar = getProgressBar();
  if (bar.innerHTML === '') {
    return;
  }
  let str = bar.innerHTML;
  if (str.includes('Rendering') || str.includes('Generating')) {
    updateProgressBar(str, 0);
  } else {
    const match = str.match(/([0-9.]+)%/);
    if (match) {
      const value = parseFloat(match[1]) / 100;
      updateProgressBar(str, value);
    }
  }
});

function runWhenReady(fn) {
  const gradioRoot = gradioApp();
  if (gradioRoot !== null && gradioRoot !== undefined) {
    fn();
  } else {
    const observer = new MutationObserver(function(mutations) {
      if (gradioRoot !== null && gradioRoot !== undefined) {
        observer.disconnect();
        fn();
      }
    });
    observer.observe(document, { childList: true, subtree: true });
  }
}

runWhenReady(function() {
  const progressbar = getProgressBar();
  if (progressbar !== null) {
    progressbar.innerHTML = '';
  }
});