import AudioToolbox
import CoreMedia
import Foundation
import MastermindPOCCore

enum SampleBufferPCMExtractor {
    static func extractFrame(from sampleBuffer: CMSampleBuffer, source: AudioSource) -> PCMFrame? {
        guard let formatDescription = CMSampleBufferGetFormatDescription(sampleBuffer),
              let streamDescription = CMAudioFormatDescriptionGetStreamBasicDescription(formatDescription)
        else {
            return nil
        }

        let asbd = streamDescription.pointee
        guard asbd.mFormatID == kAudioFormatLinearPCM else {
            return nil
        }

        var bufferListSize = 0
        var blockBuffer: CMBlockBuffer?
        var status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(
            sampleBuffer,
            bufferListSizeNeededOut: &bufferListSize,
            bufferListOut: nil,
            bufferListSize: 0,
            blockBufferAllocator: kCFAllocatorDefault,
            blockBufferMemoryAllocator: kCFAllocatorDefault,
            flags: kCMSampleBufferFlag_AudioBufferList_Assure16ByteAlignment,
            blockBufferOut: &blockBuffer
        )

        guard status == noErr, bufferListSize > 0 else {
            return nil
        }

        let bufferListPointer = UnsafeMutableRawPointer.allocate(byteCount: bufferListSize, alignment: MemoryLayout<AudioBufferList>.alignment)
        defer {
            bufferListPointer.deallocate()
        }

        let audioBufferList = bufferListPointer.bindMemory(to: AudioBufferList.self, capacity: 1)
        status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(
            sampleBuffer,
            bufferListSizeNeededOut: nil,
            bufferListOut: audioBufferList,
            bufferListSize: bufferListSize,
            blockBufferAllocator: kCFAllocatorDefault,
            blockBufferMemoryAllocator: kCFAllocatorDefault,
            flags: kCMSampleBufferFlag_AudioBufferList_Assure16ByteAlignment,
            blockBufferOut: &blockBuffer
        )

        guard status == noErr else {
            return nil
        }

        let buffers = UnsafeMutableAudioBufferListPointer(audioBufferList)
        guard let firstBuffer = buffers.first,
              let mData = firstBuffer.mData
        else {
            return nil
        }

        let sourceRate = Int(asbd.mSampleRate.rounded())
        let byteCount = Int(firstBuffer.mDataByteSize)
        let rawPointer = UnsafeRawPointer(mData)

        let samples: [Float]
        if asbd.mBitsPerChannel == 32, asbd.mFormatFlags & kAudioFormatFlagIsFloat != 0 {
            let sampleCount = byteCount / MemoryLayout<Float>.size
            let pointer = rawPointer.bindMemory(to: Float.self, capacity: sampleCount)
            samples = (0..<sampleCount).map { pointer[$0] }
        } else if asbd.mBitsPerChannel == 16, asbd.mFormatFlags & kAudioFormatFlagIsSignedInteger != 0 {
            let sampleCount = byteCount / MemoryLayout<Int16>.size
            let pointer = rawPointer.bindMemory(to: Int16.self, capacity: sampleCount)
            samples = (0..<sampleCount).map { index in
                Float(Int16(littleEndian: pointer[index])) / Float(Int16.max)
            }
        } else {
            return nil
        }

        let mono16k = PCM16LE.resampleLinear(samples: samples, sourceRate: sourceRate, targetRate: 16_000)
        let data = PCM16LE.encode(samples: mono16k)

        return PCMFrame(source: source, sampleRate: 16_000, channels: 1, pcmS16LE: data)
    }
}
