// Spectrogram Viewer
// Handles all visualization: spectrogram, partials, zoom, mouse interaction

class SpectrogramViewer {
  constructor(canvas, audioBuffer, stftCache) {
    this.canvas = canvas;
    this.ctx = canvas.getContext('2d');
    this.audioBuffer = audioBuffer;
    this.stftCache = stftCache;
    this.partials = [];
    this.frames = [];
    this.showPeaks = false;

    // Fixed time bounds
    this.t_min = 0;
    this.t_max = audioBuffer.duration;

    // View state (zoom and center)
    this.zoom_factor = 1.0; // 1.0 = full view
    this.t_center = audioBuffer.duration / 2;

    // Computed view bounds (updated by updateViewBounds)
    this.t_view_min = 0;
    this.t_view_max = audioBuffer.duration;

    // Fixed frequency bounds (log scale: freqStart must be > 0)
    this.freqStart = 20;
    this.freqEnd = 16000;

    // Tooltip
    this.tooltip = document.getElementById('tooltip');

    // Partial keep count (Infinity = all kept)
    this.keepCount = Infinity;

    // Mouse cursor overlay
    this.cursorCanvas = document.getElementById('cursorCanvas');
    this.cursorCtx = this.cursorCanvas ? this.cursorCanvas.getContext('2d') : null;
    this.mouseX = -1;

    // Playhead
    this.playheadTime = -1; // -1 = not playing

    // Spectrum viewer
    this.spectrumCanvas = document.getElementById('spectrumCanvas');
    this.spectrumCtx = this.spectrumCanvas ? this.spectrumCanvas.getContext('2d') : null;
    this.spectrumTime = 0; // Time to display spectrum for
    this.showSynthFFT = false; // Toggle: false=original, true=synth
    this.synthStftCache = null;

    // Selection and editing
    this.selectedPartial = -1;
    this.dragState = null; // {pointIndex: 0-3}
    this.onPartialSelect = null; // callback(index)
    this.onRender = null; // callback() called after each render (for synced panels)

    // Setup event handlers
    this.setupMouseHandlers();

    // Initial render
    this.updateViewBounds();
    this.render();
  }

  // --- Coordinate API ---

  // time -> canvas X
  timeToX(t) {
    return (t - this.t_view_min) / (this.t_view_max - this.t_view_min) * this.canvas.width;
  }

  // canvas X -> time
  canvasToTime(x) {
    return this.t_view_min + (x / this.canvas.width) * (this.t_view_max - this.t_view_min);
  }

  // freq -> normalized log position [0..1] within [freqStart..freqEnd]
  freqLogNorm(freq) {
    const logMin = Math.log2(this.freqStart);
    const logMax = Math.log2(this.freqEnd);
    return (Math.log2(Math.max(freq, this.freqStart)) - logMin) / (logMax - logMin);
  }

  // normalized log position [0..1] -> freq
  normToFreq(norm) {
    const logMin = Math.log2(this.freqStart);
    const logMax = Math.log2(this.freqEnd);
    return Math.pow(2, logMin + norm * (logMax - logMin));
  }

  // freq -> canvas Y (log scale)
  freqToY(freq) {
    return this.canvas.height * (1 - this.freqLogNorm(freq));
  }

  // canvas Y -> freq (log scale)
  canvasToFreq(y) {
    return this.normToFreq(1 - y / this.canvas.height);
  }

  // DB value -> normalized intensity [0..1], relative to cache maxDB over 80dB range
  normalizeDB(magDB, maxDB) {
    return Math.max(0, Math.min(1, (magDB - (maxDB - 80)) / 80));
  }

  // Partial index -> display color
  partialColor(p) {
    const colors = [
      '#f44', '#4f4', '#44f', '#ff4', '#f4f', '#4ff',
      '#fa4', '#4fa', '#a4f', '#af4', '#f4a', '#4af'
    ];
    return colors[p % colors.length];
  }

  // --- Public API ---

