app/src/App.jsx

import { useState, useEffect, useRef, useCallback, useMemo, lazy, Suspense } from 'react';
import * as Tone from 'tone';
import PianoRoll from './components/PianoRoll';
const SheetMusic = lazy(() => import('./components/SheetMusic'));
const GuitarTab = lazy(() => import('./components/GuitarTab'));
const BassTab = lazy(() => import('./components/BassTab'));
const DrumSheet = lazy(() => import('./components/DrumSheet'));
import Controls from './components/Controls';
import OctopusLogo from './components/OctopusLogo';
import { useMidi } from './hooks/useMidi';
import { usePlayback } from './hooks/usePlayback';
import { buildKeyboardLayout } from './utils/midiHelpers';

const API_BASE = import.meta.env.DEV ? 'http://localhost:8000' : '';

// Standard tunings (MIDI pitch of each open string, low to high)
const GUITAR_TUNING = [40, 45, 50, 55, 59, 64];
const BASS_TUNING = [28, 33, 38, 43];

/** Convert tab JSON events to the note format used by scheduleNotes. */
function tabToNotes(tabData, instrumentOverride) {
  if (!tabData?.events) return [];
  const tuning = tabData.instrument === 'guitar' ? GUITAR_TUNING : BASS_TUNING;
  const instrument = instrumentOverride || tabData.instrument;
  const notes = [];
  for (const event of tabData.events) {
    for (let s = 0; s < event.frets.length; s++) {
      const fret = event.frets[s];
      if (fret === null || fret === undefined) continue;
      notes.push({
        midi: tuning[s] + fret,
        time: event.time,
        duration: event.duration,
        velocity: 0.7,
        instrument,
        hand: 'right',
      });
    }
  }
  notes.sort((a, b) => a.time - b.time);
  return notes;
}

/** Convert drum tab JSON events to the note format used by scheduleNotes. */
function drumTabToNotes(drumData) {
  if (!drumData?.events) return [];
  return drumData.events.map(event => ({
    midi: 0, // not used for drums
    time: event.time,
    duration: 0.1, // drums are instantaneous
    velocity: event.velocity || 0.7,
    instrument: `drum-${event.lane}`,
    hand: 'right',
  }));
}

// App states: 'upload' -> 'loading' -> 'player'

function UploadScreen({ onFileSelected }) {
  const [isDragging, setIsDragging] = useState(false);
  const [errorMsg, setErrorMsg] = useState('');
  const [mode, setMode] = useState('solo'); // 'solo' | 'full'
  const fileInputRef = useRef(null);

  const handleFile = useCallback((file) => {
    if (!file) return;
    const ext = file.name.split('.').pop().toLowerCase();
    if (!['mp3', 'm4a', 'wav', 'ogg', 'flac'].includes(ext)) {
      setErrorMsg('Please upload an audio file (MP3, M4A, WAV, OGG, or FLAC)');
      return;
    }
    setErrorMsg('');
    onFileSelected(file, mode);
  }, [onFileSelected, mode]);

  const handleDrop = useCallback((e) => {
    e.preventDefault();
    setIsDragging(false);
    handleFile(e.dataTransfer.files[0]);
  }, [handleFile]);

  const handleDragOver = useCallback((e) => {
    e.preventDefault();
    setIsDragging(true);
  }, []);

  const handleDragLeave = useCallback(() => {
    setIsDragging(false);
  }, []);

  const handleFileSelect = useCallback((e) => {
    handleFile(e.target.files[0]);
  }, [handleFile]);

  return (
    <div className="upload-screen">
      <div className="upload-content">
        <div className="upload-logo">
          <OctopusLogo size={80} />
          <h1>Mr. Octopus</h1>
          <p className="upload-tagline">Your AI piano teacher</p>
        </div>

        <div className="upload-mode-tabs">
          <button
            className={`upload-mode-tab ${mode === 'solo' ? 'active' : ''}`}
            onClick={() => setMode('solo')}
          >
            Solo Piano
          </button>
          <button
            className={`upload-mode-tab ${mode === 'full' ? 'active' : ''}`}
            onClick={() => setMode('full')}
          >
            Full Song
          </button>
        </div>

        <p className="upload-description">
          {mode === 'solo'
            ? 'Drop a song and Mr. Octopus will transcribe it into a piano tutorial you can follow along with, note by note. Works best with clearly recorded solo piano pieces.'
            : 'Drop any song and Mr. Octopus will separate the instruments using AI, then transcribe the melody (piano, guitar, synths) and bass parts. Works with full band recordings, even AI-generated music.'}
        </p>

        <div
          className={`drop-zone ${isDragging ? 'dragging' : ''}`}
          onDrop={handleDrop}
          onDragOver={handleDragOver}
          onDragLeave={handleDragLeave}
          onClick={() => fileInputRef.current?.click()}
        >
          <div className="drop-icon">&#9835;</div>
          <p>Drag & drop an audio file</p>
          <p className="drop-hint">MP3, M4A, WAV, OGG, FLAC</p>
          <input
            ref={fileInputRef}
            type="file"
            accept="audio/*,.m4a,.mp3,.wav,.ogg,.flac"
            onChange={handleFileSelect}
            hidden
          />
        </div>

        <div className="copyright-notice">
          Please only upload audio you have the rights to use.
        </div>

        {errorMsg && (
          <div className="upload-error">{errorMsg}</div>
        )}
      </div>
    </div>
  );
}

