import AVFoundation
import Foundation

/// Converts arbitrary PCM input (any sample rate / channel count) to the
/// 16 kHz mono Float32 the models expect. Wraps AVAudioConverter with a
/// simple incremental API safe to call from one audio-callback thread.
public final class AudioResampler {
    public static let targetRate = 16_000.0

    private var converter: AVAudioConverter?
    private var sourceFormat: AVAudioFormat?
    private let targetFormat = AVAudioFormat(
        commonFormat: .pcmFormatFloat32, sampleRate: targetRate,
        channels: 1, interleaved: false)!

    public init() {}

    /// Converts one buffer. Reconfigures automatically if the input format
    /// changes (e.g. route change to a Bluetooth mic).
    public func convert(_ buffer: AVAudioPCMBuffer) throws -> [Float] {
        let format = buffer.format
        if format.sampleRate == Self.targetRate, format.channelCount == 1,
           format.commonFormat == .pcmFormatFloat32, let ch = buffer.floatChannelData {
            return Array(UnsafeBufferPointer(start: ch[0], count: Int(buffer.frameLength)))
        }
        if converter == nil || sourceFormat != format {
            guard let c = AVAudioConverter(from: format, to: targetFormat) else {
                throw SpeechError.audioFormatUnsupported(format.description)
            }
            converter = c
            sourceFormat = format
        }
        let ratio = Self.targetRate / format.sampleRate
        let capacity = AVAudioFrameCount(Double(buffer.frameLength) * ratio) + 16
        guard let out = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: capacity) else {
            throw SpeechError.resamplingFailed("buffer allocation")
        }
        var err: NSError?
        var served = false
        converter!.convert(to: out, error: &err) { _, status in
            if served { status.pointee = .noDataNow; return nil }
            served = true
            status.pointee = .haveData
            return buffer
        }
        if let err { throw SpeechError.resamplingFailed(err.localizedDescription) }
        guard let ch = out.floatChannelData else {
            throw SpeechError.resamplingFailed("no channel data")
        }
        return Array(UnsafeBufferPointer(start: ch[0], count: Int(out.frameLength)))
    }

    /// Loads an audio file and returns 16 kHz mono samples.
    public static func loadFile(_ url: URL) throws -> [Float] {
        guard let file = try? AVAudioFile(forReading: url) else {
            throw SpeechError.audioFormatUnsupported("cannot open \(url.lastPathComponent)")
        }
        guard let buf = AVAudioPCMBuffer(pcmFormat: file.processingFormat,
                                         frameCapacity: AVAudioFrameCount(file.length)) else {
            throw SpeechError.resamplingFailed("file buffer allocation")
        }
        try file.read(into: buf)
        return try AudioResampler().convert(buf)
    }
}