  setPlayheadTime(time) {
    this.playheadTime = time;
    if (time >= 0) {
      this.spectrumTime = time;
    } else if (this.mouseX >= 0) {
      this.spectrumTime = this.canvasToTime(this.mouseX);
    }
    this.render();
  }

  setPartials(partials) {
    this.partials = partials;
    this.render();
  }

  setKeepCount(n) {
    this.keepCount = n;
    this.render();
  }

  setFrames(frames) {
    this.frames = frames;
  }

  setSynthStftCache(cache) {
    this.synthStftCache = cache;
  }

  togglePeaks() {
    this.showPeaks = !this.showPeaks;
    this.render();
  }

  reset() {
    this.zoom_factor = 1.0;
    this.t_center = this.audioBuffer.duration / 2;
    this.updateViewBounds();
    this.render();
  }

  getIntensityAt(time, freq) {
    if (!this.stftCache) return -80;
    return this.stftCache.getMagnitudeDB(time, freq);
  }

  selectPartial(index) {
    this.selectedPartial = index;
    this.render();
    if (this.onPartialSelect) this.onPartialSelect(index);
  }

  // Hit-test bezier curves: returns index of nearest partial within threshold
  hitTestPartial(x, y) {
    const THRESH = 10;
    let bestIdx = -1, bestDist = THRESH;
    for (let p = 0; p < this.partials.length; ++p) {
      const curve = this.partials[p].freqCurve;
      if (!curve) continue;
      for (let i = 0; i <= 50; ++i) {
        const t = curve.t0 + (curve.t3 - curve.t0) * i / 50;
        if (t < this.t_view_min || t > this.t_view_max) continue;
        const f = evalBezier(curve, t);
        if (f < this.freqStart || f > this.freqEnd) continue;
        const px = this.timeToX(t), py = this.freqToY(f);
        const dist = Math.hypot(px - x, py - y);
        if (dist < bestDist) { bestDist = dist; bestIdx = p; }
      }
    }
    return bestIdx;
  }

  // Hit-test control points of a specific partial's freqCurve
  hitTestControlPoint(x, y, partial) {
    const curve = partial.freqCurve;
    if (!curve) return -1;
    const THRESH = 8;
    for (let i = 0; i < 4; ++i) {
      const t = curve['t' + i], v = curve['v' + i];
      if (t < this.t_view_min || t > this.t_view_max) continue;
      if (v < this.freqStart || v > this.freqEnd) continue;
      const px = this.timeToX(t), py = this.freqToY(v);
      if (Math.hypot(px - x, py - y) <= THRESH) return i;
    }
    return -1;
  }

  // --- Render ---

  render() {
    this.renderSpectrogram();
    if (this.showPeaks) this.renderPeaks();
    this.renderPartials();
    this.drawAxes();
    this.drawPlayhead();
    this.renderSpectrum();
    if (this.onRender) this.onRender();
  }

  renderPartials() {
    for (let p = 0; p < this.partials.length; ++p) {
      if (p === this.selectedPartial) continue; // draw selected last (on top)
      this._renderPartial(p, this.partials[p], false);
    }
    if (this.selectedPartial >= 0 && this.selectedPartial < this.partials.length) {
      this._renderPartial(this.selectedPartial, this.partials[this.selectedPartial], true);
    }
    this.ctx.globalAlpha = 1.0;
    this.ctx.shadowBlur = 0;
  }

