SystemAudioCaptureService.swift

@preconcurrency import AVFoundation
@preconcurrency import CoreAudio
import Foundation

/// Service for capturing system audio using Core Audio Taps (macOS 14.4+)
/// Captures all system audio output and converts to 16-bit PCM at 16kHz for transcription
@available(macOS 14.4, *)
class SystemAudioCaptureService: @unchecked Sendable {

  // MARK: - Types

  /// Callback for receiving audio chunks
  typealias AudioChunkHandler = @Sendable (Data) -> Void

  /// Callback for receiving audio levels (0.0 - 1.0)
  typealias AudioLevelHandler = @Sendable (Float) -> Void

  enum SystemAudioCaptureError: LocalizedError {
    case tapCreationFailed(OSStatus)
    case aggregateDeviceFailed(OSStatus)
    case ioProcCreationFailed(OSStatus)
    case deviceStartFailed(OSStatus)
    case formatError
    case converterCreationFailed
    case unsupportedOS

    var errorDescription: String? {
      switch self {
      case .tapCreationFailed(let status):
        return "Failed to create process tap: \(status)"
      case .aggregateDeviceFailed(let status):
        return "Failed to create aggregate device: \(status)"
      case .ioProcCreationFailed(let status):
        return "Failed to create IO proc: \(status)"
      case .deviceStartFailed(let status):
        return "Failed to start audio device: \(status)"
      case .formatError:
        return "Failed to get audio format"
      case .converterCreationFailed:
        return "Failed to create audio converter"
      case .unsupportedOS:
        return "System audio capture requires macOS 14.4 or later"
      }
    }
  }

  // MARK: - Properties

  private var tapID: AudioObjectID = kAudioObjectUnknown
  private var aggregateDeviceID: AudioObjectID = kAudioObjectUnknown
  private var ioProcID: AudioDeviceIOProcID?
  private var isCapturing = false
  private var onAudioChunk: AudioChunkHandler?
  private var onAudioLevel: AudioLevelHandler?

  /// Target sample rate for DeepGram
  private let targetSampleRate: Double = 16000

  // Resampling
  private var audioConverter: AVAudioConverter?
  private var inputFormat: AVAudioFormat?
  private var targetFormat: AVAudioFormat?
  private var sourceSampleRate: Double = 0.0
  // Flipped exactly once per convert() call by the synchronous AVAudioConverter input
  // block (a @Sendable closure) on the non-reentrant CoreAudio IO thread. nonisolated(unsafe)
  // because the IOProc is serial; it is reset at the top of every handleAudioInput() call.
  private nonisolated(unsafe) var hasConsumedInput = false

  // Tap UUID for identification
  private let tapUUID = UUID()

  /// Dedicated queue for CoreAudio device operations (start/stop)
  /// to avoid blocking the main thread on AudioDeviceStart/Stop calls.
  private let audioQueue = DispatchQueue(label: "com.omi.systemaudiocapture.device")

  // MARK: - Permission Checking

  /// Check if system audio capture permission is available
  /// Note: Core Audio Taps don't have a preflight API like screen capture.
  /// Permission is granted implicitly on first use, or may require entitlements.
  static func checkPermission() -> Bool {
    // For Core Audio Taps, there's no explicit permission API.
    // The system will prompt when we first try to create a tap.
    // Return true to indicate we can attempt capture.
    return true
  }

  /// Request system audio capture permission
  /// Returns true if permission is available (macOS 14.4+)
  static func requestPermission() async -> Bool {
    // Core Audio Taps permission is handled at capture time
    return true
  }

