import AudioToolbox
import AVFAudio
import CoreGraphics
import CoreMedia
import Foundation
import ScreenCaptureKit

@available(macOS 13.0, *)
final class LiveDeskSystemAudioCapture: NSObject, SCStreamOutput, SCStreamDelegate {
    private let sampleRate: Int
    private let channelCount: Int
    private let queue = DispatchQueue(label: "dev.livedesk.screencapturekit.audio", qos: .userInteractive)
    private let writeLock = NSLock()
    private var packetHeader = Data(count: 16)
    private var stream: SCStream?
    private var audioFrameCount = 0
    private var pcmMaterializationCount = 0
    private let packetPayloadCopyCount = 0
    private var loggedFormat = false

    init(sampleRate: Int, channelCount: Int) {
        self.sampleRate = max(8_000, min(192_000, sampleRate))
        self.channelCount = max(1, min(2, channelCount))
        super.init()
    }

    func start() async throws {
        let content = try await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: true)
        guard !content.displays.isEmpty else {
            throw LiveDeskSystemAudioError.noDisplays
        }

        let mainDisplayID = CGMainDisplayID()
        let display = content.displays.sorted { left, right in
            let leftIsMain = left.displayID == mainDisplayID
            let rightIsMain = right.displayID == mainDisplayID
            if leftIsMain != rightIsMain {
                return leftIsMain
            }
            return left.displayID < right.displayID
        }[0]

        let configuration = SCStreamConfiguration()
        configuration.width = 2
        configuration.height = 2
        configuration.minimumFrameInterval = CMTime(value: 1, timescale: 1)
        // Keep ScreenCaptureKit's own retained surface budget at one. The
        // C# handoff also owns exactly one latest-wins callback, so an
        // upstream eight-frame reservoir would only preserve stale audio and
        // inflate WindowServer/helper memory before our currentness gate.
        configuration.queueDepth = 1
        configuration.capturesAudio = true
        configuration.sampleRate = sampleRate
        configuration.channelCount = channelCount
        configuration.excludesCurrentProcessAudio = true

        let stream = SCStream(filter: SCContentFilter(display: display, excludingWindows: []), configuration: configuration, delegate: self)
        try stream.addStreamOutput(self, type: .audio, sampleHandlerQueue: queue)
        self.stream = stream
        try await stream.startCapture()
        fputs("LiveDeskSCKAudio started displayID=\(display.displayID) sampleRate=\(sampleRate) channels=\(channelCount)\n", stderr)