  _renderPartial(p, partial, isSelected) {
    const {ctx} = this;
    const color = this.partialColor(p);
    let alpha = isSelected ? 1.0 : (p < this.keepCount ? 1.0 : 0.5);
    if (partial.muted && !isSelected) alpha = 0.15;
    ctx.globalAlpha = alpha;

    // Raw trajectory
    ctx.strokeStyle = color + '44';
    ctx.lineWidth = 1;
    ctx.beginPath();
    let started = false;
    for (let i = 0; i < partial.times.length; ++i) {
      const t = partial.times[i];
      const f = partial.freqs[i];
      if (t < this.t_view_min || t > this.t_view_max) continue;
      if (f < this.freqStart || f > this.freqEnd) continue;
      const x = this.timeToX(t);
      const y = this.freqToY(f);
      if (!started) { ctx.moveTo(x, y); started = true; } else ctx.lineTo(x, y);
    }
    if (started) ctx.stroke();

    // Bezier curve
    if (partial.freqCurve) {
      const curve = partial.freqCurve;
      if (isSelected) { ctx.shadowColor = color; ctx.shadowBlur = 8; }
      ctx.strokeStyle = color;
      ctx.lineWidth = isSelected ? 3 : 2;
      ctx.beginPath();
      started = false;
      for (let i = 0; i <= 50; ++i) {
        const t = curve.t0 + (curve.t3 - curve.t0) * i / 50;
        const freq = evalBezier(curve, t);
        if (t < this.t_view_min || t > this.t_view_max) continue;
        if (freq < this.freqStart || freq > this.freqEnd) continue;
        const x = this.timeToX(t);
        const y = this.freqToY(freq);
        if (!started) { ctx.moveTo(x, y); started = true; } else ctx.lineTo(x, y);
      }
      if (started) ctx.stroke();
      if (isSelected) ctx.shadowBlur = 0;

      ctx.fillStyle = color;
      const cpR = isSelected ? 6 : 4;
      this.drawControlPoint(curve.t0, curve.v0, cpR);
      this.drawControlPoint(curve.t1, curve.v1, cpR);
      this.drawControlPoint(curve.t2, curve.v2, cpR);
      this.drawControlPoint(curve.t3, curve.v3, cpR);
    }
  }

  renderPeaks() {
    const {ctx, frames} = this;
    if (!frames || frames.length === 0) return;

    ctx.fillStyle = '#fff';
    for (const frame of frames) {
      const t = frame.time;
      if (t < this.t_view_min || t > this.t_view_max) continue;
      const x = this.timeToX(t);
      for (const peak of frame.peaks) {
        if (peak.freq < this.freqStart || peak.freq > this.freqEnd) continue;
        ctx.fillRect(x - 1, this.freqToY(peak.freq) - 1, 3, 3);
      }
    }
  }

  drawControlPoint(t, v, radius = 4) {
    if (t < this.t_view_min || t > this.t_view_max) return;
    if (v < this.freqStart || v > this.freqEnd) return;
    const x = this.timeToX(t);
    const y = this.freqToY(v);
    this.ctx.beginPath();
    this.ctx.arc(x, y, radius, 0, 2 * Math.PI);
    this.ctx.fill();
    this.ctx.strokeStyle = '#fff';
    this.ctx.lineWidth = 1;
    this.ctx.stroke();
  }

  drawMouseCursor(x) {
    if (!this.cursorCtx) return;
    const ctx = this.cursorCtx;
    const h = this.cursorCanvas.height;
    ctx.clearRect(0, 0, this.cursorCanvas.width, h);
    if (x < 0) return;
    ctx.strokeStyle = 'rgba(255, 60, 60, 0.7)';
    ctx.lineWidth = 1;
    ctx.beginPath();
    ctx.moveTo(x, 0);
    ctx.lineTo(x, h);
    ctx.stroke();
  }

  drawPlayhead() {
    if (this.playheadTime < 0) return;
    if (this.playheadTime < this.t_view_min || this.playheadTime > this.t_view_max) return;
    const {ctx, canvas} = this;
    const x = this.timeToX(this.playheadTime);
    ctx.strokeStyle = '#f00';
    ctx.lineWidth = 2;
    ctx.beginPath();
    ctx.moveTo(x, 0);
    ctx.lineTo(x, canvas.height);
    ctx.stroke();
  }