function LoadingScreen({ status, estimate }) {
  return (
    <div className="upload-screen">
      <div className="upload-processing">
        <div className="processing-logo">
          <OctopusLogo size={72} />
        </div>
        <h2>{status}</h2>
        <p className="loading-sub">{estimate || 'This usually takes 20-30 seconds'}</p>
        <div className="loading-bar">
          <div className="loading-bar-fill" />
        </div>
      </div>
    </div>
  );
}

export default function App() {
  const containerRef = useRef(null);
  const [dimensions, setDimensions] = useState({ width: 800, height: 600 });
  const [screen, setScreen] = useState('upload'); // 'upload' | 'loading' | 'player'
  const [loadingStatus, setLoadingStatus] = useState('');
  const [loadingEstimate, setLoadingEstimate] = useState('');
  const [chords, setChords] = useState([]);
  const [activeTab, setActiveTab] = useState('roll'); // 'roll' | 'sheet' | 'guitar-acoustic' | 'guitar-electric' | 'bass' | 'drums'
  const [songMode, setSongMode] = useState('solo'); // 'solo' | 'full'
  const [guitarTab, setGuitarTab] = useState(null);
  const [bassTab, setBassTab] = useState(null);
  const [drumTab, setDrumTab] = useState(null);

  const { notes, totalDuration, fileName, midiObject, loadFromUrl, loadFromBlob } = useMidi();

  const {
    isPlaying,
    currentTimeRef,
    activeNotes,
    tempo,
    samplesLoaded,
    loopStart,
    loopEnd,
    isLooping,
    originalAudioOn,
    originalVolume,
    togglePlayPause,
    pause,
    setTempo,
    seekTo,
    scheduleNotes,
    setLoopA,
    setLoopB,
    clearLoop,
    loadOriginalAudio,
    setOriginalAudioOn,
    setOriginalVolume,
  } = usePlayback();

  // In full song mode, filter bass notes from the piano roll (they have their own tab)
  const pianoRollNotes = useMemo(() => {
    if (songMode === 'full') {
      return notes.filter(n => n.instrument !== 'bass');
    }
    return notes;
  }, [notes, songMode]);

  // When samples are loaded and we have notes, transition to player
  useEffect(() => {
    if (screen === 'loading' && samplesLoaded && notes.length > 0) {
      setScreen('player');
    }
  }, [screen, samplesLoaded, notes.length]);

  const stopPlayback = useCallback(() => {
    if (isPlaying) pause();
    seekTo(0);
  }, [isPlaying, pause, seekTo]);

  const loadResult = useCallback(async (data, fileName) => {
    setLoadingStatus('Loading piano sounds...');
    const midiRes = await fetch(`${API_BASE}${data.midi_url}`);
    const blob = await midiRes.blob();
    loadFromBlob(blob, fileName.replace(/\.[^.]+$/, '.mid'));

    if (data.chords) {
      const chordList = data.chords?.chords || data.chords || [];
      setChords(Array.isArray(chordList) ? chordList : []);
    }

    // Load original audio for playback alongside transcription
    if (data.audio_url) {
      loadOriginalAudio(`${API_BASE}${data.audio_url}`);
    }

    // Full song mode: fetch tab data
    if (data.mode === 'full') {
      setSongMode('full');
      setActiveTab('roll');
      if (data.guitar_tab_url) {
        try {
          const res = await fetch(`${API_BASE}${data.guitar_tab_url}`);
          if (res.ok) setGuitarTab(await res.json());
        } catch { /* tab data optional */ }
      }
      if (data.bass_tab_url) {
        try {
          const res = await fetch(`${API_BASE}${data.bass_tab_url}`);
          if (res.ok) setBassTab(await res.json());
        } catch { /* tab data optional */ }
      }
      if (data.drum_tab_url) {
        try {
          const res = await fetch(`${API_BASE}${data.drum_tab_url}`);
          if (res.ok) setDrumTab(await res.json());
        } catch { /* tab data optional */ }
      }
    } else {
      setSongMode('solo');
      setActiveTab('roll');
    }
  }, [loadFromBlob, loadOriginalAudio]);

  const handleFileSelected = useCallback(async (file, mode = 'solo') => {
    stopPlayback();
    setScreen('loading');

    if (mode === 'full') {
      // Full song: async with polling
      setLoadingStatus('Uploading...');
      setLoadingEstimate('This usually takes 2-4 minutes');
      try {
        const form = new FormData();
        form.append('file', file);
        const res = await fetch(`${API_BASE}/api/transcribe-full`, {
          method: 'POST',
          body: form,
        });
        if (!res.ok) {
          const err = await res.json().catch(() => ({ detail: res.statusText }));
          throw new Error(err.detail || 'Failed to start transcription');
        }
        const { job_id } = await res.json();

        // Poll for status
        const poll = async () => {
          try {
            const statusRes = await fetch(`${API_BASE}/api/jobs/${job_id}/status`);
            const status = await statusRes.json();

            if (status.error) {
              throw new Error(status.error);
            }

            setLoadingStatus(status.label);

            if (status.done && status.result) {
              await loadResult(status.result, file.name);
            } else {
              setTimeout(poll, 2000);
            }
          } catch (e) {
            setScreen('upload');
            alert(e.message || 'Something went wrong. Please try again.');
          }
        };
        poll();
      } catch (e) {
        setScreen('upload');
        alert(e.message || 'Something went wrong. Please try again.');
      }
    } else {
      // Solo piano: existing synchronous flow
      setLoadingStatus('Transcribing your song...');
      setLoadingEstimate('This usually takes 20-30 seconds');
      try {
        const form = new FormData();
        form.append('file', file);
        const res = await fetch(`${API_BASE}/api/transcribe`, {
          method: 'POST',
          body: form,
        });
        if (!res.ok) {
          const err = await res.json().catch(() => ({ detail: res.statusText }));
          throw new Error(err.detail || 'Transcription failed');
        }
        const data = await res.json();
        await loadResult(data, file.name);
      } catch (e) {
        setScreen('upload');
        alert(e.message || 'Something went wrong. Please try again.');
      }
    }
  }, [loadResult, stopPlayback]);

  const handleNewSong = useCallback(() => {
    stopPlayback();
    loadOriginalAudio(null);
    setOriginalAudioOn(false);
    setScreen('upload');
    setChords([]);
    setSongMode('solo');
    setGuitarTab(null);
    setBassTab(null);
    setDrumTab(null);
    setActiveTab('roll');
  }, [stopPlayback, loadOriginalAudio, setOriginalAudioOn]);

  // Reschedule audio when notes change or active tab switches instrument
  useEffect(() => {
    let notesToPlay;
    if (activeTab === 'guitar-acoustic' && guitarTab) {
      notesToPlay = tabToNotes(guitarTab, 'guitar-acoustic');
    } else if (activeTab === 'guitar-electric' && guitarTab) {
      notesToPlay = tabToNotes(guitarTab, 'guitar-electric');
    } else if (activeTab === 'bass' && bassTab) {
      notesToPlay = tabToNotes(bassTab);
    } else if (activeTab === 'drums' && drumTab) {
      notesToPlay = drumTabToNotes(drumTab);
    } else {
      // On piano roll / sheet music: only play piano notes (filter bass in full mode)
      notesToPlay = pianoRollNotes;
    }
    if (notesToPlay.length > 0) {
      scheduleNotes(notesToPlay, totalDuration);
    }
  }, [activeTab, guitarTab, bassTab, drumTab, pianoRollNotes, totalDuration, scheduleNotes]);

  // Handle resize
  useEffect(() => {
    const el = containerRef.current;
    if (!el) return;

    const ro = new ResizeObserver(([entry]) => {
      const { width, height } = entry.contentRect;
      if (width > 0 && height > 0) {
        setDimensions({ width, height });
      }
    });
    ro.observe(el);
    return () => ro.disconnect();
  }, [screen]);

  const keyboardLayout = buildKeyboardLayout(dimensions.width);

  const handleTogglePlay = useCallback(async () => {
    if (!samplesLoaded) return;
    await Tone.start();
    togglePlayPause();
  }, [togglePlayPause, samplesLoaded]);

  if (screen === 'upload') {
    return <UploadScreen onFileSelected={handleFileSelected} />;
  }

  if (screen === 'loading') {
    return <LoadingScreen status={loadingStatus} estimate={loadingEstimate} />;
  }

  return (
    <div className="app">
      <Controls
        isPlaying={isPlaying}
        togglePlayPause={handleTogglePlay}
        tempo={tempo}
        setTempo={setTempo}
        currentTimeRef={currentTimeRef}
        totalDuration={totalDuration}
        seekTo={seekTo}
        fileName={fileName}
        onNewSong={handleNewSong}
        loopStart={loopStart}
        loopEnd={loopEnd}
        isLooping={isLooping}
        onSetLoopA={setLoopA}
        onSetLoopB={setLoopB}
        onClearLoop={clearLoop}
        originalAudioOn={originalAudioOn}
        setOriginalAudioOn={setOriginalAudioOn}
        originalVolume={originalVolume}
        setOriginalVolume={setOriginalVolume}
      />
      <div className="view-tabs">
        <button
          className={`view-tab ${activeTab === 'roll' ? 'active' : ''}`}
          onClick={() => setActiveTab('roll')}
        >
          Piano Roll
        </button>
        <button
          className={`view-tab ${activeTab === 'sheet' ? 'active' : ''}`}
          onClick={() => setActiveTab('sheet')}
        >
          Sheet Music
        </button>
        {songMode === 'full' && (
          <>
            <button
              className={`view-tab ${activeTab === 'guitar-acoustic' ? 'active' : ''}`}
              onClick={() => setActiveTab('guitar-acoustic')}
            >
              Acoustic Guitar
            </button>
            <button
              className={`view-tab ${activeTab === 'guitar-electric' ? 'active' : ''}`}
              onClick={() => setActiveTab('guitar-electric')}
            >
              Electric Guitar
            </button>
            <button
              className={`view-tab ${activeTab === 'bass' ? 'active' : ''}`}
              onClick={() => setActiveTab('bass')}
            >
              Bass Tab
            </button>
            <button
              className={`view-tab ${activeTab === 'drums' ? 'active' : ''}`}
              onClick={() => setActiveTab('drums')}
            >
              Drums
            </button>
          </>
        )}
      </div>
      {activeTab === 'roll' ? (
        <div className="canvas-container" ref={containerRef}>
          <PianoRoll
            notes={pianoRollNotes}
            currentTimeRef={currentTimeRef}
            activeNotes={activeNotes}
            keyboardLayout={keyboardLayout}
            width={dimensions.width}
            height={dimensions.height}
            loopStart={loopStart}
            loopEnd={loopEnd}
            chords={chords}
          />
        </div>
      ) : activeTab === 'sheet' ? (
        <Suspense fallback={<div className="sheet-music-empty"><p>Loading sheet music...</p></div>}>
          <SheetMusic midiObject={midiObject} fileName={fileName} currentTimeRef={currentTimeRef} isPlaying={isPlaying} />
        </Suspense>
      ) : activeTab === 'guitar-acoustic' || activeTab === 'guitar-electric' ? (
        <div className="canvas-container" ref={containerRef}>
          <Suspense fallback={<div className="tab-empty"><p>Loading guitar tab...</p></div>}>
            <GuitarTab
              tabData={guitarTab}
              currentTimeRef={currentTimeRef}
              width={dimensions.width}
              height={dimensions.height}
              chords={chords}
            />
          </Suspense>
        </div>
      ) : activeTab === 'bass' ? (
        <div className="canvas-container" ref={containerRef}>
          <Suspense fallback={<div className="tab-empty"><p>Loading bass tab...</p></div>}>
            <BassTab
              tabData={bassTab}
              currentTimeRef={currentTimeRef}
              width={dimensions.width}
              height={dimensions.height}
            />
          </Suspense>
        </div>
      ) : activeTab === 'drums' ? (
        <div className="canvas-container" ref={containerRef}>
          <Suspense fallback={<div className="tab-empty"><p>Loading drum data...</p></div>}>
            <DrumSheet
              tabData={drumTab}
              currentTimeRef={currentTimeRef}
              width={dimensions.width}
              height={dimensions.height}
            />
          </Suspense>
        </div>
      ) : null}
    </div>
  );
}