//
//  MicrophonePCMSampleVendorAT.swift
//  SwiftOpenAI
//
//  Created from AIProxySwift
//  Original: https://github.com/lzell/AIProxySwift
//

#if canImport(AVFoundation)
#if os(macOS) || os(iOS)

import AudioToolbox
@preconcurrency import AVFoundation
import Foundation
import OSLog

private let logger = Logger(subsystem: "com.swiftopenai", category: "Audio")

nonisolated private let kVoiceProcessingInputSampleRate: Double = 44100

/// This is an AudioToolbox-based implementation that vends PCM16 microphone samples at a
/// sample rate that OpenAI's realtime models expect.
///
/// ## Requirements
///
/// - Assumes an `NSMicrophoneUsageDescription` description has been added to Target > Info
/// - Assumes that microphone permissions have already been granted
///
/// ## Usage
///
///     ```
///     let microphoneVendor = MicrophonePCMSampleVendorAT()
///     let micStream = try microphoneVendor.start()
///     Task {
///         for await buffer in micStream {
///             // Use buffer
///         }
///     }
///     // ... some time later ...
///     microphoneVendor.stop()
///     ```
@RealtimeActor
class MicrophonePCMSampleVendorAT: MicrophonePCMSampleVendor {

  public init() { }

  deinit {
    logger.debug("MicrophonePCMSampleVendorAT is being freed")
  }

  public func start() throws -> AsyncStream<AVAudioPCMBuffer> {
    var desc = AudioComponentDescription(
      componentType: kAudioUnitType_Output,
      componentSubType: kAudioUnitSubType_VoiceProcessingIO,
      componentManufacturer: kAudioUnitManufacturer_Apple,
      componentFlags: 0,
      componentFlagsMask: 0)

    guard let component = AudioComponentFindNext(nil, &desc) else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not find an audio component with VoiceProcessingIO")
    }

