app/src/components/SheetMusic.jsx

import { useEffect, useRef, useMemo, useCallback } from 'react';
import { Renderer, Stave, StaveNote, Voice, Formatter, Dot, StaveConnector } from 'vexflow';
import { midiToMeasures } from '../utils/midiToNotation';

const MEASURES_PER_LINE = 4;
const STAVE_WIDTH = 250;
const TREBLE_Y_OFFSET = 0;
const BASS_Y_OFFSET = 80;
const LINE_HEIGHT = 200;
const LEFT_MARGIN = 60;
const TOP_MARGIN = 70;

const ACTIVE_COLOR = '#8b5cf6';
const DEFAULT_COLOR = '#000000';

function createStaveNote(noteData, clef) {
  const note = new StaveNote({
    keys: noteData.keys,
    duration: noteData.duration,
    clef: clef === 'treble' ? 'treble' : 'bass',
  });
  for (let d = 0; d < noteData.dots; d++) {
    Dot.buildAndAttach([note]);
  }
  return note;
}

export default function SheetMusic({ midiObject, fileName, currentTimeRef, isPlaying }) {
  const containerRef = useRef(null);
  const scrollRef = useRef(null);
  const progressRef = useRef(null);
  const noteRefsRef = useRef([]); // { svgEl, absoluteBeat, isRest }
  const animFrameRef = useRef(null);
  const measureLayoutRef = useRef([]); // { x, y, width, lineIdx, measureIdx }

  const { measures, timeSignature, bpm } = useMemo(() => {
    if (!midiObject) return { measures: [], timeSignature: [4, 4], bpm: 120 };
    return midiToMeasures(midiObject);
  }, [midiObject]);

  const beatsPerMeasure = timeSignature[0];
  const secondsPerBeat = 60 / bpm;

  const render = useCallback(() => {
    const container = containerRef.current;
    if (!container || measures.length === 0) return;

    container.innerHTML = '';
    noteRefsRef.current = [];
    measureLayoutRef.current = [];

    const numLines = Math.ceil(measures.length / MEASURES_PER_LINE);
    const totalWidth = LEFT_MARGIN + STAVE_WIDTH * MEASURES_PER_LINE + 40;
    const totalHeight = TOP_MARGIN + numLines * LINE_HEIGHT + 40;

    const renderer = new Renderer(container, Renderer.Backends.SVG);
    renderer.resize(totalWidth, totalHeight);
    const context = renderer.getContext();
    context.setFont('Arial', 10);

    // Draw title
    const title = (fileName || 'Untitled').replace(/\.[^.]+$/, '');
    const svgEl = container.querySelector('svg');
    if (svgEl) {
      const text = document.createElementNS('http://www.w3.org/2000/svg', 'text');
      text.setAttribute('x', totalWidth / 2);
      text.setAttribute('y', 30);
      text.setAttribute('text-anchor', 'middle');
      text.setAttribute('font-family', 'Inter, Arial, sans-serif');
      text.setAttribute('font-size', '20');
      text.setAttribute('font-weight', '700');
      text.setAttribute('fill', '#1a1a2e');
      text.textContent = title;
      svgEl.appendChild(text);
    }

    for (let lineIdx = 0; lineIdx < numLines; lineIdx++) {
      const startMeasure = lineIdx * MEASURES_PER_LINE;
      const endMeasure = Math.min(startMeasure + MEASURES_PER_LINE, measures.length);
      const lineY = TOP_MARGIN + lineIdx * LINE_HEIGHT;

      for (let m = startMeasure; m < endMeasure; m++) {
        const measureIdx = m - startMeasure;
        const x = LEFT_MARGIN + measureIdx * STAVE_WIDTH;
        const isFirst = measureIdx === 0;
        const measure = measures[m];

        // Store measure layout for progress line positioning
        measureLayoutRef.current.push({
          globalIdx: m,
          x,
          y: lineY,
          width: STAVE_WIDTH,
          lineIdx,
          measureStart: measure.measureStart, // in beats
        });

        // Create treble stave
        const trebleStave = new Stave(x, lineY + TREBLE_Y_OFFSET, STAVE_WIDTH);
        if (isFirst) {
          trebleStave.addClef('treble');
          if (lineIdx === 0) {
            trebleStave.addTimeSignature(`${timeSignature[0]}/${timeSignature[1]}`);
          }
        }
        trebleStave.setContext(context).draw();

        // Create bass stave
        const bassStave = new Stave(x, lineY + BASS_Y_OFFSET, STAVE_WIDTH);
        if (isFirst) {
          bassStave.addClef('bass');
          if (lineIdx === 0) {
            bassStave.addTimeSignature(`${timeSignature[0]}/${timeSignature[1]}`);
          }
        }
        bassStave.setContext(context).draw();

        // Draw brace and line connector for first measure of each line
        if (isFirst) {
          const brace = new StaveConnector(trebleStave, bassStave);
          brace.setType('brace');
          brace.setContext(context).draw();

          const lineConn = new StaveConnector(trebleStave, bassStave);
          lineConn.setType('singleLeft');
          lineConn.setContext(context).draw();
        }

        // Right barline connector
        const rightConn = new StaveConnector(trebleStave, bassStave);
        rightConn.setType('singleRight');
        rightConn.setContext(context).draw();

        // Create and draw treble voice, collecting note refs
        try {
          const trebleNotes = measure.treble.map((n, ni) => {
            const sn = createStaveNote(n, 'treble');
            sn._beatData = { absoluteBeat: measure.measureStart + (n.beatOffset || 0), isRest: n.isRest };
            return sn;
          });
          const trebleVoice = new Voice({
            num_beats: timeSignature[0],
            beat_value: timeSignature[1],
          }).setStrict(false);
          trebleVoice.addTickables(trebleNotes);
          new Formatter().joinVoices([trebleVoice]).format([trebleVoice], STAVE_WIDTH - 50);
          trebleVoice.draw(context, trebleStave);

          // Collect SVG element references
          for (const sn of trebleNotes) {
            const el = sn.getSVGElement?.();
            if (el && !sn._beatData.isRest) {
              noteRefsRef.current.push({
                el,
                absoluteBeat: sn._beatData.absoluteBeat,
              });
            }
          }
        } catch (e) {
          // Skip measures that fail to render
        }

        // Create and draw bass voice
        try {
          const bassNotes = measure.bass.map((n, ni) => {
            const sn = createStaveNote(n, 'bass');
            sn._beatData = { absoluteBeat: measure.measureStart + (n.beatOffset || 0), isRest: n.isRest };
            return sn;
          });
          const bassVoice = new Voice({
            num_beats: timeSignature[0],
            beat_value: timeSignature[1],
          }).setStrict(false);
          bassVoice.addTickables(bassNotes);
          new Formatter().joinVoices([bassVoice]).format([bassVoice], STAVE_WIDTH - 50);
          bassVoice.draw(context, bassStave);

          for (const sn of bassNotes) {
            const el = sn.getSVGElement?.();
            if (el && !sn._beatData.isRest) {
              noteRefsRef.current.push({
                el,
                absoluteBeat: sn._beatData.absoluteBeat,
              });
            }
          }
        } catch (e) {
          // Skip measures that fail to render
        }
      }
    }
  }, [measures, timeSignature, fileName]);

  useEffect(() => {
    render();
  }, [render]);

  // Animation loop: update progress line position and note colors
  useEffect(() => {
    const progressEl = progressRef.current;
    const scrollEl = scrollRef.current;
    if (!progressEl || !scrollEl) return;

    let lastColoredBeat = -1;
    let lastLineIdx = -1;

    const tick = () => {
      const currentTime = currentTimeRef?.current ?? 0;
      const currentBeat = currentTime / secondsPerBeat;
      const layout = measureLayoutRef.current;

      if (layout.length === 0) {
        progressEl.style.display = 'none';
        animFrameRef.current = requestAnimationFrame(tick);
        return;
      }

      // Find which measure we're in
      let mLayout = null;
      for (let i = layout.length - 1; i >= 0; i--) {
        if (currentBeat >= layout[i].measureStart - 0.01) {
          mLayout = layout[i];
          break;
        }
      }

      if (!mLayout) {
        mLayout = layout[0];
      }

      // Calculate x position within measure
      const beatInMeasure = currentBeat - mLayout.measureStart;
      const progress = Math.max(0, Math.min(1, beatInMeasure / beatsPerMeasure));

      // Note area starts ~50px from stave left (after clef/time sig) and ends ~10px before right
      const noteAreaStart = mLayout.x + 50;
      const noteAreaEnd = mLayout.x + mLayout.width - 10;
      const noteAreaWidth = noteAreaEnd - noteAreaStart;

      const lineX = noteAreaStart + progress * noteAreaWidth;
      const lineY = mLayout.y - 5;
      const lineHeight = BASS_Y_OFFSET + 80; // span treble + bass staves

      // Enable smooth transition within the same staff line, disable on line jumps
      if (mLayout.lineIdx !== lastLineIdx) {
        progressEl.classList.remove('smooth');
      } else {
        progressEl.classList.add('smooth');
      }
      lastLineIdx = mLayout.lineIdx;

      progressEl.style.display = 'block';
      progressEl.style.left = `${lineX + 20}px`; // +20 for container padding
      progressEl.style.top = `${lineY + 20}px`;
      progressEl.style.height = `${lineHeight}px`;

      // Auto-scroll to keep progress line visible
      const lineScreenY = lineY + 20 - scrollEl.scrollTop;
      const viewHeight = scrollEl.clientHeight;
      if (lineScreenY < 50 || lineScreenY > viewHeight - 100) {
        scrollEl.scrollTo({
          top: lineY - 80,
          behavior: 'smooth',
        });
      }

      // Color notes based on playback position
      // Only update when we've crossed a meaningful threshold
      if (Math.abs(currentBeat - lastColoredBeat) > 0.05) {
        lastColoredBeat = currentBeat;
        for (const noteRef of noteRefsRef.current) {
          const isPast = noteRef.absoluteBeat <= currentBeat + 0.1;
          const color = isPast ? ACTIVE_COLOR : DEFAULT_COLOR;
          if (noteRef._lastColor !== color) {
            noteRef._lastColor = color;
            // Color all child paths and text within the note group
            const paths = noteRef.el.querySelectorAll('path, text, circle, ellipse, line, rect');
            for (const p of paths) {
              p.setAttribute('fill', color);
              p.setAttribute('stroke', color);
            }
          }
        }
      }

      animFrameRef.current = requestAnimationFrame(tick);
    };

    animFrameRef.current = requestAnimationFrame(tick);
    return () => {
      if (animFrameRef.current) cancelAnimationFrame(animFrameRef.current);
    };
  }, [currentTimeRef, secondsPerBeat, beatsPerMeasure]);

  // Reset note colors when not playing and at time 0
  useEffect(() => {
    if (!isPlaying && currentTimeRef?.current < 0.1) {
      for (const noteRef of noteRefsRef.current) {
        noteRef._lastColor = DEFAULT_COLOR;
        const paths = noteRef.el.querySelectorAll('path, text, circle, ellipse, line, rect');
        for (const p of paths) {
          p.setAttribute('fill', DEFAULT_COLOR);
          p.setAttribute('stroke', DEFAULT_COLOR);
        }
      }
    }
  }, [isPlaying, currentTimeRef]);

  const handleDownloadPDF = useCallback(() => {
    const container = containerRef.current;
    if (!container) return;
    const svgEl = container.querySelector('svg');
    if (!svgEl) return;

    // Clone SVG and set white background
    const clone = svgEl.cloneNode(true);
    clone.setAttribute('xmlns', 'http://www.w3.org/2000/svg');

    // Add white background rect
    const bg = document.createElementNS('http://www.w3.org/2000/svg', 'rect');
    bg.setAttribute('width', '100%');
    bg.setAttribute('height', '100%');
    bg.setAttribute('fill', 'white');
    clone.insertBefore(bg, clone.firstChild);

    const svgData = new XMLSerializer().serializeToString(clone);
    const svgBlob = new Blob([svgData], { type: 'image/svg+xml;charset=utf-8' });
    const url = URL.createObjectURL(svgBlob);

    const img = new Image();
    img.onload = () => {
      const canvas = document.createElement('canvas');
      const scale = 2;
      canvas.width = img.width * scale;
      canvas.height = img.height * scale;
      const ctx = canvas.getContext('2d');
      ctx.scale(scale, scale);
      ctx.fillStyle = 'white';
      ctx.fillRect(0, 0, img.width, img.height);
      ctx.drawImage(img, 0, 0);

      canvas.toBlob((blob) => {
        const a = document.createElement('a');
        a.href = URL.createObjectURL(blob);
        const name = (fileName || 'sheet-music').replace(/\.[^.]+$/, '');
        a.download = `${name} - Sheet Music.png`;
        a.click();
        URL.revokeObjectURL(a.href);
      }, 'image/png');

      URL.revokeObjectURL(url);
    };
    img.src = url;
  }, [fileName]);

  if (!midiObject || measures.length === 0) {
    return (
      <div className="sheet-music-empty">
        <p>No sheet music to display. Transcribe a song first.</p>
      </div>
    );
  }

  return (
    <div className="sheet-music-wrapper">
      <div className="sheet-music-toolbar">
        <span className="sheet-music-info">
          {measures.length} measures | {timeSignature[0]}/{timeSignature[1]} | {Math.round(bpm)} BPM
        </span>
        <button className="btn btn-download" onClick={handleDownloadPDF}>
          <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2">
            <path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4" />
            <polyline points="7 10 12 15 17 10" />
            <line x1="12" y1="15" x2="12" y2="3" />
          </svg>
          Download PNG
        </button>
      </div>
      <div className="sheet-music-scroll" ref={scrollRef}>
        <div className="sheet-music-container" ref={containerRef} />
        <div className="sheet-music-progress-line" ref={progressRef} />
      </div>
    </div>
  );
}