// AACFileWriter.swift
//
// Serial, direct-to-file AAC-LC encoder that consumes the same PCM
// AVAudioPCMBuffer delivered to the main capture tap and writes ADTS-framed
// AAC frames to a single file on disk. AAC bytes never cross the JS bridge.
//
// Crashproof: ADTS frames carry a sync word (0xFFF9) at every frame, so the
// file on disk is decodable up to the last complete frame at any moment.
// No global header, no footer, no finalisation needed.

import AVFoundation
import CoreAudio
import Foundation

final class AACFileWriter {
    private static let className = "AACFileWriter"

    /// ADTS sample-rate table (MPEG-4 Part 3, Table 1.18).
    static let adtsSampleRateTable: [Int] = [
        96_000, 88_200, 64_000, 48_000, 44_100, 32_000,
        24_000, 22_050, 16_000, 12_000, 11_025,  8_000, 7_350,
    ]

    static func adtsSampleRateIndex(for sampleRate: Int) -> Int? {
        return adtsSampleRateTable.firstIndex(of: sampleRate)
    }

    enum WriterError: Swift.Error, CustomStringConvertible {
        case invalidSampleRate(Int)
        case invalidChannels(Int)
        case fileCreationFailed(URL)
        case fileHandleFailed(URL)

        var description: String {
            switch self {
            case .invalidSampleRate(let rate):
                return "Sample rate \(rate) Hz is not in the ADTS sample-rate table"
            case .invalidChannels(let ch):
                return "Unsupported channel count: \(ch) (only 1 or 2)"
            case .fileCreationFailed(let url):
                return "Failed to create file at \(url.path)"
            case .fileHandleFailed(let url):
                return "Failed to open write handle on \(url.path)"
            }
        }
    }

    // Serial queue: all converter + file I/O happens here, never on the audio thread.
    private let queue = DispatchQueue(label: "com.siteed.aac-encoder", qos: .utility)

    // State (only touched on `queue`).
    private var converter: AVAudioConverter?
    private var fileHandle: FileHandle?
    private var outputBuffer: AVAudioCompressedBuffer?
    private var aacFormat: AVAudioFormat?
    private var sampleRate: Int = 0
    private var channels: Int = 1
    private var bitrate: Int = 0
    private var sampleRateIndex: Int = 0
    private var channelConfig: Int = 1
    private var isRunning: Bool = false
    private var setupFailed: Bool = false
    private var _fileURL: URL?

    /// URL of the file being written. Valid after a successful `start(...)`.
    /// Thread-safe.
    var fileURL: URL? {
        return queue.sync { _fileURL }
    }

    /// Starts a new encoder session. The converter is built lazily on the
    /// first `encode(...)` call so that the input format matches the actual
    /// `AVAudioPCMBuffer.format` produced by the capture pipeline after
    /// resampling (including the `isInterleaved` flag).
    ///
    /// - Parameters:
    ///   - sampleRate: Target output sample rate. Must be present in the
    ///                 ADTS sample-rate table.
    ///   - channels: 1 (mono) or 2 (stereo).
    ///   - bitrate: AAC-LC target bitrate in bps.
    ///   - outputURL: Destination file. Truncated if it exists.
    func start(sampleRate: Int, channels: Int, bitrate: Int, outputURL: URL) throws {
        guard let rateIndex = Self.adtsSampleRateIndex(for: sampleRate) else {
            throw WriterError.invalidSampleRate(sampleRate)
        }
        guard channels == 1 || channels == 2 else {
            throw WriterError.invalidChannels(channels)
        }

        let fileManager = FileManager.default
        try fileManager.createDirectory(
            at: outputURL.deletingLastPathComponent(),
            withIntermediateDirectories: true,
            attributes: nil
        )
        if fileManager.fileExists(atPath: outputURL.path) {
            try? fileManager.removeItem(at: outputURL)
        }
        guard fileManager.createFile(atPath: outputURL.path, contents: nil, attributes: nil) else {
            throw WriterError.fileCreationFailed(outputURL)
        }
        let handle: FileHandle
        do {
            handle = try FileHandle(forWritingTo: outputURL)
        } catch {
            throw WriterError.fileHandleFailed(outputURL)
        }

        queue.sync {
            self.sampleRate = sampleRate
            self.channels = channels
            self.bitrate = bitrate
            self.sampleRateIndex = rateIndex
            self.channelConfig = channels
            self._fileURL = outputURL
            self.fileHandle = handle
            self.converter = nil
            self.outputBuffer = nil
            self.aacFormat = nil
            self.setupFailed = false
            self.isRunning = true
        }

        Logger.info(Self.className,
                    "Started \(outputURL.lastPathComponent) @ \(sampleRate) Hz, \(channels) ch, \(bitrate) bps")
    }

