// 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;

    // 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;

    // Setup event handlers
    this.setupMouseHandlers();

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

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

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

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

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

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

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

  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;

    // Clamp to [t_min, t_max]
    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;

    // Recompute zoom_factor and t_center from clamped values
    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;
  }

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

  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 timeDuration = this.t_view_max - this.t_view_min;
    const x = (this.playheadTime - this.t_view_min) / timeDuration * canvas.width;

    ctx.strokeStyle = '#f00';
    ctx.lineWidth = 2;
    ctx.beginPath();
    ctx.moveTo(x, 0);
    ctx.lineTo(x, canvas.height);
    ctx.stroke();
  }

  // Render spectrogram background
  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 viewDuration = this.t_view_max - this.t_view_min;

    // Map view bounds to frame indices
    const startFrameIdx = Math.floor(this.t_view_min * sampleRate / hopSize);
    const endFrameIdx = Math.ceil(this.t_view_max * sampleRate / hopSize);
    const numFrames = stftCache.getNumFrames();

    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;

      // Compute frame time range
      const frameTime = frame.time;
      const frameTimeEnd = frameTime + frameDuration;

      const xStart = Math.floor((frameTime - this.t_view_min) / viewDuration * width);
      const xEnd = Math.ceil((frameTimeEnd - this.t_view_min) / viewDuration * width);
      const frameWidth = Math.max(1, xEnd - xStart);

      // Draw frequency bins
      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 normalized = (magDB - (stftCache.maxDB - 80)) / 80;
        const clamped = Math.max(0, Math.min(1, normalized));
        // Power law for better peak visibility
        const intensity = Math.pow(clamped, 2.);

        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 color = this.getSpectrogramColor(intensity);
        ctx.fillStyle = `rgba(${color.r},${color.g},${color.b}, 0.5)`;
        ctx.fillRect(xStart, y0, frameWidth, binHeight);
      }
    }
  }

  // Render extracted partials
  renderPartials() {
    const {ctx, canvas, partials} = this;
    const width = canvas.width;
    const height = canvas.height;

    const colors = [
      '#f44', '#4f4', '#44f', '#ff4', '#f4f', '#4ff',
      '#fa4', '#4fa', '#a4f', '#af4', '#f4a', '#4af'
    ];

    const timeDuration = this.t_view_max - this.t_view_min;

    for (let p = 0; p < partials.length; ++p) {
      const partial = partials[p];
      const color = colors[p % colors.length];
      ctx.globalAlpha = p < this.keepCount ? 1.0 : 0.5;

      // Draw 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 = (t - this.t_view_min) / timeDuration * width;
        const y = this.freqToY(f);

        if (!started) {
          ctx.moveTo(x, y);
          started = true;
        } else {
          ctx.lineTo(x, y);
        }
      }

      if (started) ctx.stroke();

      // Draw bezier curve
      if (partial.freqCurve) {
        ctx.strokeStyle = color;
        ctx.lineWidth = 2;
        ctx.beginPath();

        const curve = partial.freqCurve;
        const numSteps = 50;

        started = false;
        for (let i = 0; i <= numSteps; ++i) {
          const t = curve.t0 + (curve.t3 - curve.t0) * i / numSteps;
          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 = (t - this.t_view_min) / timeDuration * width;
          const y = this.freqToY(freq);

          if (!started) {
            ctx.moveTo(x, y);
            started = true;
          } else {
            ctx.lineTo(x, y);
          }
        }

        if (started) ctx.stroke();

        // Draw control points
        ctx.fillStyle = color;
        this.drawControlPoint(curve.t0, curve.v0);
        this.drawControlPoint(curve.t1, curve.v1);
        this.drawControlPoint(curve.t2, curve.v2);
        this.drawControlPoint(curve.t3, curve.v3);
      }
    }

    ctx.globalAlpha = 1.0;
  }

  // Render raw peaks from mq_extract (before partial tracking)
  renderPeaks() {
    const {ctx, canvas, frames} = this;
    if (!frames || frames.length === 0) return;

    const timeDuration = this.t_view_max - this.t_view_min;

    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 = (t - this.t_view_min) / timeDuration * canvas.width;

      for (const peak of frame.peaks) {
        if (peak.freq < this.freqStart || peak.freq > this.freqEnd) continue;

        const y = this.freqToY(peak.freq);
        ctx.fillRect(x - 1, y - 1, 3, 3);
      }
    }
  }

  // Draw control point
  drawControlPoint(t, v) {
    if (t < this.t_view_min || t > this.t_view_max) return;
    if (v < this.freqStart || v > this.freqEnd) return;

    const timeDuration = this.t_view_max - this.t_view_min;

    const x = (t - this.t_view_min) / timeDuration * this.canvas.width;
    const y = this.freqToY(v);

    this.ctx.beginPath();
    this.ctx.arc(x, y, 4, 0, 2 * Math.PI);
    this.ctx.fill();

    this.ctx.strokeStyle = '#fff';
    this.ctx.lineWidth = 1;
    this.ctx.stroke();
  }

  // Draw axes with ticks and labels
  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;

    const timeDuration = this.t_view_max - this.t_view_min;

    // Time axis
    const timeStep = this.getAxisStep(timeDuration);
    let t = Math.ceil(this.t_view_min / timeStep) * timeStep;
    while (t <= this.t_view_max) {
      const x = (t - this.t_view_min) / timeDuration * width;

      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)
    const freqTicks = [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 16000];
    for (const f of freqTicks) {
      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();

      const label = f >= 1000 ? (f/1000).toFixed(0) + 'k' : f.toFixed(0);
      ctx.fillText(label + 'Hz', 2, y - 2);
    }
  }

  // Setup mouse event handlers
  setupMouseHandlers() {
    const {canvas, tooltip} = this;

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

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

      // Update spectrum time when not playing
      if (this.playheadTime < 0) {
        this.spectrumTime = time;
        this.renderSpectrum();
      }

      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', () => {
      tooltip.style.display = 'none';
    });

    // Mouse wheel: scroll (default) or zoom (shift)
    canvas.addEventListener('wheel', (e) => {
      e.preventDefault();

      const delta = e.deltaY !== 0 ? e.deltaY : e.deltaX;

      if (e.shiftKey) {
        // Zoom in/out around mouse position
        const rect = canvas.getBoundingClientRect();
        const x = e.clientX - rect.left;
        const mouseTime = this.canvasToTime(x);

        const zoomDelta = delta > 0 ? 1.2 : 1 / 1.2;
        const newZoomFactor = this.zoom_factor * zoomDelta;

        // Clamp zoom to [1.0, inf]
        if (newZoomFactor < 1.0) {
          this.zoom_factor = 1.0;
          this.t_center = (this.t_max + this.t_min) / 2;
        } else {
          this.zoom_factor = newZoomFactor;
          const newDuration = (this.t_max - this.t_min) / this.zoom_factor;
          const oldDuration = this.t_view_max - this.t_view_min;
          const mouseRatio = (mouseTime - this.t_view_min) / oldDuration;
          this.t_center = mouseTime - newDuration * (mouseRatio - 0.5);
        }
      } else {
        // Scroll left/right
        const timeDuration = this.t_view_max - this.t_view_min;
        const scrollAmount = (delta / 100) * timeDuration * 0.1;
        this.t_center += scrollAmount;
      }

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

  // Coordinate conversion
  canvasToTime(x) {
    return this.t_view_min + (x / this.canvas.width) * (this.t_view_max - this.t_view_min);
  }

  // freq -> canvas Y (log scale)
  freqToY(freq) {
    const logMin = Math.log2(this.freqStart);
    const logMax = Math.log2(this.freqEnd);
    const norm = (Math.log2(Math.max(freq, this.freqStart)) - logMin) / (logMax - logMin);
    return this.canvas.height * (1 - norm);
  }

  // canvas Y -> freq (log scale, inverse of freqToY)
  canvasToFreq(y) {
    const logMin = Math.log2(this.freqStart);
    const logMax = Math.log2(this.freqEnd);
    const norm = 1 - (y / this.canvas.height);
    return Math.pow(2, logMin + norm * (logMax - logMin));
  }

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

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

  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;

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

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

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

    // Log-scale frequency mapping (matches main spectrogram)
    const logMin = Math.log2(this.freqStart);
    const logMax = Math.log2(this.freqEnd);
    const freqToX = (freq) => {
      const norm = (Math.log2(Math.max(freq, this.freqStart)) - logMin) / (logMax - logMin);
      return norm * width;
    };

    // Draw histogram bars — one per pixel column
    for (let px = 0; px < width; ++px) {
      const fStart = Math.pow(2, logMin + (px / width) * (logMax - logMin));
      const fEnd   = Math.pow(2, logMin + ((px + 1) / width) * (logMax - logMin));

      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 intensity = Math.max(0, Math.min(1, (magDB - (cache.maxDB - 80)) / 80));
      const barHeight = Math.round(intensity * 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 as red histogram bars
    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 = '#f44';
      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(peak.freq * Math.pow(2, 1 / width))));
        ctx.fillRect(x0, 0, x1 - x0, height);
      }
    }

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

  // 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 targetSteps = 8;
    const targetStep = range / targetSteps;

    for (const step of steps) {
      if (step >= targetStep) return step;
    }

    return steps[steps.length - 1];
  }

  getSpectrogramColor(intensity) {
    const v = Math.floor(intensity * 255);
    return {r: v, g: v, b: v};
  }
}

// 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;
}