import AudioToolbox
import CoreAudio

/// Captures up to two channels of an input device (e.g. headset mic) into an `AudioFanout`.
final class InputEngine: @unchecked Sendable {
    let device: AudioDevice
    let fanout = AudioFanout()
    private let sampleRate: Double
    private let channels: Int
    private var unit: AudioUnit?
    private let maxFrames = 4096
    private let samples: UnsafeMutablePointer<Float>
    private let bufferList: UnsafeMutableAudioBufferListPointer

    init(device: AudioDevice, bufferFrames: UInt32 = 128) throws {
        self.device = device
        self.sampleRate = device.sampleRate
        self.channels = min(max(device.inputChannels, 1), 2)
        samples = .allocate(capacity: maxFrames * 2)
        samples.initialize(repeating: 0, count: maxFrames * 2)
        bufferList = AudioBufferList.allocate(maximumBuffers: channels)

        var description = AudioComponentDescription(
            componentType: kAudioUnitType_Output, componentSubType: kAudioUnitSubType_HALOutput,
            componentManufacturer: kAudioUnitManufacturer_Apple, componentFlags: 0, componentFlagsMask: 0)
        guard let component = AudioComponentFindNext(nil, &description) else { throw AudioError.noHAL }
        var newUnit: AudioUnit?
        try check(AudioComponentInstanceNew(component, &newUnit), "instance")
        guard let newUnit else { throw AudioError.noHAL }
        unit = newUnit

        var enable: UInt32 = 1, disable: UInt32 = 0
        let size = UInt32(MemoryLayout<UInt32>.size)
        try check(AudioUnitSetProperty(newUnit, kAudioOutputUnitProperty_EnableIO, kAudioUnitScope_Input, 1, &enable, size), "enable input")
        try check(AudioUnitSetProperty(newUnit, kAudioOutputUnitProperty_EnableIO, kAudioUnitScope_Output, 0, &disable, size), "disable output")

        var deviceID = device.deviceID
        try check(AudioUnitSetProperty(newUnit, kAudioOutputUnitProperty_CurrentDevice, kAudioUnitScope_Global, 0,
                                       &deviceID, UInt32(MemoryLayout<AudioDeviceID>.size)), "device")

        var format = AudioStreamBasicDescription.floatNonInterleaved(sampleRate: sampleRate, channels: channels)
        try check(AudioUnitSetProperty(newUnit, kAudioUnitProperty_StreamFormat, kAudioUnitScope_Output, 1,
                                       &format, UInt32(MemoryLayout<AudioStreamBasicDescription>.size)), "format")

        var frames = bufferFrames
        AudioUnitSetProperty(newUnit, kAudioDevicePropertyBufferFrameSize, kAudioUnitScope_Global, 0, &frames, size)

        var callback = AURenderCallbackStruct(inputProc: inputCaptured, inputProcRefCon: Unmanaged.passUnretained(self).toOpaque())
        try check(AudioUnitSetProperty(newUnit, kAudioOutputUnitProperty_SetInputCallback, kAudioUnitScope_Global, 0,
                                       &callback, UInt32(MemoryLayout<AURenderCallbackStruct>.size)), "callback")
        try check(AudioUnitInitialize(newUnit), "initialize")
        try check(AudioOutputUnitStart(newUnit), "start")
    }

    deinit {
        stop()
        samples.deallocate()
        free(bufferList.unsafeMutablePointer)
    }

    func stop() {
        guard let unit else { return }
        AudioOutputUnitStop(unit)
        AudioUnitUninitialize(unit)
        AudioComponentInstanceDispose(unit)
        self.unit = nil
    }

    fileprivate func capture(flags: UnsafeMutablePointer<AudioUnitRenderActionFlags>, timestamp: UnsafePointer<AudioTimeStamp>, frames: UInt32) -> OSStatus {
        guard let unit, Int(frames) <= maxFrames else { return noErr }
        for channel in 0..<channels {
            bufferList[channel] = AudioBuffer(mNumberChannels: 1, mDataByteSize: frames * 4,
                                              mData: UnsafeMutableRawPointer(samples + channel * maxFrames))
        }
        let status = AudioUnitRender(unit, flags, timestamp, 1, frames, bufferList.unsafeMutablePointer)
        guard status == noErr else { return status }
        fanout.push(planar: samples, channelStride: maxFrames, channels: channels, frames: Int(frames), sampleRate: sampleRate)
        return noErr
    }
}

private func inputCaptured(
    refCon: UnsafeMutableRawPointer,
    flags: UnsafeMutablePointer<AudioUnitRenderActionFlags>,
    timestamp: UnsafePointer<AudioTimeStamp>,
    bus: UInt32,
    frames: UInt32,
    data: UnsafeMutablePointer<AudioBufferList>?
) -> OSStatus {
    Unmanaged<InputEngine>.fromOpaque(refCon).takeUnretainedValue().capture(flags: flags, timestamp: timestamp, frames: frames)
}