    /// Feeds one PCM buffer into the encoder. Returns immediately; the actual
    /// encoding + file I/O happens asynchronously on the internal serial
    /// queue so the audio capture thread is never blocked.
    func encode(_ pcmBuffer: AVAudioPCMBuffer) {
        queue.async { [weak self] in
            guard let self = self, self.isRunning, !self.setupFailed else { return }

            if self.converter == nil {
                if !self.setupConverter(inputFormat: pcmBuffer.format) {
                    self.setupFailed = true
                    return
                }
            }
            guard let converter = self.converter else { return }
            self.encodeAndWrite(pcmBuffer: pcmBuffer, converter: converter, endOfStream: false)
        }
    }

    /// Synchronizes completed writes to disk without finalizing the AAC stream.
    /// Used on pause / interruption boundaries where encoding must resume into
    /// the same continuous ADTS file.
    func synchronize() {
        queue.sync { [weak self] in
            guard let self = self, self.isRunning, !self.setupFailed else { return }
            try? self.fileHandle?.synchronize()
        }
    }

    /// Flushes, closes the file and releases the encoder. Returns the final
    /// file URL so callers can surface it in stop results.
    @discardableResult
    func stop() -> URL? {
        return queue.sync { () -> URL? in
            let url = self._fileURL
            guard self.isRunning else { return url }

            if let converter = self.converter {
                self.encodeAndWrite(pcmBuffer: nil, converter: converter, endOfStream: true)
            }
            do {
                try self.fileHandle?.synchronize()
                try self.fileHandle?.close()
            } catch {
                Logger.warn(Self.className, "Close failed: \(error.localizedDescription)")
            }
            self.fileHandle = nil
            self.converter = nil
            self.outputBuffer = nil
            self.aacFormat = nil
            self.isRunning = false
            Logger.info(Self.className, "Stopped: \(url?.lastPathComponent ?? "<nil>")")
            return url
        }
    }

    // MARK: - Private

    private func setupConverter(inputFormat: AVAudioFormat) -> Bool {
        var outputASBD = AudioStreamBasicDescription(
            mSampleRate: Float64(sampleRate),
            mFormatID: kAudioFormatMPEG4AAC,
            mFormatFlags: AudioFormatFlags(MPEG4ObjectID.AAC_LC.rawValue),
            mBytesPerPacket: 0,
            mFramesPerPacket: 1024,
            mBytesPerFrame: 0,
            mChannelsPerFrame: UInt32(channels),
            mBitsPerChannel: 0,
            mReserved: 0
        )
        guard let aacFormat = AVAudioFormat(streamDescription: &outputASBD) else {
            Logger.error(Self.className, "Failed to build AAC output format (\(sampleRate) Hz, \(channels) ch)")
            return false
        }
        guard let converter = AVAudioConverter(from: inputFormat, to: aacFormat) else {
            Logger.error(Self.className, "AVAudioConverter init failed (input: \(inputFormat))")
            return false
        }
        converter.bitRate = bitrate

        let buffer = AVAudioCompressedBuffer(
            format: aacFormat,
            packetCapacity: 16,
            maximumPacketSize: 4096
        )

        self.aacFormat = aacFormat
        self.converter = converter
        self.outputBuffer = buffer
        Logger.info(Self.className,
                    "Converter ready — input=\(inputFormat), bitrate=\(bitrate) bps")
        return true
    }

