import { useState, useRef, useEffect, useCallback } from 'react';
import { INSTRUMENTS, INITIAL_TEMPO, INITIAL_STEPS, MIN_STEPS, MAX_STEPS, NOTE_FREQ_MAP } from '../constants';
import { Grid, BeatData, BassLineGrid } from '../types';

const createEmptyGrid = (steps: number): Grid => {
  return Array.from({ length: INSTRUMENTS.length }, () => Array(steps).fill(false));
};

const createEmptyBassLine = (steps: number): BassLineGrid => {
    return Array.from({ length: steps }, () => []);
}

const playKick = (audioContext: AudioContext, time: number, destination: AudioNode) => {
    const osc = audioContext.createOscillator();
    const subOsc = audioContext.createOscillator();
    const gain = audioContext.createGain();
    const subGain = audioContext.createGain();
    const shaper = audioContext.createWaveShaper();

    shaper.curve = new Float32Array(65536).map((_, i) => {
        const x = (i / 32768) - 1;
        return Math.tanh(3 * x);
    });
    shaper.oversample = '4x';

    osc.connect(gain);
    subOsc.connect(subGain);
    gain.connect(shaper);
    subGain.connect(shaper);
    shaper.connect(destination);

    const duration = 0.6;
    const finalStopTime = time + duration;

    osc.type = 'sine';
    osc.frequency.setValueAtTime(180, time);
    osc.frequency.exponentialRampToValueAtTime(30, time + 0.1);
    gain.gain.setValueAtTime(1.8, time);
    gain.gain.setTargetAtTime(0, time, 0.1);

    subOsc.type = 'sine';
    subOsc.frequency.setValueAtTime(50, time);
    subGain.gain.setValueAtTime(0.35, time);
    subGain.gain.exponentialRampToValueAtTime(0.001, time + 0.2);

    osc.start(time);
    subOsc.start(time);
    osc.stop(finalStopTime);
    subOsc.stop(finalStopTime);
};

const playSnare = (audioContext: AudioContext, time: number, destination: AudioNode) => {
    const noiseGain = audioContext.createGain();
    const noiseFilter = audioContext.createBiquadFilter();
    const osc = audioContext.createOscillator();
    const oscGain = audioContext.createGain();

    const noiseBuffer = audioContext.createBuffer(1, audioContext.sampleRate * 0.5, audioContext.sampleRate);
    const output = noiseBuffer.getChannelData(0);
    for (let i = 0; i < output.length; i++) {
        output[i] = Math.random() * 2 - 1;
    }
    const noiseSource = audioContext.createBufferSource();
    noiseSource.buffer = noiseBuffer;

    noiseFilter.type = 'highpass';
    noiseFilter.frequency.value = 1000;
    noiseSource.connect(noiseFilter);
    noiseFilter.connect(noiseGain);
    noiseGain.connect(destination);

    osc.type = 'triangle';
    osc.frequency.setValueAtTime(200, time);
    osc.connect(oscGain);
    oscGain.connect(destination);

    const duration = 0.2;
    noiseGain.gain.setValueAtTime(1, time);
    noiseGain.gain.exponentialRampToValueAtTime(0.01, time + duration);
    oscGain.gain.setValueAtTime(0.7, time);
    oscGain.gain.exponentialRampToValueAtTime(0.01, time + duration / 2);

    noiseSource.start(time);
    osc.start(time);
    noiseSource.stop(time + duration);
    osc.stop(time + duration);
};