        while !Task.isCancelled {
            try await Task.sleep(nanoseconds: 1_000_000_000)
        }
    }

    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) {
        guard type == .audio, sampleBuffer.isValid, CMSampleBufferDataIsReady(sampleBuffer) else {
            return
        }
        writePcm(sampleBuffer)
    }

    func stream(_ stream: SCStream, didStopWithError error: Error) {
        fputs("LiveDeskSCKAudio stream stopped: \(error)\n", stderr)
        exit(3)
    }

    private func writePcm(_ sampleBuffer: CMSampleBuffer) {
        guard let description = sampleBuffer.formatDescription?.audioStreamBasicDescription else {
            return
        }

        let channels = Int(description.mChannelsPerFrame)
        let frames = Int(CMSampleBufferGetNumSamples(sampleBuffer))
        let isFloat32 = description.mFormatID == kAudioFormatLinearPCM
            && description.mBitsPerChannel == 32
            && (description.mFormatFlags & kAudioFormatFlagIsFloat) != 0
        let isInterleaved = (description.mFormatFlags & kAudioFormatFlagIsNonInterleaved) == 0
        guard isFloat32, frames > 0, channels > 0, channels <= 8 else {
            if !loggedFormat {
                loggedFormat = true
                fputs("LiveDeskSCKAudio unsupported PCM format id=\(description.mFormatID) flags=\(description.mFormatFlags) bits=\(description.mBitsPerChannel) channels=\(channels)\n", stderr)
            }
            return
        }

        do {
            try sampleBuffer.withAudioBufferList { audioBufferList, _ in
                let bytesPerSample = MemoryLayout<Float>.size
                let bytesPerChannel = frames * bytesPerSample
                let outputBytes = bytesPerChannel * channels
                guard audioBufferList.count > 0, outputBytes > 0 else {
                    return
                }

                let output: Data
                if isInterleaved {
                    let sourceBuffer = audioBufferList[0]
                    guard let source = sourceBuffer.mData, Int(sourceBuffer.mDataByteSize) >= outputBytes else {
                        return
                    }
                    output = Data(bytes: source, count: outputBytes)
                } else {
                    guard audioBufferList.count >= channels else {
                        return
                    }
                    var interleaved = Data(count: outputBytes)
                    var copied = true
                    interleaved.withUnsafeMutableBytes { rawBuffer in
                        guard let destination = rawBuffer.baseAddress?.assumingMemoryBound(to: Float.self) else {
                            copied = false
                            return
                        }
                        for channelIndex in 0..<channels {
                            let sourceBuffer = audioBufferList[channelIndex]
                            guard let source = sourceBuffer.mData,
                                  Int(sourceBuffer.mDataByteSize) >= bytesPerChannel else {
                                copied = false
                                return
                            }
                            let sourceSamples = source.assumingMemoryBound(to: Float.self)
                            for frameIndex in 0..<frames {
                                destination[frameIndex * channels + channelIndex] = sourceSamples[frameIndex]
                            }
                        }
                    }
                    guard copied else {
                        return
                    }
                    output = interleaved
                }

                if !loggedFormat {
                    loggedFormat = true
                    fputs("LiveDeskSCKAudio format sampleRate=\(Int(description.mSampleRate.rounded())) channels=\(channels) frames=\(frames) interleaved=\(isInterleaved ? 1 : 0)\n", stderr)
                }
                writePacket(output, sampleRate: Int(description.mSampleRate.rounded()), channels: channels)
                audioFrameCount += 1
                pcmMaterializationCount += 1
                if audioFrameCount == 1 || audioFrameCount % 250 == 0 {
                    fputs(
                        "LiveDeskSCKAudio frames=\(audioFrameCount) bytes=\(output.count) "
                        + "pcmMaterializations=\(pcmMaterializationCount) "
                        + "packetPayloadCopies=\(packetPayloadCopyCount)\n",
                        stderr)
                }
            }
        } catch {
            fputs("LiveDeskSCKAudio audio buffer error: \(error)\n", stderr)
        }
    }

    private func writePacket(_ pcm: Data, sampleRate: Int, channels: Int) {
        guard sampleRate > 0, channels > 0, pcm.count <= Int(UInt32.max) else {
            return
        }
        writeLock.lock()
        packetHeader[0] = 0x4c
        packetHeader[1] = 0x44
        packetHeader[2] = 0x41
        packetHeader[3] = 0x55
        writeLittleEndian(UInt16(1), at: 4, to: &packetHeader)
        writeLittleEndian(
            UInt16(min(channels, Int(UInt16.max))),
            at: 6,
            to: &packetHeader)
        writeLittleEndian(
            UInt32(min(sampleRate, Int(UInt32.max))),
            at: 8,
            to: &packetHeader)
        writeLittleEndian(
            UInt32(pcm.count),
            at: 12,
            to: &packetHeader)
        // The ScreenCaptureKit callback requires one stable PCM owner because
        // the AudioBufferList lifetime ends on return. Write that owner
        // directly after the reusable header instead of allocating a second
        // packet Data and copying the full PCM payload into it.
        FileHandle.standardOutput.write(packetHeader)
        FileHandle.standardOutput.write(pcm)
        writeLock.unlock()
    }

    private func writeLittleEndian<T: FixedWidthInteger>(
        _ value: T,
        at offset: Int,
        to data: inout Data
    ) {
        let littleEndian = value.littleEndian
        for byteOffset in 0..<MemoryLayout<T>.size {
            data[offset + byteOffset] = UInt8(
                truncatingIfNeeded: littleEndian >> (byteOffset * 8))
        }
    }
}

enum LiveDeskSystemAudioError: Error {
    case noDisplays
}

func intArg(_ name: String, default defaultValue: Int) -> Int {
    let args = CommandLine.arguments
    guard let index = args.firstIndex(of: name), index + 1 < args.count else {
        return defaultValue
    }
    return Int(args[index + 1]) ?? defaultValue
}

@main
struct LiveDeskSystemAudioMain {
    static func main() async {
        guard #available(macOS 13.0, *) else {
            fputs("LiveDeskSCKAudio requires macOS 13 or later.\n", stderr)
            exit(2)
        }
        guard CGPreflightScreenCaptureAccess() else {
            fputs("LiveDeskSCKAudio Screen & System Audio Recording permission is not granted.\n", stderr)
            exit(10)
        }

        let sampleRate = intArg("--sample-rate", default: 48_000)
        let channels = intArg("--channels", default: 2)
        do {
            let capture = LiveDeskSystemAudioCapture(sampleRate: sampleRate, channelCount: channels)
            try await capture.start()
        } catch {
            fputs("LiveDeskSCKAudio failed: \(error)\n", stderr)
            exit(1)
        }
    }
}