import AVFoundation
import CoreMedia
import Darwin
import Foundation
import ScreenCaptureKit

private enum CaptureFailure: Error, CustomStringConvertible {
    case usage
    case invalidFormat
    case microphonePermission
    case noDisplay
    case audioConverter(String)
    case systemAudio(String)

    var description: String {
        switch self {
        case .usage:
            return "Usage: macos-meeting-audio-capture <microphone|system> <sample-rate> <channels>"
        case .invalidFormat:
            return "Sample rate must be positive and channels must be 1 or 2."
        case .microphonePermission:
            return "Microphone permission is required. Enable Meeting Assistant under System Settings > Privacy & Security > Microphone."
        case .noDisplay:
            return "ScreenCaptureKit could not find a display for system-audio capture."
        case .audioConverter(let detail):
            return "Native PCM conversion failed: \(detail)"
        case .systemAudio(let detail):
            return "Screen Recording/System Audio permission is required. Enable Meeting Assistant under System Settings > Privacy & Security > Screen & System Audio Recording. \(detail)"
        }
    }
}

private final class ErrorReporter: @unchecked Sendable {
    private let lock = NSLock()

    func write(_ message: String) {
        lock.lock()
        defer { lock.unlock() }
        let line = message.hasSuffix("\n") ? message : message + "\n"
        if let data = line.data(using: .utf8) {
            try? FileHandle.standardError.write(contentsOf: data)
        }
    }

    func exit(with failure: CaptureFailure, code: Int32) -> Never {
        write(failure.description)
        Darwin.exit(code)
    }
}

private final class PcmWriter: @unchecked Sendable {
    private let targetFormat: AVAudioFormat
    private let output = FileHandle.standardOutput
    private let reporter: ErrorReporter
    private let lock = NSLock()
    private var converter: AVAudioConverter?

    init(sampleRate: Int, channels: Int, reporter: ErrorReporter) throws {
        guard let format = AVAudioFormat(
            commonFormat: .pcmFormatInt16,
            sampleRate: Double(sampleRate),
            channels: AVAudioChannelCount(channels),
            interleaved: true
        ) else {
            throw CaptureFailure.invalidFormat
        }

        targetFormat = format
        self.reporter = reporter
    }

    func write(_ input: AVAudioPCMBuffer) {
        lock.lock()
        defer { lock.unlock() }

        do {
            let outputBuffer = try convert(input)
            let buffers = UnsafeMutableAudioBufferListPointer(outputBuffer.mutableAudioBufferList)
            guard let audioBuffer = buffers.first, let bytes = audioBuffer.mData else {
                return
            }

            let data = Data(bytes: bytes, count: Int(audioBuffer.mDataByteSize))
            try output.write(contentsOf: data)
        } catch {
            reporter.exit(with: .audioConverter(error.localizedDescription), code: 70)
        }
    }

    private func convert(_ input: AVAudioPCMBuffer) throws -> AVAudioPCMBuffer {
        let sourceFormat = input.format
        let activeConverter: AVAudioConverter
        if let converter {
            activeConverter = converter
        } else {
            guard let created = AVAudioConverter(from: sourceFormat, to: targetFormat) else {
                throw CaptureFailure.audioConverter("AVAudioConverter could not convert \(sourceFormat) to \(targetFormat).")
            }
            converter = created
            activeConverter = created
        }

        let ratio = targetFormat.sampleRate / sourceFormat.sampleRate
        let capacity = max(1, AVAudioFrameCount(ceil(Double(input.frameLength) * ratio)) + 32)
        guard let converted = AVAudioPCMBuffer(pcmFormat: targetFormat, frameCapacity: capacity) else {
            throw CaptureFailure.audioConverter("Could not allocate the target PCM buffer.")
        }

        var suppliedInput = false
        var conversionError: NSError?
        let status = activeConverter.convert(to: converted, error: &conversionError) { _, inputStatus in
            if suppliedInput {
                inputStatus.pointee = .noDataNow
                return nil
            }

            suppliedInput = true
            inputStatus.pointee = .haveData
            return input
        }

        if status == .error {
            throw conversionError ?? CaptureFailure.audioConverter("AVAudioConverter returned an unknown error.")
        }

        return converted
    }
}

private final class MicrophoneCapture: @unchecked Sendable {
    private let engine = AVAudioEngine()
    private let writer: PcmWriter

    init(writer: PcmWriter) {
        self.writer = writer
    }

    func run() async throws -> Never {
        let authorized = await withCheckedContinuation { continuation in
            AVCaptureDevice.requestAccess(for: .audio) { granted in
                continuation.resume(returning: granted)
            }
        }
        guard authorized else {
            throw CaptureFailure.microphonePermission
        }

        let input = engine.inputNode
        let inputFormat = input.outputFormat(forBus: 0)
        input.installTap(onBus: 0, bufferSize: 1_024, format: inputFormat) { [writer] buffer, _ in
            writer.write(buffer)
        }
        engine.prepare()
        try engine.start()

        while true {
            try await Task.sleep(for: .seconds(3_600))
        }
    }
}

private final class SystemAudioCapture: NSObject, SCStreamOutput, SCStreamDelegate, @unchecked Sendable {
    private let writer: PcmWriter
    private let reporter: ErrorReporter
    private let sampleQueue = DispatchQueue(label: "meeting-assistant.system-audio")
    private var stream: SCStream?