const createHiHatSound = (audioContext: AudioContext, time: number, destination: AudioNode, duration: number) => {
    const fundamental = 40;
    const ratios = [2, 3, 4.16, 5.43, 6.79, 8.21];
    const gain = audioContext.createGain();
    const bandpass = audioContext.createBiquadFilter();
    const highpass = audioContext.createBiquadFilter();

    bandpass.type = 'bandpass';
    bandpass.frequency.value = 10000;
    bandpass.Q.value = 0.5;

    highpass.type = 'highpass';
    highpass.frequency.value = 7000;

    gain.connect(bandpass);
    bandpass.connect(highpass);
    highpass.connect(destination);

    ratios.forEach(ratio => {
        const osc = audioContext.createOscillator();
        osc.type = 'square';
        osc.frequency.value = (fundamental * ratio) + (Math.random() * fundamental * 0.1);
        osc.connect(gain);
        osc.start(time);
        osc.stop(time + duration);
    });

    gain.gain.setValueAtTime(0.00001, time);
    gain.gain.exponentialRampToValueAtTime(0.4, time + 0.02);
    gain.gain.exponentialRampToValueAtTime(0.00001, time + duration);
};

const playHiHat = (audioContext: AudioContext, time: number, destination: AudioNode) => {
    createHiHatSound(audioContext, time, destination, 0.08);
};

const playOpenHat = (audioContext: AudioContext, time: number, destination: AudioNode) => {
    createHiHatSound(audioContext, time, destination, 0.8);
};

const playRide = (audioContext: AudioContext, time: number, destination: AudioNode) => {
    const masterGain = audioContext.createGain();
    const highpass = audioContext.createBiquadFilter();

    highpass.type = 'highpass';
    highpass.frequency.setValueAtTime(800, time);
    highpass.Q.value = 0.8;

    masterGain.connect(highpass);
    highpass.connect(destination);
    
    const tickOsc = audioContext.createOscillator();
    const tickGain = audioContext.createGain();
    tickOsc.type = 'square';
    tickOsc.frequency.setValueAtTime(1200, time);
    
    tickGain.gain.setValueAtTime(0.5, time);
    tickGain.gain.exponentialRampToValueAtTime(0.0001, time + 0.02);

    tickOsc.connect(tickGain);
    tickGain.connect(masterGain);

    const fundamental = 120;
    const ratios = [1.00, 1.41, 2.23, 2.77, 3.14, 4.01];
    ratios.forEach(ratio => {
        const osc = audioContext.createOscillator();
        osc.type = 'square';
        osc.frequency.value = fundamental * ratio + (Math.random() - 0.5) * 5;
        osc.connect(masterGain);
        osc.start(time);
        osc.stop(time + 1.2);
    });

    masterGain.gain.setValueAtTime(0.0001, time);
    masterGain.gain.exponentialRampToValueAtTime(0.6, time + 0.005);
    masterGain.gain.exponentialRampToValueAtTime(0.2, time + 0.1);
    masterGain.gain.exponentialRampToValueAtTime(0.0001, time + 1.2);

    tickOsc.start(time);
    tickOsc.stop(time + 0.03);
};