  /// Prime the Core Audio process-tap consent DURING onboarding.
  ///
  /// The system-audio capture path uses a global process tap
  /// (`CATapDescription` → `AudioHardwareCreateProcessTap`, see
  /// `startCaptureOnQueue`) whose consent — "<app> is requesting to bypass the
  /// system private window picker and directly access your screen and audio" —
  /// is SEPARATE from the Screen-Recording TCC grant onboarding already
  /// requests, and only fires the first time a tap is actually created. Because
  /// onboarding never creates a tap, that consent would otherwise be deferred
  /// and fire AGAIN on the first real capture, post-onboarding.
  ///
  /// This creates the SAME tap + aggregate device the real capture path builds
  /// (mirroring `startCaptureOnQueue`), so macOS shows and persists the exact
  /// same consent, then IMMEDIATELY and FULLY tears everything down (mirroring
  /// `cleanupTap()`/`cleanup()`). It never creates an IO proc or starts real
  /// capture. Idempotent and safe to call when permission is already granted
  /// (no crash, no lingering device), and safe to call on a background queue.
  /// Returns quickly.
  ///
  /// - Returns: `true` if the tap was created (consent granted or already
  ///   granted); `false` if tap creation failed (e.g. the user denied consent).
  @discardableResult
  static func primePermission() -> Bool {
    // Local IDs only — never touch instance state, so this is safe to call
    // statically and concurrently with a live capture on another instance.
    var tapID: AudioObjectID = kAudioObjectUnknown
    var aggregateDeviceID: AudioObjectID = kAudioObjectUnknown

    // Guarantee full teardown on every exit path. Mirrors cleanupTap()/cleanup()
    // teardown order (destroy aggregate device, then tap, null out IDs). No
    // AudioDeviceStop / AudioDeviceDestroyIOProcID — we never create an IO proc.
    func teardown() {
      if aggregateDeviceID != kAudioObjectUnknown {
        AudioHardwareDestroyAggregateDevice(aggregateDeviceID)
        aggregateDeviceID = kAudioObjectUnknown
      }
      if tapID != kAudioObjectUnknown {
        AudioHardwareDestroyProcessTap(tapID)
        tapID = kAudioObjectUnknown
      }
    }

    // 1. Same tap description as startCaptureOnQueue (global tap, unmuted).
    let tapUUID = UUID()
    let tapDescription = CATapDescription(stereoGlobalTapButExcludeProcesses: [])
    tapDescription.uuid = tapUUID
    tapDescription.name = "OMI System Audio Tap (permission prime)"
    tapDescription.muteBehavior = .unmuted  // Don't mute playback

    // 2. Create the process tap — this is what triggers/persists the consent.
    var status = AudioHardwareCreateProcessTap(tapDescription, &tapID)
    guard status == noErr else {
      log("SystemAudioCapture: primePermission tap creation failed (\(status)) — consent likely denied")
      teardown()
      return false
    }
    log("SystemAudioCapture: primePermission created tap with ID \(tapID)")

    // 3. Create the same aggregate device the real path creates. Some macOS
    //    versions only register the tap consent once a tap-backed aggregate
    //    device is instantiated, so mirror startCaptureOnQueue exactly.
    let aggregateDescription: [String: Any] = [
      kAudioAggregateDeviceNameKey as String: "OMI System Audio Tap Device",
      kAudioAggregateDeviceUIDKey as String: "omi.systemaudio.\(tapUUID.uuidString)",
      kAudioAggregateDeviceIsPrivateKey as String: true,
      kAudioAggregateDeviceTapListKey as String: [
        [
          kAudioSubTapUIDKey as String: tapUUID.uuidString,
          kAudioSubTapDriftCompensationKey as String: NSNumber(value: 1),
          kAudioSubTapDriftCompensationQualityKey as String:
            NSNumber(value: kAudioAggregateDriftCompensationMaxQuality),
        ]
      ],
      kAudioAggregateDeviceTapAutoStartKey as String: true,
    ]

    status = AudioHardwareCreateAggregateDevice(aggregateDescription as CFDictionary, &aggregateDeviceID)
    guard status == noErr else {
      log("SystemAudioCapture: primePermission aggregate device creation failed (\(status))")
      // The tap itself was created, so the consent already fired — treat as
      // primed after tearing the tap back down.
      teardown()
      return true
    }
    log("SystemAudioCapture: primePermission created aggregate device with ID \(aggregateDeviceID)")

    // 4. Immediately and fully tear down — no IO proc, no capture ever started.
    teardown()
    log("SystemAudioCapture: primePermission complete, tap + aggregate device destroyed")
    return true
  }