    AudioComponentInstanceNew(component, &audioUnit)
    guard let audioUnit else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not instantiate an audio component with VoiceProcessingIO")
    }

    var one: UInt32 = 1
    var err = AudioUnitSetProperty(
      audioUnit,
      kAudioOutputUnitProperty_EnableIO,
      kAudioUnitScope_Input,
      1,
      &one,
      UInt32(MemoryLayout.size(ofValue: one)))

    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not enable the input scope of the microphone bus")
    }

    var zero: UInt32 = 0
    err = AudioUnitSetProperty(
      audioUnit,
      kAudioOutputUnitProperty_EnableIO,
      kAudioUnitScope_Output,
      0,
      &zero,
      UInt32(MemoryLayout.size(ofValue: one)))

    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not disable the output scope of the speaker bus")
    }

    var hardwareASBD = AudioStreamBasicDescription()
    var size = UInt32(MemoryLayout<AudioStreamBasicDescription>.size)
    let _ = AudioUnitGetProperty(
      audioUnit,
      kAudioUnitProperty_StreamFormat,
      kAudioUnitScope_Input,
      1,
      &hardwareASBD,
      &size)
    logger.debug("Hardware mic is natively at \(hardwareASBD.mSampleRate) sample rate")

    var ioFormat = AudioStreamBasicDescription(
      mSampleRate: kVoiceProcessingInputSampleRate,
      mFormatID: kAudioFormatLinearPCM,
      mFormatFlags: kAudioFormatFlagIsSignedInteger | kAudioFormatFlagIsPacked,
      mBytesPerPacket: 2,
      mFramesPerPacket: 1,
      mBytesPerFrame: 2,
      mChannelsPerFrame: 1,
      mBitsPerChannel: 16,
      mReserved: 0)

    err = AudioUnitSetProperty(
      audioUnit,
      kAudioUnitProperty_StreamFormat,
      kAudioUnitScope_Output,
      1,
      &ioFormat,
      UInt32(MemoryLayout<AudioStreamBasicDescription>.size))
    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not set ASBD on the output scope of the mic bus")
    }

    #if os(macOS)
    if let deviceID = AudioUtils.getDefaultAudioInputDevice() {
      // Try to get 50ms updates.
      // 50ms is half the granularity of our target accumulator (we accumulate into 100ms payloads that we send up to OpenAI)
      var bufferSize = UInt32(kVoiceProcessingInputSampleRate / 20)
      var propertyAddress = AudioObjectPropertyAddress(
        mSelector: kAudioDevicePropertyBufferFrameSize,
        mScope: kAudioDevicePropertyScopeInput,
        mElement: kAudioObjectPropertyElementMain)

      let size = UInt32(MemoryLayout.size(ofValue: bufferSize))
      let status = AudioObjectSetPropertyData(
        deviceID,
        &propertyAddress,
        0,
        nil,
        size,
        &bufferSize)

      if status != noErr {
        logger.debug("Could not set desired buffer size")
      }
    }
    #endif

    var inputCallbackStruct = AURenderCallbackStruct(
      inputProc: audioRenderCallback,
      inputProcRefCon: Unmanaged.passUnretained(self).toOpaque())
    err = AudioUnitSetProperty(
      audioUnit,
      kAudioOutputUnitProperty_SetInputCallback,
      kAudioUnitScope_Global,
      1,
      &inputCallbackStruct,
      UInt32(MemoryLayout<AURenderCallbackStruct>.size))

    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not set the render callback on the voice processing audio unit")
    }

    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not configure auto gain control")
    }

    err = AudioUnitInitialize(audioUnit)
    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not initialize the audio unit")
    }

    err = AudioOutputUnitStart(audioUnit)
    guard err == noErr else {
      throw MicrophonePCMSampleVendorError.couldNotConfigureAudioUnit(
        "Could not start the audio unit")
    }

    return AsyncStream<AVAudioPCMBuffer> { [weak self] continuation in
      self?.continuation = continuation
    }
  }

  public func stop() {
    continuation?.finish()
    continuation = nil

    // Capture the audio unit reference
    if let au = audioUnit {
      // Clear the property immediately on RealtimeActor
      audioUnit = nil

      // Dispatch the actual AudioToolbox cleanup to a background queue
      // to avoid priority inversion with the audio render thread.
      // Note: AudioUnit is a C pointer type (non-Sendable), but the AudioToolbox
      // cleanup APIs are thread-safe and we've already cleared our reference.
      DispatchQueue.global(qos: .utility).async {
        AudioOutputUnitStop(au)
        AudioUnitUninitialize(au)
        AudioComponentInstanceDispose(au)
      }
    }

    microphonePCMSampleVendorCommon.audioConverter = nil
  }

  fileprivate func didReceiveRenderCallback(
    _ ioActionFlags: UnsafeMutablePointer<AudioUnitRenderActionFlags>,
    _ inTimeStamp: UnsafePointer<AudioTimeStamp>,
    _ inBusNumber: UInt32,
    _ inNumberFrames: UInt32)
  {
    guard let audioUnit else {
      logger.error("There is no audioUnit attached to the sample vendor. Render callback should not be called")
      return
    }
    var bufferList = AudioBufferList(
      mNumberBuffers: 1,
      mBuffers: AudioBuffer(
        mNumberChannels: 1,
        mDataByteSize: inNumberFrames * 2,
        mData: UnsafeMutableRawPointer.allocate(
          byteCount: Int(inNumberFrames) * 2,
          alignment: MemoryLayout<Int16>.alignment)))

    let status = AudioUnitRender(
      audioUnit,
      ioActionFlags,
      inTimeStamp,
      inBusNumber,
      inNumberFrames,
      &bufferList)

    guard status == noErr else {
      logger.error("Could not render voice processed audio data to bufferList")
      return
    }

    guard
      let audioFormat = AVAudioFormat(
        commonFormat: .pcmFormatInt16,
        sampleRate: kVoiceProcessingInputSampleRate,
        channels: 1,
        interleaved: true)
    else {
      logger.error("Could not create audio format inside render callback.")
      return
    }

    if
      let sampleBuffer = AVAudioPCMBuffer(pcmFormat: audioFormat, bufferListNoCopy: &bufferList),
      let accumulatedBuffer = microphonePCMSampleVendorCommon.resampleAndAccumulate(sampleBuffer)
    {
      // If the buffer has accumulated to a sufficient level, give it back to the caller
      Task { @RealtimeActor in
        self.continuation?.yield(accumulatedBuffer)
      }
    }
  }

  private var audioUnit: AudioUnit?
  private let microphonePCMSampleVendorCommon = MicrophonePCMSampleVendorCommon()
  private var continuation: AsyncStream<AVAudioPCMBuffer>.Continuation?

}

/// This @RealtimeActor annotation is a lie.
@RealtimeActor private let audioRenderCallback: AURenderCallback = {
  inRefCon,
    ioActionFlags,
    inTimeStamp,
    inBusNumber,
    inNumberFrames,
    _ in
  let microphonePCMSampleVendor = Unmanaged<MicrophonePCMSampleVendorAT>
    .fromOpaque(inRefCon)
    .takeUnretainedValue()
  microphonePCMSampleVendor.didReceiveRenderCallback(
    ioActionFlags,
    inTimeStamp,
    inBusNumber,
    inNumberFrames)
  return noErr
}
#endif
#endif
