audiotee/Sources/AudioTeeCLI/AudioTee.swift

import AudioTeeCore
import CoreAudio
import Foundation

// Set by the SIGINT/SIGTERM handlers, which — being passed to the C `signal()`
// API — cannot capture `self` and so can't touch instance state directly.
private var shouldStop = false

struct AudioTee {
  var includeProcesses: [Int32] = []
  var excludeProcesses: [Int32] = []
  var mute: Bool = false
  var stereo: Bool = false
  var sampleRate: Double?
  var chunkDuration: Double = 0.2
  var captureMic: Bool = false
  var micOutputPath: String?

  init() {}

  static func main() {
    let parser = SimpleArgumentParser(
      programName: "audiotee",
      abstract: "Capture system audio and stream to stdout",
      discussion: """
        AudioTee captures system audio using Core Audio taps and streams it as structured output.

        Process filtering:
        • include-processes: Only tap specified process IDs (empty = all processes)
        • exclude-processes: Tap all processes except specified ones
        • mute: How to handle processes being tapped

        Examples:
          audiotee                              # Auto format, tap all processes
          audiotee --sample-rate 16000          # Convert to 16kHz mono for ASR
          audiotee --sample-rate 8000           # Convert to 8kHz for telephony
          audiotee --include-processes 1234     # Only tap process 1234
          audiotee --include-processes 1234 5678 9012  # Tap only these processes
          audiotee --exclude-processes 1234 5678       # Tap everything except these
          audiotee --mute                       # Mute processes being tapped
          audiotee --capture-mic --mic-output mic.pcm > system.pcm
                                                 # Capture system audio and mic to separate files
        """
    )

    // Configure arguments
    parser.addArrayOption(
      name: "include-processes",
      help: "Process IDs to include (space-separated, empty = all processes)")
    parser.addArrayOption(
      name: "exclude-processes", help: "Process IDs to exclude (space-separated)")
    parser.addFlag(name: "mute", help: "Mute processes being tapped")
    parser.addFlag(name: "stereo", help: "Records in stereo")
    parser.addOption(
      name: "sample-rate",
      help: "Target sample rate (8000, 16000, 22050, 24000, 32000, 44100, 48000)")
    parser.addOption(
      name: "chunk-duration", help: "Audio chunk duration in seconds", defaultValue: "0.2")
    parser.addFlag(
      name: "capture-mic",
      help: "Also capture the default input device (microphone) as a second track")
    parser.addOption(
      name: "mic-output",
      help: "File path to write microphone PCM audio to (required with --capture-mic)")

    // Parse arguments
    do {
      try parser.parse()

      var audioTee = AudioTee()

      // Extract values
      audioTee.includeProcesses = try parser.getArrayValue("include-processes", as: Int32.self)
      audioTee.excludeProcesses = try parser.getArrayValue("exclude-processes", as: Int32.self)
      audioTee.mute = parser.getFlag("mute")
      audioTee.stereo = parser.getFlag("stereo")
      audioTee.sampleRate = try parser.getOptionalValue("sample-rate", as: Double.self)
      audioTee.chunkDuration = try parser.getValue("chunk-duration", as: Double.self)
      audioTee.captureMic = parser.getFlag("capture-mic")
      audioTee.micOutputPath = try parser.getOptionalValue("mic-output", as: String.self)

      // Validate
      try audioTee.validate()

      // Run
      try audioTee.run()

    } catch ArgumentParserError.helpRequested {
      parser.printHelp()
      exit(0)
    } catch ArgumentParserError.validationFailed(let message) {
      print("Error: \(message)", to: &standardError)
      exit(1)
    } catch let error as ArgumentParserError {
      print("Error: \(error.description)", to: &standardError)
      parser.printHelp()
      exit(1)
    } catch {
      print("Error: \(error)", to: &standardError)
      exit(1)
    }
  }

  func validate() throws {
    if !includeProcesses.isEmpty && !excludeProcesses.isEmpty {
      throw ArgumentParserError.validationFailed(
        "Cannot specify both --include-processes and --exclude-processes")
    }
    if captureMic && micOutputPath == nil {
      throw ArgumentParserError.validationFailed(
        "--mic-output is required when --capture-mic is set")
    }
    if !captureMic && micOutputPath != nil {
      throw ArgumentParserError.validationFailed("--mic-output requires --capture-mic")
    }
  }