  /// Async convenience wrapper around `primePermission()`. The CoreAudio HAL
  /// calls are synchronous IPC to coreaudiod and can block (seconds after wake
  /// from sleep), so this offloads them onto a background queue — mirroring the
  /// off-main dispatch the real capture path uses in `startCapture`. Never
  /// throws; returns the same Bool as the sync form.
  @discardableResult
  static func primePermission() async -> Bool {
    await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
      DispatchQueue.global(qos: .userInitiated).async {
        continuation.resume(returning: primePermission())
      }
    }
  }

  // MARK: - Public Methods

  /// Start capturing system audio
  /// - Parameters:
  ///   - onAudioChunk: Callback receiving 16-bit PCM audio data chunks at 16kHz mono
  ///   - onAudioLevel: Optional callback receiving normalized audio level (0.0 - 1.0)
  func startCapture(onAudioChunk: @escaping AudioChunkHandler, onAudioLevel: AudioLevelHandler? = nil) async throws {
    guard !isCapturing else {
      log("SystemAudioCapture: Already capturing")
      return
    }

    // All CoreAudio HAL calls (CreateTap, CreateAggregateDevice, AudioDeviceStart) are
    // synchronous IPC to coreaudiod via mach_msg. After wake from sleep the daemon can
    // take seconds to respond, blocking the caller. Dispatch the entire setup to audioQueue,
    // mirroring the pattern already used in stopCapture().
    // The callback assignments live on audioQueue too: the serial queue is the single
    // owner of all state the HAL IO thread reads (callbacks, converter, formats), so a
    // stop's deferred clear and a restart's fresh assignment can never interleave.
    try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
      audioQueue.async { [weak self] in
        guard let self else {
          continuation.resume()
          return
        }
        self.onAudioChunk = onAudioChunk
        self.onAudioLevel = onAudioLevel
        do {
          try self.startCaptureOnQueue()
          continuation.resume()
        } catch {
          continuation.resume(throwing: error)
        }
      }
    }
  }

  /// Performs all blocking CoreAudio HAL setup. Must be called on audioQueue, not the main thread.
  private func startCaptureOnQueue() throws {
    // 1. Create tap description for all system audio
    let tapDescription = CATapDescription(stereoGlobalTapButExcludeProcesses: [])
    tapDescription.uuid = tapUUID
    tapDescription.name = "OMI System Audio Tap"
    tapDescription.muteBehavior = .unmuted  // Don't mute playback

    // 2. Create the process tap
    var status = AudioHardwareCreateProcessTap(tapDescription, &tapID)
    guard status == noErr else {
      throw SystemAudioCaptureError.tapCreationFailed(status)
    }
    log("SystemAudioCapture: Created tap with ID \(tapID)")

    // 3. Create aggregate device with tap
    //
    // IMPORTANT: drift compensation is enabled per-tap via kAudioSubTapDriftCompensationKey.
    // Without it, the aggregate device's clock can drift relative to the real output device,
    // and the system resamples on every IO cycle to compensate. That resampling produces
    // periodic crackling/artifacts in *all* system audio playback (music, calls, etc.) even
    // though we're only reading from the tap. Enabling drift compensation tells CoreAudio
    // to reconcile clocks at the sub-tap level, eliminating the artifacts.
    // CoreAudio expects a CFNumber here ("non-zero value indicates that drift compensation
    // is enabled" — see <CoreAudio/AudioHardware.h>), not a CFBoolean.
    let aggregateDescription: [String: Any] = [
      kAudioAggregateDeviceNameKey as String: "OMI System Audio Tap Device",
      kAudioAggregateDeviceUIDKey as String: "omi.systemaudio.\(tapUUID.uuidString)",
      kAudioAggregateDeviceIsPrivateKey as String: true,
      kAudioAggregateDeviceTapListKey as String: [
        [
          kAudioSubTapUIDKey as String: tapUUID.uuidString,
          kAudioSubTapDriftCompensationKey as String: NSNumber(value: 1),
          kAudioSubTapDriftCompensationQualityKey as String:
            NSNumber(value: kAudioAggregateDriftCompensationMaxQuality),
        ]
      ],
      kAudioAggregateDeviceTapAutoStartKey as String: true,
    ]

    status = AudioHardwareCreateAggregateDevice(aggregateDescription as CFDictionary, &aggregateDeviceID)
    guard status == noErr else {
      cleanupTap()
      throw SystemAudioCaptureError.aggregateDeviceFailed(status)
    }
    log("SystemAudioCapture: Created aggregate device with ID \(aggregateDeviceID)")

    // 4. Get audio format from the tap
    guard let format = getStreamFormat(for: aggregateDeviceID) else {
      cleanup()
      throw SystemAudioCaptureError.formatError
    }

    sourceSampleRate = format.mSampleRate
    log(
      "SystemAudioCapture: Source format - \(format.mSampleRate)Hz, \(format.mChannelsPerFrame) channels, \(format.mBitsPerChannel) bits"
    )

    // 5. Create AVAudioFormat for conversion (see makeConverterInputFormat — MONO).
    guard let inputFmt = Self.makeConverterInputFormat(sampleRate: format.mSampleRate) else {
      cleanup()
      throw SystemAudioCaptureError.formatError
    }
    self.inputFormat = inputFmt

    // Target format: 16kHz mono Float32 (we'll convert to Int16 manually)
    guard
      let targetFmt = AVAudioFormat(
        standardFormatWithSampleRate: targetSampleRate,
        channels: 1
      )
    else {
      cleanup()
      throw SystemAudioCaptureError.converterCreationFailed
    }
    self.targetFormat = targetFmt

    // Create audio converter for resampling
    guard let converter = AVAudioConverter(from: inputFmt, to: targetFmt) else {
      cleanup()
      throw SystemAudioCaptureError.converterCreationFailed
    }
    self.audioConverter = converter

    // 6. Create IO proc for audio callbacks
    status = AudioDeviceCreateIOProcIDWithBlock(&ioProcID, aggregateDeviceID, nil) {
      [weak self] inNow, inInputData, inInputTime, outOutputData, inOutputTime in
      self?.handleAudioInput(inInputData, timestamp: inInputTime)
    }

    guard status == noErr else {
      cleanup()
      throw SystemAudioCaptureError.ioProcCreationFailed(status)
    }

    // 7. Start the device
    status = AudioDeviceStart(aggregateDeviceID, ioProcID)
    guard status == noErr else {
      cleanup()
      throw SystemAudioCaptureError.deviceStartFailed(status)
    }

    isCapturing = true
    log("SystemAudioCapture: Started capturing system audio")
  }

  /// Stop capturing system audio
  func stopCapture() {
    guard isCapturing else { return }
    isCapturing = false

    // Capture values for background cleanup to avoid blocking main thread
    let procID = self.ioProcID
    let aggDevID = self.aggregateDeviceID
    let tID = self.tapID

    self.ioProcID = nil
    self.aggregateDeviceID = kAudioObjectUnknown
    self.tapID = kAudioObjectUnknown

    // Do NOT release the state handleAudioInput reads (callbacks, converter,
    // formats, sourceSampleRate) here: the CoreAudio HAL IO thread can still be
    // inside handleAudioInput until AudioDeviceStop below returns, and nil-ing
    // ARC references from the main thread races its reads (torn retain/release →
    // intermittent crash on stop; sourceSampleRate=0 would divide-by-zero the
    // resample math). The serial audioQueue quiesces the IOProc first and then
    // clears; startCapture assigns this state on the same queue, so a stop's
    // clear and a restart's fresh assignment can never interleave.
    // AudioDeviceStop can block — run off main thread.
    audioQueue.async { [weak self] in
      if let procID = procID, aggDevID != kAudioObjectUnknown {
        AudioDeviceStop(aggDevID, procID)
        AudioDeviceDestroyIOProcID(aggDevID, procID)
      }
      if aggDevID != kAudioObjectUnknown {
        AudioHardwareDestroyAggregateDevice(aggDevID)
      }
      if tID != kAudioObjectUnknown {
        AudioHardwareDestroyProcessTap(tID)
      }
      guard let self else { return }
      self.onAudioChunk = nil
      self.onAudioLevel = nil
      self.audioConverter = nil
      self.inputFormat = nil
      self.targetFormat = nil
      self.sourceSampleRate = 0.0
    }

    log("SystemAudioCapture: Stopped capturing")
  }

  /// Check if currently capturing
  var capturing: Bool {
    return isCapturing
  }

  // MARK: - Private Methods

  /// Get stream format for a device
  private func getStreamFormat(for deviceID: AudioObjectID) -> AudioStreamBasicDescription? {
    var address = AudioObjectPropertyAddress(
      mSelector: kAudioDevicePropertyStreamFormat,
      mScope: kAudioDevicePropertyScopeInput,
      mElement: kAudioObjectPropertyElementMain
    )

    var format = AudioStreamBasicDescription()
    var size = UInt32(MemoryLayout<AudioStreamBasicDescription>.size)

    let status = AudioObjectGetPropertyData(
      deviceID,
      &address,
      0,
      nil,
      &size,
      &format
    )

    return status == noErr ? format : nil
  }

  /// Builds the converter input format. It is ALWAYS mono: `handleAudioInput`
  /// down-mixes stereo source frames into a single channel (channel 0) and never
  /// fills a second channel, so the converter must see a mono input and perform only
  /// sample-rate conversion. Declaring the source channel count here made the
  /// converter run its own stereo→mono downmix against the unwritten channel 1,
  /// averaging our real mix against silence (~6 dB attenuation of all system audio
  /// fed to transcription). The microphone path (AudioCaptureService) likewise builds
  /// its input format with channels: 1. `static` so the mono contract is unit-testable.
  static func makeConverterInputFormat(sampleRate: Double) -> AVAudioFormat? {
    AVAudioFormat(
      commonFormat: .pcmFormatFloat32,
      sampleRate: sampleRate,
      channels: 1,
      interleaved: false
    )
  }

  /// Frames a `sourceSampleRate`→`targetSampleRate` conversion of `frameCount` input
  /// frames produces. Returns 0 when the source rate is unknown (0): dividing by a
  /// zero rate yields `.infinity`, and `AVAudioFrameCount(.infinity)` traps the
  /// real-time tap thread — reachable when the tap reports no rate at start or
  /// `stopCapture` zeroes `sourceSampleRate` mid-callback. Pure math, unit-testable.
  static func resampledFrameCapacity(
    frameCount: AVAudioFrameCount, sourceSampleRate: Double, targetSampleRate: Double
  ) -> AVAudioFrameCount {
    guard frameCount > 0, sourceSampleRate > 0, targetSampleRate > 0 else { return 0 }
    return AVAudioFrameCount(ceil(Double(frameCount) * targetSampleRate / sourceSampleRate))
  }

  /// Handle incoming audio data from the tap
  private func handleAudioInput(_ inputData: UnsafePointer<AudioBufferList>?, timestamp: UnsafePointer<AudioTimeStamp>?)
  {
    guard isCapturing,
      let bufferList = inputData?.pointee,
      let converter = audioConverter,
      let targetFmt = targetFormat
    else { return }

    // Get the first buffer (interleaved or first channel)
    let buffer = bufferList.mBuffers

    guard let data = buffer.mData, buffer.mDataByteSize > 0 else { return }

    // Calculate frame count
    let bytesPerFrame = UInt32(MemoryLayout<Float32>.size) * buffer.mNumberChannels
    guard bytesPerFrame > 0 else { return }
    let frameCount = buffer.mDataByteSize / bytesPerFrame

    guard frameCount > 0 else { return }

    // Create input AVAudioPCMBuffer
    guard let inputFmt = inputFormat,
      let inputBuffer = AVAudioPCMBuffer(pcmFormat: inputFmt, frameCapacity: frameCount)
    else { return }

    inputBuffer.frameLength = frameCount

    // Copy data to input buffer
    // System audio is typically interleaved stereo Float32
    let srcPtr = data.assumingMemoryBound(to: Float32.self)
    let channelCount = Int(buffer.mNumberChannels)

    if channelCount >= 2 {
      // Mix stereo to mono by averaging channels
      guard let floatData = inputBuffer.floatChannelData else { return }
      let monoPtr = floatData[0]

      for i in 0..<Int(frameCount) {
        let left = srcPtr[i * channelCount]
        let right = srcPtr[i * channelCount + 1]
        monoPtr[i] = (left + right) / 2.0
      }
    } else {
      // Already mono, just copy
      guard let floatData = inputBuffer.floatChannelData else { return }
      memcpy(floatData[0], srcPtr, Int(buffer.mDataByteSize))
    }

    // Calculate output frame count based on sample rate conversion
    let outputFrameCapacity = Self.resampledFrameCapacity(
      frameCount: frameCount, sourceSampleRate: sourceSampleRate, targetSampleRate: targetSampleRate)
    guard outputFrameCapacity > 0,
      let outputBuffer = AVAudioPCMBuffer(pcmFormat: targetFmt, frameCapacity: outputFrameCapacity)
    else { return }

    // Convert using input block pattern
    var error: NSError?
    hasConsumedInput = false

    let inputBlock: AVAudioConverterInputBlock = { inNumPackets, outStatus in
      if self.hasConsumedInput {
        outStatus.pointee = .noDataNow
        return nil
      }
      self.hasConsumedInput = true
      outStatus.pointee = .haveData
      return inputBuffer
    }

    converter.convert(to: outputBuffer, error: &error, withInputFrom: inputBlock)

    if let error = error {
      logError("SystemAudioCapture: Conversion error", error: error)
      return
    }

    // Convert Float32 to Int16 (linear16 PCM for DeepGram)
    guard let channelData = outputBuffer.floatChannelData?[0] else { return }

    let processedFrameLength = Int(outputBuffer.frameLength)
    var pcmData = [Int16]()
    pcmData.reserveCapacity(processedFrameLength)

    for i in 0..<processedFrameLength {
      let sample = channelData[i]
      // Clamp and convert to Int16 range (-32768 to 32767)
      let pcmSample = Int16(max(-32768, min(32767, sample * 32767)))
      pcmData.append(pcmSample)
    }

    // Convert to Data
    let byteData = pcmData.withUnsafeBufferPointer { buffer in
      return Data(buffer: buffer)
    }

    // Calculate and report audio level (RMS normalized to 0.0 - 1.0)
    if let levelHandler = onAudioLevel, !pcmData.isEmpty {
      let sumOfSquares: Float = pcmData.reduce(0.0) { acc, sample in
        let normalized = Float(sample) / 32767.0
        return acc + normalized * normalized
      }
      let rms = sqrt(sumOfSquares / Float(pcmData.count))
      // Clamp to 0.0 - 1.0 range
      let level = min(Float(1.0), max(Float(0.0), rms))
      DispatchQueue.main.async {
        levelHandler(level)
      }
    }

    // Send to callback
    onAudioChunk?(byteData)
  }

  /// Clean up tap resources
  private func cleanupTap() {
    if tapID != kAudioObjectUnknown {
      AudioHardwareDestroyProcessTap(tapID)
      tapID = kAudioObjectUnknown
    }
  }

  /// Clean up all resources
  private func cleanup() {
    if let procID = ioProcID, aggregateDeviceID != kAudioObjectUnknown {
      AudioDeviceStop(aggregateDeviceID, procID)
      AudioDeviceDestroyIOProcID(aggregateDeviceID, procID)
      ioProcID = nil
    }

    if aggregateDeviceID != kAudioObjectUnknown {
      AudioHardwareDestroyAggregateDevice(aggregateDeviceID)
      aggregateDeviceID = kAudioObjectUnknown
    }

    cleanupTap()

    audioConverter = nil
    inputFormat = nil
    targetFormat = nil
    sourceSampleRate = 0.0
  }

  deinit {
    // Call the HAL teardown directly — do NOT `audioQueue.sync` here.
    // deinit can run *on* audioQueue (the last strong reference is captured by
    // the `audioQueue.async` block in startCapture and released when that block
    // finishes), and dispatching sync to the current serial queue deadlocks it
    // permanently. Mirrors AudioCaptureService.deinit. The object has no
    // remaining references, so no concurrent audioQueue work can touch it, and
    // these HAL calls are thread-safe — a direct call completes cleanup before
    // deallocation (which `audioQueue.async` could not guarantee).
    let procID = self.ioProcID
    let aggDevID = self.aggregateDeviceID
    let tID = self.tapID
    if let procID = procID, aggDevID != kAudioObjectUnknown {
      AudioDeviceStop(aggDevID, procID)
      AudioDeviceDestroyIOProcID(aggDevID, procID)
    }
    if aggDevID != kAudioObjectUnknown {
      AudioHardwareDestroyAggregateDevice(aggDevID)
    }
    if tID != kAudioObjectUnknown {
      AudioHardwareDestroyProcessTap(tID)
    }
  }
}