    private func encodeAndWrite(pcmBuffer: AVAudioPCMBuffer?,
                                converter: AVAudioConverter,
                                endOfStream: Bool) {
        guard let outputBuffer = self.outputBuffer, let handle = self.fileHandle else { return }

        var inputDelivered = false
        let maxLoops = 64 // safety rail
        var loop = 0

        while loop < maxLoops {
            loop += 1
            outputBuffer.packetCount = 0
            outputBuffer.byteLength = 0
            var convError: NSError?

            let status = converter.convert(to: outputBuffer, error: &convError) { _, outStatus in
                if endOfStream {
                    outStatus.pointee = .endOfStream
                    return nil
                }
                if let buf = pcmBuffer, !inputDelivered {
                    inputDelivered = true
                    outStatus.pointee = .haveData
                    return buf
                }
                outStatus.pointee = .noDataNow
                return nil
            }

            if let err = convError {
                Logger.error(Self.className, "convert() failed: \(err.localizedDescription)")
                break
            }

            writePackets(outputBuffer: outputBuffer, handle: handle)

            if status == .endOfStream || status == .inputRanDry { break }
            if outputBuffer.packetCount == 0 { break }
        }
    }

    private func writePackets(outputBuffer: AVAudioCompressedBuffer, handle: FileHandle) {
        guard outputBuffer.packetCount > 0,
              let packetDescriptions = outputBuffer.packetDescriptions else { return }
        let basePtr = outputBuffer.data.assumingMemoryBound(to: UInt8.self)
        let count = Int(outputBuffer.packetCount)

        // Batch ADTS frames into a single write to minimize syscalls.
        var assembled = Data()
        for i in 0..<count {
            let desc = packetDescriptions[i]
            let packetLength = Int(desc.mDataByteSize)
            let offset = Int(desc.mStartOffset)
            if packetLength <= 0 { continue }
            let header = Self.makeADTSHeader(
                packetLength: packetLength,
                sampleRateIndex: sampleRateIndex,
                channelConfig: channelConfig
            )
            assembled.append(header)
            assembled.append(Data(bytes: basePtr.advanced(by: offset), count: packetLength))
        }
        if assembled.isEmpty { return }
        do {
            try handle.write(contentsOf: assembled)
        } catch {
            Logger.error(Self.className, "File write failed: \(error.localizedDescription)")
        }
    }

    /// Builds an ADTS (Audio Data Transport Stream) frame header for one
    /// AAC-LC packet. 7 bytes, no CRC, MPEG-2.
    ///
    /// - Parameters:
    ///   - packetLength: Length of the AAC raw data block.
    ///   - sampleRateIndex: Index in the ADTS sample-rate table.
    ///   - channelConfig: 1 for mono, 2 for stereo.
    static func makeADTSHeader(
        packetLength: Int,
        sampleRateIndex: Int,
        channelConfig: Int
    ) -> Data {
        let profile = 2 // AAC-LC
        let frameLength = packetLength + 7
        var header = Data(count: 7)
        header[0] = 0xFF
        header[1] = 0xF9 // MPEG-2, no CRC
        header[2] = UInt8(((profile - 1) << 6) | (sampleRateIndex << 2) | (channelConfig >> 2))
        header[3] = UInt8(((channelConfig & 0x3) << 6) | ((frameLength >> 11) & 0x3))
        header[4] = UInt8((frameLength >> 3) & 0xFF)
        header[5] = UInt8(((frameLength & 0x7) << 5) | 0x1F)
        header[6] = 0xFC
        return header
    }
}