    init(writer: PcmWriter, reporter: ErrorReporter) {
        self.writer = writer
        self.reporter = reporter
    }

    func run(sampleRate: Int, channels: Int) async throws -> Never {
        let content: SCShareableContent
        do {
            content = try await SCShareableContent.excludingDesktopWindows(
                false,
                onScreenWindowsOnly: false
            )
        } catch {
            throw CaptureFailure.systemAudio(error.localizedDescription)
        }

        guard let display = content.displays.first else {
            throw CaptureFailure.noDisplay
        }

        let filter = SCContentFilter(
            display: display,
            excludingApplications: [],
            exceptingWindows: []
        )
        let configuration = SCStreamConfiguration()
        configuration.capturesAudio = true
        configuration.excludesCurrentProcessAudio = false
        configuration.sampleRate = sampleRate
        configuration.channelCount = channels
        configuration.width = 2
        configuration.height = 2
        configuration.showsCursor = false
        configuration.minimumFrameInterval = CMTime(value: 1, timescale: 1)

        let captureStream = SCStream(filter: filter, configuration: configuration, delegate: self)
        try captureStream.addStreamOutput(self, type: .audio, sampleHandlerQueue: sampleQueue)
        do {
            try await captureStream.startCapture()
        } catch {
            throw CaptureFailure.systemAudio(error.localizedDescription)
        }
        stream = captureStream

        while true {
            try await Task.sleep(for: .seconds(3_600))
        }
    }

    func stream(
        _ stream: SCStream,
        didOutputSampleBuffer sampleBuffer: CMSampleBuffer,
        of outputType: SCStreamOutputType
    ) {
        guard outputType == .audio,
              sampleBuffer.isValid,
              CMSampleBufferDataIsReady(sampleBuffer),
              let description = CMSampleBufferGetFormatDescription(sampleBuffer),
              let streamDescription = CMAudioFormatDescriptionGetStreamBasicDescription(description)
        else {
            return
        }

        let frameCount = AVAudioFrameCount(CMSampleBufferGetNumSamples(sampleBuffer))
        guard frameCount > 0,
              let sourceFormat = AVAudioFormat(streamDescription: streamDescription)
        else {
            return
        }

        var bufferListSize = 0
        var retainedBlockBuffer: CMBlockBuffer?
        var status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(
            sampleBuffer,
            bufferListSizeNeededOut: &bufferListSize,
            bufferListOut: nil,
            bufferListSize: 0,
            blockBufferAllocator: nil,
            blockBufferMemoryAllocator: nil,
            flags: 0,
            blockBufferOut: &retainedBlockBuffer
        )
        guard status == noErr else {
            reporter.exit(
                with: .audioConverter("CoreMedia could not describe system audio (OSStatus \(status))."),
                code: 71
            )
        }

        let rawBufferList = UnsafeMutableRawPointer.allocate(
            byteCount: bufferListSize,
            alignment: MemoryLayout<AudioBufferList>.alignment
        )
        defer { rawBufferList.deallocate() }
        let audioBufferList = rawBufferList.bindMemory(to: AudioBufferList.self, capacity: 1)
        status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(
            sampleBuffer,
            bufferListSizeNeededOut: nil,
            bufferListOut: audioBufferList,
            bufferListSize: bufferListSize,
            blockBufferAllocator: nil,
            blockBufferMemoryAllocator: nil,
            flags: UInt32(kCMSampleBufferFlag_AudioBufferList_Assure16ByteAlignment),
            blockBufferOut: &retainedBlockBuffer
        )
        guard status == noErr,
              let buffer = AVAudioPCMBuffer(
                pcmFormat: sourceFormat,
                bufferListNoCopy: audioBufferList,
                deallocator: nil
              )
        else {
            reporter.exit(
                with: .audioConverter("CoreMedia could not expose system audio (OSStatus \(status))."),
                code: 71
            )
        }

        buffer.frameLength = frameCount
        writer.write(buffer)
    }

    func stream(_ stream: SCStream, didStopWithError error: Error) {
        reporter.exit(with: .systemAudio(error.localizedDescription), code: 72)
    }
}

@main
private struct MacOsMeetingAudioCapture {
    static func main() async {
        let reporter = ErrorReporter()
        do {
            let arguments = CommandLine.arguments
            guard arguments.count == 4,
                  let sampleRate = Int(arguments[2]),
                  let channels = Int(arguments[3])
            else {
                throw CaptureFailure.usage
            }
            guard sampleRate > 0, channels == 1 || channels == 2 else {
                throw CaptureFailure.invalidFormat
            }

            let writer = try PcmWriter(sampleRate: sampleRate, channels: channels, reporter: reporter)
            switch arguments[1] {
            case "microphone":
                try await MicrophoneCapture(writer: writer).run()
            case "system":
                try await SystemAudioCapture(writer: writer, reporter: reporter)
                    .run(sampleRate: sampleRate, channels: channels)
            default:
                throw CaptureFailure.usage
            }
        } catch let failure as CaptureFailure {
            reporter.exit(with: failure, code: 64)
        } catch {
            reporter.write(error.localizedDescription)
            Darwin.exit(1)
        }
    }
}