  func run() throws {
    setupSignalHandlers()

    AudioTeeLogging.logger.info("Starting AudioTee...")

    // Validate chunk duration
    guard chunkDuration > 0 && chunkDuration <= 5.0 else {
      AudioTeeLogging.logger.error(
        "Invalid chunk duration",
        context: ["chunk_duration": String(chunkDuration), "valid_range": "0.0 < duration <= 5.0"])
      throw ExitCode.failure
    }

    // Convert include/exclude processes to TapConfiguration format
    let (processes, isExclusive) = convertProcessFlags()

    let tapConfig = TapConfiguration(
      processes: processes,
      muteBehavior: mute ? .muted : .unmuted,
      isExclusive: isExclusive,
      isMono: !stereo
    )

    let audioTapManager = AudioTapManager()
    do {
      try audioTapManager.setupAudioTap(with: tapConfig)
    } catch AudioTeeError.pidTranslationFailed(let failedPIDs) {
      AudioTeeLogging.logger.error(
        "Failed to translate process IDs to audio objects",
        context: [
          "failed_pids": failedPIDs.map(String.init).joined(separator: ", "),
          "suggestion": "Check that the process IDs exist and are running",
        ])
      throw ExitCode.failure
    } catch {
      AudioTeeLogging.logger.error(
        "Failed to setup audio tap", context: ["error": String(describing: error)])
      throw ExitCode.failure
    }

    guard let deviceID = audioTapManager.getDeviceID() else {
      AudioTeeLogging.logger.error("Failed to get device ID from audio tap manager")
      throw ExitCode.failure
    }

    let outputHandler = BinaryAudioOutputHandler()
    let recorder = try AudioRecorder(
      deviceID: deviceID, outputHandler: outputHandler, convertToSampleRate: sampleRate,
      chunkDuration: chunkDuration)
    try recorder.startRecording()

    let micRecorder = try setupMicRecorderIfNeeded()
    try micRecorder?.startRecording()

    // Run until the run loop is stopped (by signal handler). shouldStop is
    // checked on every iteration (not just the CFRunLoopRun result) because
    // a signal can arrive during setup, before this loop is ever entered —
    // CFRunLoopStop has no lasting effect on a run loop that isn't running yet.
    while !shouldStop {
      let result = CFRunLoopRunInMode(CFRunLoopMode.defaultMode, 0.1, false)
      if result == CFRunLoopRunResult.stopped || result == CFRunLoopRunResult.finished {
        break
      }
    }

    AudioTeeLogging.logger.info("Shutting down...")
    recorder.stopRecording()
    micRecorder?.stopRecording()
  }

  /// Sets up a second, independent recording pipeline reading from the
  /// default input device (microphone) when --capture-mic was requested.
  /// Its audio is written to its own file rather than stdout: writes from
  /// two concurrent Core Audio IO threads interleaved on one fd/stream
  /// would otherwise corrupt both tracks.
  private func setupMicRecorderIfNeeded() throws -> AudioRecorder? {
    guard captureMic, let micOutputPath = micOutputPath else {
      return nil
    }

    let micDeviceID: AudioObjectID
    do {
      micDeviceID = try InputDeviceResolver.defaultInputDevice()
    } catch {
      AudioTeeLogging.logger.error(
        "Failed to resolve default input device", context: ["error": String(describing: error)])
      throw ExitCode.failure
    }

    let micFd = open(micOutputPath, O_WRONLY | O_CREAT | O_TRUNC, 0o644)
    guard micFd >= 0 else {
      AudioTeeLogging.logger.error(
        "Failed to open mic output file",
        context: ["path": micOutputPath, "errno": String(errno)])
      throw ExitCode.failure
    }

    let micOutputHandler = BinaryAudioOutputHandler(fd: micFd, source: "mic")
    return try AudioRecorder(
      deviceID: micDeviceID, outputHandler: micOutputHandler, convertToSampleRate: sampleRate,
      chunkDuration: chunkDuration)
  }

  private func setupSignalHandlers() {
    signal(SIGINT) { _ in
      AudioTeeLogging.logger.info("Received SIGINT, initiating graceful shutdown...")
      shouldStop = true
      CFRunLoopStop(CFRunLoopGetMain())
    }
    signal(SIGTERM) { _ in
      AudioTeeLogging.logger.info("Received SIGTERM, initiating graceful shutdown...")
      shouldStop = true
      CFRunLoopStop(CFRunLoopGetMain())
    }
  }

  private func convertProcessFlags() -> ([Int32], Bool) {
    if !includeProcesses.isEmpty {
      // Include specific processes only
      return (includeProcesses, false)
    } else if !excludeProcesses.isEmpty {
      // Exclude specific processes (tap everything except these)
      return (excludeProcesses, true)
    } else {
      // Default: tap everything
      return ([], true)
    }
  }
}

// Helper for stderr output
var standardError = FileHandle.standardError

extension FileHandle: TextOutputStream {
  public func write(_ string: String) {
    let data = Data(string.utf8)
    self.write(data)
  }
}

// Exit code handling
enum ExitCode: Error {
  case failure
}

extension ExitCode {
  var code: Int32 {
    switch self {
    case .failure:
      return 1
    }
  }
}