  renderSpectrogram() {
    const {canvas, ctx, stftCache} = this;
    const width = canvas.width;
    const height = canvas.height;

    ctx.fillStyle = '#000';
    ctx.fillRect(0, 0, width, height);

    if (!stftCache) return;

    const sampleRate = this.audioBuffer.sampleRate;
    const hopSize = stftCache.hopSize;
    const fftSize = stftCache.fftSize;
    const frameDuration = hopSize / sampleRate;
    const numFrames = stftCache.getNumFrames();

    const startFrameIdx = Math.floor(this.t_view_min * sampleRate / hopSize);
    const endFrameIdx = Math.ceil(this.t_view_max * sampleRate / hopSize);

    for (let frameIdx = startFrameIdx; frameIdx < endFrameIdx; ++frameIdx) {
      if (frameIdx < 0 || frameIdx >= numFrames) continue;

      const frame = stftCache.getFrameAtIndex(frameIdx);
      if (!frame) continue;

      const squaredAmp = frame.squaredAmplitude;
      const xStart = Math.floor(this.timeToX(frame.time));
      const xEnd   = Math.ceil(this.timeToX(frame.time + frameDuration));
      const frameWidth = Math.max(1, xEnd - xStart);

      const numBins = fftSize / 2;
      const binFreqWidth = sampleRate / fftSize;

      for (let bin = 0; bin < numBins; ++bin) {
        const freq     = bin * binFreqWidth;
        const freqNext = (bin + 1) * binFreqWidth;
        if (freqNext < this.freqStart || freq > this.freqEnd) continue;

        const magDB    = 10 * Math.log10(Math.max(squaredAmp[bin], 1e-20));
        const intensity = Math.pow(this.normalizeDB(magDB, stftCache.maxDB), 2.0);

        const y0 = Math.floor(this.freqToY(freqNext));
        const y1 = Math.floor(this.freqToY(Math.max(freq, this.freqStart)));
        const binHeight = Math.max(1, y1 - y0);

        const v = Math.floor(intensity * 255);
        ctx.fillStyle = `rgba(${v},${v},${v}, 0.5)`;
        ctx.fillRect(xStart, y0, frameWidth, binHeight);
      }
    }
  }

  renderSpectrum() {
    if (!this.spectrumCtx || !this.stftCache) return;

    const useSynth = this.showSynthFFT && this.synthStftCache;
    const cache = useSynth ? this.synthStftCache : this.stftCache;

    const canvas = this.spectrumCanvas;
    const ctx = this.spectrumCtx;
    const width = canvas.width;
    const height = canvas.height;

    ctx.fillStyle = '#1e1e1e';
    ctx.fillRect(0, 0, width, height);

    const squaredAmp = cache.getSquaredAmplitude(this.spectrumTime);
    if (!squaredAmp) return;

    const numBins = cache.fftSize / 2;
    const binWidth = cache.sampleRate / cache.fftSize;

    // freq -> mini-spectrum X using same log scale as main view
    const freqToX = (freq) => this.freqLogNorm(freq) * width;

    // Draw histogram bars — one per pixel column
    for (let px = 0; px < width; ++px) {
      const fStart = this.normToFreq(px / width);
      const fEnd   = this.normToFreq((px + 1) / width);

      const bStart = Math.max(0, Math.floor(fStart / binWidth));
      const bEnd   = Math.min(numBins - 1, Math.ceil(fEnd / binWidth));

      let sum = 0, count = 0;
      for (let b = bStart; b <= bEnd; ++b) { sum += squaredAmp[b]; ++count; }
      if (count === 0) continue;

      const magDB = 10 * Math.log10(Math.max(sum / count, 1e-20));
      const barHeight = Math.round(this.normalizeDB(magDB, cache.maxDB) * height);
      if (barHeight === 0) continue;

      const gradient = ctx.createLinearGradient(0, height - barHeight, 0, height);
      if (useSynth) {
        gradient.addColorStop(0, '#4f8'); gradient.addColorStop(1, '#af4');
      } else {
        gradient.addColorStop(0, '#4af'); gradient.addColorStop(1, '#fa4');
      }
      ctx.fillStyle = gradient;
      ctx.fillRect(px, height - barHeight, 1, barHeight);
    }

    // Overlay extracted peaks (green)
    if (this.frames && this.frames.length > 0) {
      let bestFrame = this.frames[0];
      let bestDt = Math.abs(bestFrame.time - this.spectrumTime);
      for (const f of this.frames) {
        const dt = Math.abs(f.time - this.spectrumTime);
        if (dt < bestDt) { bestDt = dt; bestFrame = f; }
      }
      ctx.fillStyle = '#4f4';
      for (const peak of bestFrame.peaks) {
        if (peak.freq < this.freqStart || peak.freq > this.freqEnd) continue;
        const x0 = Math.floor(freqToX(peak.freq));
        const x1 = Math.max(x0 + 1, Math.floor(freqToX(this.normToFreq(this.freqLogNorm(peak.freq) + 1 / width))));
        ctx.fillRect(x0, 0, x1 - x0, height);
      }
    }

    // Overlay partial f0 markers (red, 2px)
    ctx.strokeStyle = '#f44';
    ctx.lineWidth = 2;
    for (let p = 0; p < this.partials.length && p < this.keepCount; ++p) {
      const curve = this.partials[p].freqCurve;
      if (!curve || this.spectrumTime < curve.t0 || this.spectrumTime > curve.t3) continue;
      const freq = evalBezier(curve, this.spectrumTime);
      if (freq < this.freqStart || freq > this.freqEnd) continue;
      const x = Math.round(freqToX(freq));
      ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, height); ctx.stroke();
    }

