audiotee/Sources/AudioTeeCore/Core/AudioFormatConverter.swift

import AVFoundation
import CoreAudio
import Foundation

/// Audio format converter using AVFoundation's AVAudioConverter.
///
/// Pre-allocates input/output buffers on first use and reuses them across
/// transform() calls. This eliminates two AVAudioPCMBuffer heap allocations
/// per chunk — significant when chunks are small (50ms = 20 calls/sec).
public class AudioFormatConverter {
  private let avConverter: AVAudioConverter
  private let sourceFormat: AVAudioFormat
  private let targetFormat: AVAudioFormat

  /// Pre-allocated buffers reused across transform() calls. Lazily created
  /// on first transform() since we need the actual input frame count to
  /// size them correctly.
  private var cachedInputBuffer: AVAudioPCMBuffer?
  private var cachedOutputBuffer: AVAudioPCMBuffer?

  public init(sourceFormat: AudioStreamBasicDescription, targetFormat: AudioStreamBasicDescription)
    throws
  {
    var mutableSourceFormat = sourceFormat
    var mutableTargetFormat = targetFormat

    guard let sourceAVFormat = AVAudioFormat(streamDescription: &mutableSourceFormat),
      let targetAVFormat = AVAudioFormat(streamDescription: &mutableTargetFormat)
    else {
      throw AudioConverterError.invalidFormat
    }

    guard let converter = AVAudioConverter(from: sourceAVFormat, to: targetAVFormat) else {
      throw AudioConverterError.creationFailed
    }

    self.sourceFormat = sourceAVFormat
    self.targetFormat = targetAVFormat
    self.avConverter = converter

    AudioTeeLogging.logger.debug(
      "Audio converter created",
      context: [
        "source_sample_rate": String(sourceAVFormat.sampleRate),
        "target_sample_rate": String(targetAVFormat.sampleRate),
        "source_channels": String(sourceAVFormat.channelCount),
        "target_channels": String(targetAVFormat.channelCount),
      ])

    // Warn about upsampling once during initialization
    if targetAVFormat.sampleRate > sourceAVFormat.sampleRate {
      AudioTeeLogging.logger.info(
        "Upsampling audio - this doesn't add frequency content above the original Nyquist limit",
        context: [
          "source_rate": String(sourceAVFormat.sampleRate),
          "target_rate": String(targetAVFormat.sampleRate),
        ])
    }
  }

  /// The source format this converter reads from.
  public var sourceFormatDescription: AudioStreamBasicDescription {
    return sourceFormat.streamDescription.pointee
  }

  /// The target format this converter produces.
  public var targetFormatDescription: AudioStreamBasicDescription {
    return targetFormat.streamDescription.pointee
  }

  /// Returns pre-allocated input and output buffers sized for the given
  /// input frame count. Allocates once on first call; reuses on subsequent
  /// calls when capacity is sufficient. Re-allocates if a larger frame
  /// count arrives (shouldn't happen with fixed chunk sizes, but handled
  /// gracefully).
  private func getBuffers(inputFrameCount: AVAudioFrameCount)
    -> (input: AVAudioPCMBuffer, output: AVAudioPCMBuffer)?
  {
    // ceil() prevents float-to-int truncation from undersizing the buffer
    // by one frame (e.g. 3199.9999 → 3199 instead of 3200).
    let outputFrameCount = AVAudioFrameCount(
      ceil(Double(inputFrameCount) * (targetFormat.sampleRate / sourceFormat.sampleRate))
    )

    // Reuse cached buffers if they have sufficient capacity
    if let inputBuf = cachedInputBuffer,
      let outputBuf = cachedOutputBuffer,
      inputBuf.frameCapacity >= inputFrameCount,
      outputBuf.frameCapacity >= outputFrameCount
    {
      // Reset frame lengths for reuse — the underlying memory is retained,
      // we just tell AVAudioPCMBuffer how many frames are valid this time.
      inputBuf.frameLength = 0
      outputBuf.frameLength = 0
      return (inputBuf, outputBuf)
    }

    // Allocate new buffers (first call, or unexpected capacity increase)
    guard
      let inputBuf = AVAudioPCMBuffer(
        pcmFormat: sourceFormat, frameCapacity: inputFrameCount)
    else {
      AudioTeeLogging.logger.error("Failed to create input buffer")
      return nil
    }

    guard
      let outputBuf = AVAudioPCMBuffer(
        pcmFormat: targetFormat, frameCapacity: outputFrameCount)
    else {
      AudioTeeLogging.logger.error("Failed to create output buffer")
      return nil
    }

    // Cache for reuse on subsequent calls
    cachedInputBuffer = inputBuf
    cachedOutputBuffer = outputBuf

    AudioTeeLogging.logger.debug(
      "Allocated converter buffers",
      context: [
        "input_frame_capacity": String(inputFrameCount),
        "output_frame_capacity": String(outputFrameCount),
      ])

    return (inputBuf, outputBuf)
  }

