import AVFoundation
import Metal

public final class AudioFeatureExtractor {
    public let sampleRate: Int
    public let nMels: Int
    public let nFft: Int
    public let hopLength: Int
    public let fMin: Float
    public let fMax: Float
    
    public init(
        sampleRate: Int = 16000,
        nMels: Int = 128,
        nFft: Int = 400,
        hopLength: Int = 160,
        fMin: Float = 0,
        fMax: Float = 8000
    ) {
        self.sampleRate = sampleRate
        self.nMels = nMels
        self.nFft = nFft
        self.hopLength = hopLength
        self.fMin = fMin
        self.fMax = fMax
    }
    
    public func extractMelSpectrogram(from audioData: [Float]) -> [[Float]] {
        let numFrames = (audioData.count - nFft) / hopLength + 1
        var melSpec = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames)
        
        for frameIdx in 0..<numFrames {
            let startIdx = frameIdx * hopLength
            let endIdx = min(startIdx + nFft, audioData.count)
            
            // Zero-pad if frame is shorter than nFft
            var frame = [Float](repeating: 0, count: nFft)
            for i in startIdx..<endIdx {
                frame[i - startIdx] = audioData[i]
            }
            
            let windowedFrame = applyHannWindow(frame)
            let spectrum = computeSpectrum(windowedFrame)
            let melEnergies = computeMelEnergies(spectrum)
            
            melSpec[frameIdx] = melEnergies
        }
        