export const useDrumMachine = (lastMessage: any, sendMessage: (message: any) => void) => {
  const [steps, setSteps] = useState(INITIAL_STEPS);
  const [grid, setGrid] = useState<Grid>(createEmptyGrid(INITIAL_STEPS));
  const [bassLine, setBassLine] = useState<BassLineGrid>(() => createEmptyBassLine(INITIAL_STEPS));
  const [isPlaying, setIsPlaying] = useState(false);
  const [tempo, setTempo] = useState(INITIAL_TEMPO);
  const [currentStep, setCurrentStep] = useState<number | null>(null);
  const [mutes, setMutes] = useState<boolean[]>(() => Array(INSTRUMENTS.length).fill(false));
  const [drumVolume, setDrumVolume] = useState(1);
  const [bassVolume, setBassVolume] = useState(0.4);

  const audioContextRef = useRef<AudioContext | null>(null);
  const audioBuffersRef = useRef<Map<string, AudioBuffer>>(new Map());
  const timerRef = useRef<number | null>(null);
  const lookahead = 25.0;
  const scheduleAheadTime = 0.1;
  const nextNoteTimeRef = useRef<number>(0.0);
  const sequenceStepRef = useRef<number>(0);

  const activeOscillatorsRef = useRef<Map<string, { osc: OscillatorNode; gain: GainNode }>>(new Map());
  const drumMasterGainRef = useRef<GainNode | null>(null);
  const bassMasterGainRef = useRef<GainNode | null>(null);

  const loadSamples = useCallback(async (context: AudioContext) => {
    const newBuffers = new Map<string, AudioBuffer>();
    for (const instrument of INSTRUMENTS) {
      if (!instrument.sampleUrl) continue;
      try {
        const response = await fetch(instrument.sampleUrl);
        const arrayBuffer = await response.arrayBuffer();
        const decodedData = await context.decodeAudioData(arrayBuffer);
        newBuffers.set(instrument.name, decodedData);
      } catch (error) {
        console.error(`Error loading sample for ${instrument.name}:`, error);
      }
    }
    audioBuffersRef.current = newBuffers;
  }, []);
  
  const initAudio = useCallback(async () => {
    if (!audioContextRef.current) {
        audioContextRef.current = new window.AudioContext();
        await loadSamples(audioContextRef.current);
        
        drumMasterGainRef.current = audioContextRef.current.createGain();
        drumMasterGainRef.current.gain.value = drumVolume;
        drumMasterGainRef.current.connect(audioContextRef.current.destination);

        bassMasterGainRef.current = audioContextRef.current.createGain();
        bassMasterGainRef.current.gain.value = bassVolume;
        bassMasterGainRef.current.connect(audioContextRef.current.destination);
    }
  }, [loadSamples, drumVolume, bassVolume]);

  const stopAllBassNotes = useCallback(() => {
    if (!audioContextRef.current || !bassMasterGainRef.current) return;
    const now = audioContextRef.current.currentTime;

    activeOscillatorsRef.current.forEach(({ osc, gain }) => {
        try {
            gain.gain.cancelScheduledValues(now);
            gain.gain.setTargetAtTime(0, now, 0.01);
            osc.stop(now + 0.1);
        } catch (e) { /* Ignore errors */ }
    });
    activeOscillatorsRef.current.clear();
  }, []);

  const setStep = (instrumentIndex: number, stepIndex: number, isActive: boolean) => {
    const newGrid = grid.map(row => [...row]);
    if(newGrid[instrumentIndex]){
        newGrid[instrumentIndex][stepIndex] = isActive;
    }
    setGrid(newGrid);
    sendMessage({ type: 'grid', payload: { grid: newGrid } });
  };

  const setBassNote = (stepIndex: number, note: string) => {
    const newBassLine = bassLine.map(stepNotes => [...stepNotes]);
    const stepNotes = newBassLine[stepIndex];
    const noteIndex = stepNotes.indexOf(note);

    if (noteIndex > -1) {
        stepNotes.splice(noteIndex, 1);
    } else {
        stepNotes.push(note);
    }
    setBassLine(newBassLine);
    sendMessage({ type: 'bassLine', payload: { bassLine: newBassLine } });
  };

  const clearPattern = () => {
    if (timerRef.current) {
        window.clearTimeout(timerRef.current);
        timerRef.current = null;
    }
    const newGrid = createEmptyGrid(steps);
    const newBassLine = createEmptyBassLine(steps);
    setGrid(newGrid);
    setBassLine(newBassLine);
    stopAllBassNotes();
    sendMessage({ type: 'clear', payload: { steps } });
  };

  const toggleMute = (instrumentIndex: number) => {
    const newMutes = [...mutes];
    newMutes[instrumentIndex] = !newMutes[instrumentIndex];
    setMutes(newMutes);
    sendMessage({ type: 'mutes', payload: { mutes: newMutes } });
  };
  
  const resizeGrid = useCallback((newSteps: number) => {
    setGrid(prevGrid => {
        const newGrid = createEmptyGrid(newSteps);
        for(let i=0; i<INSTRUMENTS.length; i++){
            for(let j=0; j<Math.min(prevGrid[i]?.length ?? 0, newSteps); j++){
                newGrid[i][j] = prevGrid[i][j];
            }
        }
        return newGrid;
    });
    setBassLine(prevBassLine => {
        const newBassLine = createEmptyBassLine(newSteps);
        for (let i = 0; i < Math.min(prevBassLine.length, newSteps); i++) {
            newBassLine[i] = prevBassLine[i];
        }
        return newBassLine;
    });
  }, []);

  const handleStepsChange = (newSteps: number) => {
    const clampedSteps = Math.max(MIN_STEPS, Math.min(MAX_STEPS, newSteps));
    if (clampedSteps === steps) return;
    
    setSteps(clampedSteps);
    resizeGrid(clampedSteps);
    sendMessage({ type: 'steps', payload: { steps: clampedSteps } });
  };

  const scheduleNote = useCallback((beatNumber: number, time: number) => {
    const audioContext = audioContextRef.current;
    if (!audioContext || !bassMasterGainRef.current || !drumMasterGainRef.current) return;

    for (let i = 0; i < INSTRUMENTS.length; i++) {
      const instrumentName = INSTRUMENTS[i].name;
      if (grid[i]?.[beatNumber] && !mutes[i]) {
        if (instrumentName === 'Kick') playKick(audioContext, time, drumMasterGainRef.current);
        else if (instrumentName === 'Snare') playSnare(audioContext, time, drumMasterGainRef.current);
        else if (instrumentName === 'Hi-Hat') playHiHat(audioContext, time, drumMasterGainRef.current);
        else if (instrumentName === 'Open Hat') playOpenHat(audioContext, time, drumMasterGainRef.current);
        else if (instrumentName === 'Ride') playRide(audioContext, time, drumMasterGainRef.current);
      }
    }

    const previousBeatNumber = (beatNumber - 1 + steps) % steps;
    const currentNotes = bassLine[beatNumber] || [];
    const previousNotes = bassLine[previousBeatNumber] || [];

    previousNotes.forEach(note => {
      if (!currentNotes.includes(note)) {
        const activeNode = activeOscillatorsRef.current.get(note);
        if (activeNode) {
          const { osc, gain } = activeNode;
          gain.gain.cancelScheduledValues(time);
          gain.gain.setValueAtTime(gain.gain.value, time);
          gain.gain.linearRampToValueAtTime(0, time + 0.02);
          osc.stop(time + 0.02);
          activeOscillatorsRef.current.delete(note);
        }
      }
    });

    currentNotes.forEach(note => {
      if (!previousNotes.includes(note)) {
        const freq = NOTE_FREQ_MAP[note];
        if (freq) {
          const oscillator = audioContext.createOscillator();
          const gainNode = audioContext.createGain();

          oscillator.type = 'sine';
          oscillator.frequency.setValueAtTime(freq, time);
          
          gainNode.connect(bassMasterGainRef.current as GainNode);

          gainNode.gain.setValueAtTime(0, time);
          gainNode.gain.linearRampToValueAtTime(0.3, time + 0.01);

          oscillator.connect(gainNode);
          oscillator.start(time);

          activeOscillatorsRef.current.set(note, { osc: oscillator, gain: gainNode });
        }
      }
    });
  }, [grid, mutes, bassLine, steps]);

  const scheduler = useCallback(() => {
    const audioContext = audioContextRef.current;
    if (!audioContext) return;

    while (nextNoteTimeRef.current < audioContext.currentTime + scheduleAheadTime) {
      scheduleNote(sequenceStepRef.current, nextNoteTimeRef.current);
      setCurrentStep(sequenceStepRef.current);

      const secondsPerBeat = 60.0 / tempo;
      const secondsPerStep = secondsPerBeat / 4;
      nextNoteTimeRef.current += secondsPerStep;
      sequenceStepRef.current = (sequenceStepRef.current + 1) % steps;
    }

    timerRef.current = window.setTimeout(scheduler, lookahead);
  }, [tempo, steps, scheduleNote]);

  const startPlayback = async () => {
    await initAudio();
    const audioContext = audioContextRef.current;
    if (!audioContext || !bassMasterGainRef.current) return;
    
    if (audioContext.state === 'suspended') await audioContext.resume();
    
    const now = audioContext.currentTime;
    bassMasterGainRef.current.gain.cancelScheduledValues(now);
    bassMasterGainRef.current.gain.setTargetAtTime(bassVolume, now, 0.01);
    
    sequenceStepRef.current = 0;
    setCurrentStep(null);
    nextNoteTimeRef.current = now;
    
    setIsPlaying(true);
    sendMessage({ type: 'playback', payload: { isPlaying: true } });
  };

  const stopPlayback = () => {
    setIsPlaying(false);
    sendMessage({ type: 'playback', payload: { isPlaying: false } });
  };

  const exportBeat = () => {
    const beatData: BeatData = { tempo, steps, grid, mutes, bassLine };
    const jsonString = JSON.stringify(beatData, null, 2);
    const blob = new Blob([jsonString], { type: 'application/json' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = 'my-beat.json';
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  };

  const importBeat = (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0];
    if (!file) return;

    const reader = new FileReader();
    reader.onload = (e) => {
      try {
        const text = e.target?.result;
        if (typeof text !== 'string') throw new Error("File is not valid text");
        const data: BeatData = JSON.parse(text);
        
        if(typeof data.tempo !== 'number' || typeof data.steps !== 'number' || !Array.isArray(data.grid)) {
          throw new Error("Invalid beat file format.");
        }

        setTempo(data.tempo);
        setSteps(data.steps);
        setGrid(data.grid);
        setMutes(data.mutes || Array(INSTRUMENTS.length).fill(false));
        setBassLine(data.bassLine || createEmptyBassLine(data.steps));

        sendMessage({ type: 'import', payload: data });

        alert('Beat imported successfully!');
      } catch (error) {
        console.error("Failed to import beat:", error);
        alert('Failed to import beat. The file may be invalid.');
      }
    };
    reader.readAsText(file);
    event.target.value = '';
  };

  const handleDrumVolumeChange = useCallback((newVolume: number) => {
    setDrumVolume(newVolume);
    if (drumMasterGainRef.current && audioContextRef.current) {
        drumMasterGainRef.current.gain.setTargetAtTime(newVolume, audioContextRef.current.currentTime, 0.01);
    }
    sendMessage({ type: 'drumVolume', payload: { drumVolume: newVolume } });
  }, [sendMessage]);

  const handleBassVolumeChange = useCallback((newVolume: number) => {
    setBassVolume(newVolume);
    if (bassMasterGainRef.current && audioContextRef.current) {
        bassMasterGainRef.current.gain.setTargetAtTime(newVolume, audioContextRef.current.currentTime, 0.01);
    }
    sendMessage({ type: 'bassVolume', payload: { bassVolume: newVolume } });
  }, [sendMessage]);

  const triggerSound = useCallback(async (instrumentName: string) => {
    await initAudio();
    const audioContext = audioContextRef.current;
    if (!audioContext || !drumMasterGainRef.current) return;
    const time = audioContext.currentTime;

    if (instrumentName === 'Kick') playKick(audioContext, time, drumMasterGainRef.current);
    else if (instrumentName === 'Snare') playSnare(audioContext, time, drumMasterGainRef.current);
    else if (instrumentName === 'Hi-Hat') playHiHat(audioContext, time, drumMasterGainRef.current);
    else if (instrumentName === 'Open Hat') playOpenHat(audioContext, time, drumMasterGainRef.current);
    else if (instrumentName === 'Ride') playRide(audioContext, time, drumMasterGainRef.current);

  }, [initAudio]);

  const triggerBassNote = useCallback(async (note: string) => {
    await initAudio();
    const audioContext = audioContextRef.current;
    if (!audioContext || !bassMasterGainRef.current) return;
    
    const time = audioContext.currentTime;

    bassMasterGainRef.current.gain.cancelScheduledValues(time);
    bassMasterGainRef.current.gain.setTargetAtTime(bassVolume, time, 0.01);

    const freq = NOTE_FREQ_MAP[note];
    if (freq) {
      const oscillator = audioContext.createOscillator();
      const gainNode = audioContext.createGain();

      oscillator.type = 'sine';
      oscillator.frequency.setValueAtTime(freq, time);
      
      gainNode.connect(bassMasterGainRef.current as GainNode);

      gainNode.gain.setValueAtTime(0, time);
      gainNode.gain.linearRampToValueAtTime(0.3, time + 0.01);
      gainNode.gain.linearRampToValueAtTime(0, time + 1.0);

      oscillator.connect(gainNode);
      oscillator.start(time);
      oscillator.stop(time + 1.0);
    }
  }, [initAudio, bassVolume]);

  const setTempoCallback = useCallback((value: React.SetStateAction<number>) => {
    setTempo(prevTempo => {
        const newTempo = typeof value === 'function' ? value(prevTempo) : value;
        sendMessage({ type: 'tempo', payload: { tempo: newTempo } });
        return newTempo;
    });
  }, [sendMessage]);

  useEffect(() => {
    if (isPlaying) {
      if (timerRef.current) clearTimeout(timerRef.current);
      scheduler();
    } else {
      if (timerRef.current) clearTimeout(timerRef.current);
      stopAllBassNotes();
      setCurrentStep(null);
    }
    
    return () => {
      if (timerRef.current) window.clearTimeout(timerRef.current);
    };
  }, [isPlaying, scheduler, stopAllBassNotes]);

  useEffect(() => {
    if (lastMessage) {
        const { type, payload } = lastMessage;
        switch (type) {
            case 'state':
                if (payload.grid) setGrid(payload.grid);
                if (payload.bassLine) setBassLine(payload.bassLine);
                if (payload.tempo) setTempo(payload.tempo);
                if (payload.isPlaying) setIsPlaying(payload.isPlaying);
                if (payload.mutes) setMutes(payload.mutes);
                if (payload.drumVolume) setDrumVolume(payload.drumVolume);
                if (payload.bassVolume) setBassVolume(payload.bassVolume);
                if (payload.steps) {
                    setSteps(payload.steps);
                    resizeGrid(payload.steps);
                }
                break;
            case 'grid': setGrid(payload.grid); break;
            case 'bassLine': setBassLine(payload.bassLine); break;
            case 'tempo': setTempo(payload.tempo); break;
            case 'playback': setIsPlaying(payload.isPlaying); break;
            case 'mutes': setMutes(payload.mutes); break;
            case 'drumVolume': setDrumVolume(payload.drumVolume); break;
            case 'bassVolume': setBassVolume(payload.bassVolume); break;
            case 'steps': 
                setSteps(payload.steps);
                resizeGrid(payload.steps);
                break;
            case 'clear':
                setGrid(createEmptyGrid(payload.steps));
                setBassLine(createEmptyBassLine(payload.steps));
                break;
            case 'import':
                setTempo(payload.tempo);
                setSteps(payload.steps);
                setGrid(payload.grid);
                setMutes(payload.mutes);
                setBassLine(payload.bassLine);
                break;
        }
    }
  }, [lastMessage, resizeGrid]);

  useEffect(() => {
    return () => {
      stopAllBassNotes();
      if (audioContextRef.current?.state !== 'closed') {
        audioContextRef.current?.close();
      }
    };
  }, [stopAllBassNotes]);

  return {
    steps,
    grid,
    bassLine,
    isPlaying,
    tempo,
    currentStep,
    mutes,
    drumVolume,
    bassVolume,
    setStep,
    setBassNote,
    clearPattern,
    setTempo: setTempoCallback,
    handleStepsChange,
    startPlayback,
    stopPlayback,
    exportBeat,
    importBeat,
    toggleMute,
    handleDrumVolumeChange,
    handleBassVolumeChange,
    triggerSound,
    triggerBassNote,
  };
};