import Foundation

/// Manages audio capture to WAV files for debugging purposes.
/// Extracted from SherpaOnnxOfflineStt.swift to reduce file size.
class AudioCaptureManager {
    private var isEnabled: Bool = false
    private var fileHandle: FileHandle?
    private var filePath: String?
    private var samplesWritten: Int64 = 0
    private var sampleRate: Int

    init(sampleRate: Int = STTConstants.defaultSampleRate) {
        self.sampleRate = sampleRate
    }

    var captureEnabled: Bool {
        return isEnabled
    }

    var currentFilePath: String? {
        return filePath
    }

    var totalSamplesWritten: Int64 {
        return samplesWritten
    }

    /// Start capturing audio to a WAV file
    func startCapture(to path: String) throws {
        guard !isEnabled else {
            throw AudioCaptureError.alreadyCapturing
        }

        let fileManager = FileManager.default
        let dirPath = (path as NSString).deletingLastPathComponent

        // Create directory if needed
        try fileManager.createDirectory(atPath: dirPath, withIntermediateDirectories: true)

        // Create file with placeholder header
        fileManager.createFile(atPath: path, contents: nil)
        let handle = try FileHandle(forWritingTo: URL(fileURLWithPath: path))

        // Write placeholder header (44 bytes)
        let header = createWavHeader(dataSize: 0)
        handle.write(header)

        fileHandle = handle
        filePath = path
        samplesWritten = 0
        isEnabled = true

        STTLogger.i("AudioCapture", "Started capture: \(path)")
    }

    /// Stop capturing and finalize the WAV file
    func stopCapture() throws -> Bool {
        guard isEnabled else {
            return false
        }

        isEnabled = false

        guard let handle = fileHandle, let path = filePath else {
            return false
        }

        let dataSize = Int(samplesWritten * 2) // 16-bit = 2 bytes per sample

        // Seek to beginning and write updated header with correct size
        handle.seek(toFileOffset: 0)
        let header = createWavHeader(dataSize: dataSize)
        handle.write(header)

        handle.closeFile()
        STTLogger.i("AudioCapture", "Stopped capture: \(samplesWritten) samples written to \(path)")

        fileHandle = nil
        filePath = nil
        samplesWritten = 0

        return true
    }

    /// Write audio samples to the capture file
    func writeSamples(_ samples: [Float]) {
        guard isEnabled, let handle = fileHandle else { return }

        var data = Data(capacity: samples.count * 2)
        for sample in samples {
            // Convert float [-1.0, 1.0] to 16-bit PCM
            let pcmValue = Int16(max(-32768, min(32767, sample * 32767)))
            var littleEndian = pcmValue.littleEndian
            data.append(Data(bytes: &littleEndian, count: 2))
        }

        handle.seekToEndOfFile()
        handle.write(data)
        samplesWritten += Int64(samples.count)
    }

    /// Create a WAV file header
    private func createWavHeader(dataSize: Int) -> Data {
        var header = Data(capacity: 44)

        let channels = 1
        let bitsPerSample = 16
        let byteRate = sampleRate * channels * bitsPerSample / 8
        let blockAlign = channels * bitsPerSample / 8

        // RIFF header
        header.append("RIFF".data(using: .ascii)!)
        header.append(UInt32(36 + dataSize).littleEndianData)
        header.append("WAVE".data(using: .ascii)!)

        // fmt subchunk
        header.append("fmt ".data(using: .ascii)!)
        header.append(UInt32(16).littleEndianData) // Subchunk1Size
        header.append(UInt16(1).littleEndianData) // AudioFormat (PCM)
        header.append(UInt16(channels).littleEndianData) // NumChannels
        header.append(UInt32(sampleRate).littleEndianData) // SampleRate
        header.append(UInt32(byteRate).littleEndianData) // ByteRate
        header.append(UInt16(blockAlign).littleEndianData) // BlockAlign
        header.append(UInt16(bitsPerSample).littleEndianData) // BitsPerSample

        // data subchunk
        header.append("data".data(using: .ascii)!)
        header.append(UInt32(dataSize).littleEndianData)

        return header
    }

    /// Update sample rate (call before starting capture)
    func setSampleRate(_ rate: Int) {
        self.sampleRate = rate
    }
}

/// Errors that can occur during audio capture
enum AudioCaptureError: Error, LocalizedError {
    case alreadyCapturing
    case notCapturing
    case fileCreationFailed
    case writeFailed

    var errorDescription: String? {
        switch self {
        case .alreadyCapturing:
            return "Audio capture is already in progress"
        case .notCapturing:
            return "No audio capture in progress"
        case .fileCreationFailed:
            return "Failed to create capture file"
        case .writeFailed:
            return "Failed to write audio data"
        }
    }
}

// MARK: - Extensions for WAV header creation

extension UInt32 {
    var littleEndianData: Data {
        var value = self.littleEndian
        return Data(bytes: &value, count: 4)
    }
}

extension UInt16 {
    var littleEndianData: Data {
        var value = self.littleEndian
        return Data(bytes: &value, count: 2)
    }
}