  public func transform(_ packet: AudioPacket) -> AudioPacket {
    let inputData = packet.data

    // Calculate frame count from the input data size
    let bytesPerFrame = Int(sourceFormat.streamDescription.pointee.mBytesPerFrame)
    let inputFrameCount = AVAudioFrameCount(inputData.count / bytesPerFrame)

    // Get or create pre-allocated buffers
    guard let (inputBuffer, outputBuffer) = getBuffers(inputFrameCount: inputFrameCount) else {
      return packet
    }

    // Copy input data into the reusable input buffer
    inputData.withUnsafeBytes { bytes in
      let dest = inputBuffer.audioBufferList.pointee.mBuffers.mData!
      dest.copyMemory(from: bytes.baseAddress!, byteCount: inputData.count)
    }
    inputBuffer.frameLength = inputFrameCount

    // Perform conversion — the block-based API lets AVAudioConverter pull
    // input data as needed. We do NOT call avConverter.reset() between
    // calls because the resampler maintains internal state for continuity
    // across chunks (avoiding discontinuity artifacts).
    var error: NSError?

    let status = avConverter.convert(to: outputBuffer, error: &error) {
      requestedPackets, outStatus in
      outStatus.pointee = .haveData
      return inputBuffer
    }

    // Check if conversion produced output (regardless of status code)
    guard outputBuffer.frameLength > 0 else {
      AudioTeeLogging.logger.error(
        "Audio conversion produced no output",
        context: [
          "status": String(describing: status),
          "error": String(describing: error),
          "input_frames": String(inputBuffer.frameLength),
          "output_capacity": String(outputBuffer.frameCapacity),
        ])
      return packet
    }

    // Extract converted data from the reusable output buffer
    let outputData = Data(
      bytes: outputBuffer.audioBufferList.pointee.mBuffers.mData!,
      count: Int(outputBuffer.frameLength * targetFormat.streamDescription.pointee.mBytesPerFrame))

    return AudioPacket(
      timestamp: packet.timestamp,
      duration: packet.duration,
      data: outputData
    )
  }

  public static func toSampleRate(
    _ sampleRate: Double, from sourceFormat: AudioStreamBasicDescription
  ) throws -> AudioFormatConverter {
    var targetFormat = AudioStreamBasicDescription()
    targetFormat.mSampleRate = sampleRate
    targetFormat.mFormatID = kAudioFormatLinearPCM
    targetFormat.mFormatFlags = kAudioFormatFlagIsPacked | kAudioFormatFlagIsSignedInteger
    targetFormat.mFramesPerPacket = 1
    targetFormat.mBitsPerChannel = 16
    targetFormat.mChannelsPerFrame = sourceFormat.mChannelsPerFrame
    targetFormat.mBytesPerFrame =
      (targetFormat.mBitsPerChannel / 8) * sourceFormat.mChannelsPerFrame
    targetFormat.mBytesPerPacket = targetFormat.mFramesPerPacket * targetFormat.mBytesPerFrame

    return try AudioFormatConverter(sourceFormat: sourceFormat, targetFormat: targetFormat)
  }

  public static func isValidSampleRate(_ sampleRate: Double) -> Bool {
    return [8000, 16000, 22050, 24000, 32000, 44100, 48000].contains(sampleRate)
  }
}