wisper-key/Sources/WhisperKey/Audio.swift

import AVFoundation
import Foundation

struct RecordedChunk: Sendable {
    let url: URL
    let continuesPhrase: Bool
}

final class AudioRecorder: @unchecked Sendable {
    private var engine: AVAudioEngine?
    private let lock = NSLock()
    private var noSpeechTimeout: DispatchWorkItem?
    private var currentURL: URL?
    private var outputFile: AVAudioFile?
    private var currentFrames: AVAudioFramePosition = 0
    private var currentChunkHasSpeech = false
    private var heardSpeech = false
    private var silentFrames: AVAudioFramePosition = 0
    private var chunkSilentFrames: AVAudioFramePosition = 0
    private var silenceTriggered = false
    private var onChunk: (@Sendable (RecordedChunk) -> Void)?
    private var onSilence: (@Sendable () -> Void)?
    private var recordingFormat: AVAudioFormat?

    func start(
        onChunk: @escaping @Sendable (RecordedChunk) -> Void,
        onSilence: @escaping @Sendable () -> Void
    ) throws {
        guard engine?.isRunning != true else { return }
        let engine = AVAudioEngine()
        self.engine = engine
        let input = engine.inputNode
        let format = input.outputFormat(forBus: 0)
        self.onChunk = onChunk
        self.onSilence = onSilence
        recordingFormat = format
        try beginChunk(format: format)
        input.installTap(onBus: 0, bufferSize: 4096, format: format) { [weak self] buffer, _ in
            self?.consume(buffer, format: format)
        }
        do {
            engine.prepare()
            try engine.start()
            scheduleNoSpeechTimeout()
        } catch {
            input.removeTap(onBus: 0)
            engine.reset()
            self.engine = nil
            if let currentURL { try? FileManager.default.removeItem(at: currentURL) }
            outputFile = nil
            currentURL = nil
            throw error
        }
    }

    func stop() throws -> URL? {
        noSpeechTimeout?.cancel()
        noSpeechTimeout = nil
        if let engine {
            if engine.isRunning {
                engine.inputNode.removeTap(onBus: 0)
                engine.stop()
            }
            engine.reset()
            self.engine = nil
        }
        lock.lock()
        defer { lock.unlock() }
        let finalURL = currentFrames > 0 && currentChunkHasSpeech ? currentURL : nil
        if finalURL == nil, let currentURL {
            try? FileManager.default.removeItem(at: currentURL)
        }
        outputFile = nil
        currentURL = nil
        currentFrames = 0
        currentChunkHasSpeech = false
        heardSpeech = false
        silentFrames = 0
        chunkSilentFrames = 0
        silenceTriggered = false
        return finalURL
    }

    private func scheduleNoSpeechTimeout() {
        let work = DispatchWorkItem { [weak self] in
            guard let self else { return }
            lock.lock()
            let shouldStop = engine?.isRunning == true && !heardSpeech && !silenceTriggered
            if shouldStop { silenceTriggered = true }
            lock.unlock()
            if shouldStop { onSilence?() }
        }
        noSpeechTimeout = work
        DispatchQueue.global(qos: .utility).asyncAfter(deadline: .now() + 8, execute: work)
    }

    private func beginChunk(format: AVAudioFormat) throws {
        let url = FileManager.default.temporaryDirectory.appendingPathComponent("whisperkey-\(UUID().uuidString).caf")
        outputFile = try AVAudioFile(forWriting: url, settings: format.settings)
        currentURL = url
        currentFrames = 0
        currentChunkHasSpeech = false
        chunkSilentFrames = 0
    }

    private func consume(_ buffer: AVAudioPCMBuffer, format: AVAudioFormat) {
        lock.lock()
        defer { lock.unlock() }
        guard let activeOutputFile = outputFile else { return }
        try? activeOutputFile.write(from: buffer)
        currentFrames += AVAudioFramePosition(buffer.frameLength)

        let rms = bufferRMS(buffer)
        if rms > 0.006 {
            heardSpeech = true
            currentChunkHasSpeech = true
            silentFrames = 0
            chunkSilentFrames = 0
        } else if heardSpeech {
            silentFrames += AVAudioFramePosition(buffer.frameLength)
            if currentChunkHasSpeech {
                chunkSilentFrames += AVAudioFramePosition(buffer.frameLength)
            }
        }

        let silenceLimit = AVAudioFramePosition(format.sampleRate * 5.0)
        if heardSpeech, silentFrames >= silenceLimit, !silenceTriggered {
            silenceTriggered = true
            onSilence?()
            return
        }

        // Keep enough context for Whisper, but insert completed phrases while
        // the user is still dictating. Prefer a short natural pause; cap a
        // continuous phrase so the input never stays blank indefinitely.
        let minimumChunk = AVAudioFramePosition(format.sampleRate * 2.5)
        let pauseBoundary = AVAudioFramePosition(format.sampleRate * 1.0)
        let maximumChunk = AVAudioFramePosition(format.sampleRate * 4.0)
        let reachedPause =
            currentFrames >= minimumChunk && chunkSilentFrames >= pauseBoundary
        let reachedMaximum = currentFrames >= maximumChunk
        let shouldRotate = currentChunkHasSpeech && (reachedPause || reachedMaximum)

        if shouldRotate, let finishedURL = currentURL {
            let finishedHasSpeech = currentChunkHasSpeech
            let continuesPhrase = reachedMaximum && !reachedPause
            outputFile = nil
            do {
                try beginChunk(format: format)
                if finishedHasSpeech {
                    onChunk?(RecordedChunk(url: finishedURL, continuesPhrase: continuesPhrase))
                } else {
                    try? FileManager.default.removeItem(at: finishedURL)
                }
            } catch {
                try? FileManager.default.removeItem(at: finishedURL)
                silenceTriggered = true
                onSilence?()
            }
        }
    }

    private func bufferRMS(_ buffer: AVAudioPCMBuffer) -> Float {
        guard let channels = buffer.floatChannelData, buffer.frameLength > 0 else { return 0 }
        let samples = channels[0]
        var sum: Float = 0
        for index in 0..<Int(buffer.frameLength) {
            let value = samples[index]
            sum += value * value
        }
        return sqrt(sum / Float(buffer.frameLength))
    }
}

enum AudioConverter {
    static func toWhisperWAV(_ input: URL) throws -> URL {
        let output = FileManager.default.temporaryDirectory.appendingPathComponent("whisperkey-\(UUID().uuidString).wav")
        let ffmpegCandidates = ["/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg"]
        guard let ffmpeg = ffmpegCandidates.first(where: FileManager.default.isExecutableFile(atPath:)) else {
            throw WhisperKeyError.missingDependency("ffmpeg")
        }
        let process = Process()
        process.executableURL = URL(fileURLWithPath: ffmpeg)
        process.arguments = ["-hide_banner", "-loglevel", "error", "-y", "-i", input.path, "-ar", "16000", "-ac", "1", "-c:a", "pcm_s16le", output.path]
        let stderr = Pipe()
        process.standardError = stderr
        try process.run()
        process.waitUntilExit()
        guard process.terminationStatus == 0 else {
            let text = String(data: stderr.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8) ?? "ffmpeg error"
            throw WhisperKeyError.processFailed(text)
        }
        return output
    }
}