        return melSpec
    }
    
    private func applyHannWindow(_ frame: [Float]) -> [Float] {
        let n = frame.count
        return frame.enumerated().map { i, val in
            val * 0.5 * (1.0 - cos(2.0 * Float.pi * Float(i) / Float(n - 1)))
        }
    }
    
    private func computeSpectrum(_ frame: [Float]) -> [Float] {
        let n = frame.count
        var spectrum = [Float](repeating: 0, count: n / 2 + 1)
        
        for k in 0..<spectrum.count {
            var real: Float = 0
            var imag: Float = 0
            for i in 0..<n {
                let angle = -2.0 * Float.pi * Float(k) * Float(i) / Float(n)
                real += frame[i] * cos(angle)
                imag += frame[i] * sin(angle)
            }
            spectrum[k] = sqrt(real * real + imag * imag)
        }
        
        return spectrum
    }
    
    private func computeMelEnergies(_ spectrum: [Float]) -> [Float] {
        var melEnergies = [Float](repeating: 0, count: nMels)
        
        let melPoints = createMelFilterbank()
        
        for melIdx in 0..<nMels {
            var sum: Float = 0
            for fftIdx in 0..<spectrum.count {
                sum += spectrum[fftIdx] * melPoints[melIdx][fftIdx]
            }
            melEnergies[melIdx] = log10(max(sum, 1e-10))
        }
        
        return melEnergies
    }
    
    private func createMelFilterbank() -> [[Float]] {
        var filterbank = [[Float]](repeating: [Float](repeating: 0, count: nFft / 2 + 1), count: nMels)
        
        let melMin = hzToMel(fMin)
        let melMax = hzToMel(fMax)
        
        let melPoints = (0..<nMels + 2).map { i in
            melMin + Float(i) * (melMax - melMin) / Float(nMels + 1)
        }
        
        let hzPoints = melPoints.map { melToHz($0) }
        let binPoints = hzPoints.map { Int(round($0 * Float(nFft) / Float(sampleRate))) }
        
        for melIdx in 0..<nMels {
            let leftBin = binPoints[melIdx]
            let centerBin = binPoints[melIdx + 1]
            let rightBin = binPoints[melIdx + 2]
            
            for bin in leftBin..<centerBin {
                if bin < filterbank[melIdx].count {
                    filterbank[melIdx][bin] = Float(bin - leftBin) / Float(centerBin - leftBin)
                }
            }
            for bin in centerBin..<rightBin {
                if bin < filterbank[melIdx].count {
                    filterbank[melIdx][bin] = Float(rightBin - bin) / Float(rightBin - centerBin)
                }
            }
        }
        
        return filterbank
    }
    
    private func hzToMel(_ hz: Float) -> Float {
        2595.0 * log10(1.0 + hz / 700.0)
    }
    
    private func melToHz(_ mel: Float) -> Float {
        700.0 * (pow(10.0, mel / 2595.0) - 1.0)
    }
    
    // ── GPU-accelerated mel spectrogram ──
    
    public func extractMelSpectrogramGPU(
        engine: MarkBaseEngine,
        audioData: [Float]
    ) throws -> [[Float]] {
        let device = engine.device
        let spectrumSize = nFft / 2 + 1
        let numFrames = (audioData.count - nFft) / hopLength + 1
        let melBufferSize = numFrames * nMels

        let filterbank2D = createMelFilterbank()
        var flatFilterbank = [Float](repeating: 0, count: nMels * spectrumSize)
        for m in 0..<nMels {
            for b in 0..<spectrumSize {
                flatFilterbank[m * spectrumSize + b] = filterbank2D[m][b]
            }
        }

        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        let audioBuf = device.makeBuffer(bytes: audioData, length: audioData.count * 4)!
        let filterbankBuf = device.makeBuffer(bytes: flatFilterbank, length: flatFilterbank.count * 4)!
        let spectrumBuf = device.makeBuffer(length: numFrames * spectrumSize * 4)!
        let melBuf = device.makeBuffer(length: melBufferSize * 4)!

        let psoDFT = try engine.pipeline(named: "audio_dft_magnitude")
        let psoMel = try engine.pipeline(named: "audio_mel_filterbank")

        do {
            let enc = cmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(psoDFT)
            enc.setBuffer(audioBuf, offset: 0, index: 0)
            enc.setBuffer(spectrumBuf, offset: 0, index: 1)
            var n = UInt32(nFft); enc.setBytes(&n, length: 4, index: 2)
            var h = UInt32(hopLength); enc.setBytes(&h, length: 4, index: 3)
            var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 4)
            var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 5)
            var al = UInt32(audioData.count); enc.setBytes(&al, length: 4, index: 6)
            let grid = MTLSize(width: numFrames, height: spectrumSize, depth: 1)
            let tg = MTLSize(width: 8, height: 8, depth: 1)
            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
            enc.endEncoding()
        }
        do {
            let enc = cmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(psoMel)
            enc.setBuffer(spectrumBuf, offset: 0, index: 0)
            enc.setBuffer(filterbankBuf, offset: 0, index: 1)
            enc.setBuffer(melBuf, offset: 0, index: 2)
            var ss = UInt32(spectrumSize); enc.setBytes(&ss, length: 4, index: 3)
            var nm = UInt32(nMels); enc.setBytes(&nm, length: 4, index: 4)
            var nf = UInt32(numFrames); enc.setBytes(&nf, length: 4, index: 5)
            let grid = MTLSize(width: numFrames, height: nMels, depth: 1)
            let tg = MTLSize(width: 8, height: 8, depth: 1)
            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
            enc.endEncoding()
        }

        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()

        let ptr = melBuf.contents().assumingMemoryBound(to: Float.self)
        let flat = Array(UnsafeBufferPointer(start: ptr, count: melBufferSize))
        var result = [[Float]](repeating: [Float](repeating: 0, count: nMels), count: numFrames)
        for f in 0..<numFrames {
            for m in 0..<nMels {
                result[f][m] = flat[f * nMels + m]
            }
        }
        return result
    }

    public func loadAudioFile(url: URL) throws -> [Float] {
        let asset = AVURLAsset(url: url)
        let reader = try AVAssetReader(asset: asset)
        
        let output = AVAssetReaderAudioMixOutput(audioTracks: asset.tracks, audioSettings: nil)
        reader.add(output)
        reader.startReading()
        
        var samples: [Float] = []
        while reader.status == .reading {
            let buffer = output.copyNextSampleBuffer()
            if let buffer = buffer {
                let blockBuffer = CMSampleBufferGetDataBuffer(buffer)
                if let blockBuffer = blockBuffer {
                    let length = CMBlockBufferGetDataLength(blockBuffer)
                    var data = [Float](repeating: 0, count: length / MemoryLayout<Float>.stride)
                    CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data)
                    samples.append(contentsOf: data)
                }
            }
        }
        
        return samples
    }
}