    // Label
    ctx.fillStyle = useSynth ? '#4f8' : '#4af';
    ctx.font = '9px monospace';
    ctx.fillText(useSynth ? 'SYNTH [a]' : 'ORIG [a]', 4, 10);
  }

  // --- View management ---

  updateViewBounds() {
    const full_duration = this.t_max - this.t_min;
    const view_duration = full_duration / this.zoom_factor;

    let t_view_min = this.t_center - view_duration / 2;
    let t_view_max = this.t_center + view_duration / 2;

    if (t_view_min < this.t_min) { t_view_min = this.t_min; t_view_max = this.t_min + view_duration; }
    if (t_view_max > this.t_max) { t_view_max = this.t_max; t_view_min = this.t_max - view_duration; }
    if (t_view_min < this.t_min) t_view_min = this.t_min;
    if (t_view_max > this.t_max) t_view_max = this.t_max;

    this.t_view_min = t_view_min;
    this.t_view_max = t_view_max;

    const actual_duration = this.t_view_max - this.t_view_min;
    this.zoom_factor = full_duration / actual_duration;
    this.t_center = (this.t_view_min + this.t_view_max) / 2;
  }

  setupMouseHandlers() {
    const {canvas, tooltip} = this;

    canvas.addEventListener('mousedown', (e) => {
      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;

      // Check control point drag on selected partial
      if (this.selectedPartial >= 0 && this.selectedPartial < this.partials.length) {
        const ptIdx = this.hitTestControlPoint(x, y, this.partials[this.selectedPartial]);
        if (ptIdx >= 0) {
          this.dragState = { pointIndex: ptIdx };
          canvas.style.cursor = 'grabbing';
          e.preventDefault();
          return;
        }
      }

      // Otherwise: select partial by click
      const idx = this.hitTestPartial(x, y);
      this.selectPartial(idx);
    });

    canvas.addEventListener('mousemove', (e) => {
      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;

      if (this.dragState) {
        const t = Math.max(0, Math.min(this.t_max, this.canvasToTime(x)));
        const v = Math.max(this.freqStart, Math.min(this.freqEnd, this.canvasToFreq(y)));
        const partial = this.partials[this.selectedPartial];
        const i = this.dragState.pointIndex;
        partial.freqCurve['t' + i] = t;
        partial.freqCurve['v' + i] = v;
        this.render();
        e.preventDefault();
        return;
      }

      this.mouseX = x;
      this.drawMouseCursor(x);

      const time = this.canvasToTime(x);
      const freq = this.canvasToFreq(y);
      const intensity = this.getIntensityAt(time, freq);

      if (this.playheadTime < 0) {
        this.spectrumTime = time;
        this.renderSpectrum();
      }

      // Cursor hint for control points
      if (this.selectedPartial >= 0 && this.selectedPartial < this.partials.length) {
        const ptIdx = this.hitTestControlPoint(x, y, this.partials[this.selectedPartial]);
        canvas.style.cursor = ptIdx >= 0 ? 'grab' : 'crosshair';
      } else {
        canvas.style.cursor = 'crosshair';
      }

      tooltip.style.left = (e.clientX + 10) + 'px';
      tooltip.style.top  = (e.clientY + 10) + 'px';
      tooltip.style.display = 'block';
      tooltip.textContent = `${time.toFixed(3)}s, ${freq.toFixed(1)}Hz, ${intensity.toFixed(1)}dB`;
    });

    canvas.addEventListener('mouseleave', () => {
      this.mouseX = -1;
      this.drawMouseCursor(-1);
      tooltip.style.display = 'none';
    });

    canvas.addEventListener('mouseup', () => {
      if (this.dragState) {
        this.dragState = null;
        canvas.style.cursor = 'crosshair';
        if (this.onPartialSelect) this.onPartialSelect(this.selectedPartial);
      }
    });

    canvas.addEventListener('wheel', (e) => {
      e.preventDefault();
      const delta = e.deltaY !== 0 ? e.deltaY : e.deltaX;

      if (e.shiftKey) {
        const rect = canvas.getBoundingClientRect();
        const mouseTime = this.canvasToTime(e.clientX - rect.left);
        const zoomDelta = delta > 0 ? 1.2 : 1 / 1.2;
        const newZoom = this.zoom_factor * zoomDelta;
        if (newZoom < 1.0) {
          this.zoom_factor = 1.0;
          this.t_center = (this.t_max + this.t_min) / 2;
        } else {
          this.zoom_factor = newZoom;
          const newDuration = (this.t_max - this.t_min) / this.zoom_factor;
          const mouseRatio = (mouseTime - this.t_view_min) / (this.t_view_max - this.t_view_min);
          this.t_center = mouseTime - newDuration * (mouseRatio - 0.5);
        }
      } else {
        const timeDuration = this.t_view_max - this.t_view_min;
        this.t_center += (delta / 100) * timeDuration * 0.1;
      }

      this.updateViewBounds();
      this.render();
    });
  }

  // --- Utilities ---

  getAxisStep(range) {
    const steps = [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 5000];
    const targetStep = range / 8;
    for (const step of steps) { if (step >= targetStep) return step; }
    return steps[steps.length - 1];
  }

  drawAxes() {
    const {ctx, canvas} = this;
    const width = canvas.width;
    const height = canvas.height;

    ctx.strokeStyle = '#666';
    ctx.fillStyle = '#aaa';
    ctx.font = '11px monospace';
    ctx.lineWidth = 1;

    // Time axis
    const timeDuration = this.t_view_max - this.t_view_min;
    const timeStep = this.getAxisStep(timeDuration);
    let t = Math.ceil(this.t_view_min / timeStep) * timeStep;
    while (t <= this.t_view_max) {
      const x = this.timeToX(t);
      ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, height); ctx.stroke();
      ctx.fillText(t.toFixed(2) + 's', x + 2, height - 4);
      t += timeStep;
    }

    // Frequency axis (log-spaced ticks)
    for (const f of [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 16000]) {
      if (f < this.freqStart || f > this.freqEnd) continue;
      const y = this.freqToY(f);
      ctx.beginPath(); ctx.moveTo(0, y); ctx.lineTo(width, y); ctx.stroke();
      ctx.fillText((f >= 1000 ? (f/1000).toFixed(0) + 'k' : f.toFixed(0)) + 'Hz', 2, y - 2);
    }
  }
}

// Bezier evaluation (shared utility)
function evalBezier(curve, t) {
  let u = (t - curve.t0) / (curve.t3 - curve.t0);
  u = Math.max(0, Math.min(1, u));
  const u1 = 1 - u;
  return u1*u1*u1 * curve.v0 +
         3*u1*u1*u * curve.v1 +
         3*u1*u*u  * curve.v2 +
         u*u*u     * curve.v3;
}