import AVFoundation
import ScreenCaptureKit

// ScreenAudioCapture — macOS system/desktop audio capture for the SCREEN_AUDIO stream.
//
// The macOS analog of the Windows WASAPI loopback path (docs/voice.md §9). ScreenCaptureKit
// (macOS 13+) captures whatever the system is playing; we convert each audio CMSampleBuffer
// (Float32) → int16 interleaved and push 20 ms frames (960 samples/channel @ 48 kHz) into the
// core via `vc_stream_feed_pcm` (exposed as `VoiceCatClient.feedPcm`). The core then runs the
// same Opus-encode → media-AEAD → UDP path as any other stream — only the *source* is
// platform-specific (architecture.md §4).
//
// Audio-only: we request a 2×2 video plane at 1 fps purely because SCStream needs a video
// configuration, and we never add a `.screen` output — only `.audio`. `excludesCurrentProcess
// Audio` prevents the self-echo loop of re-capturing our own incoming voice mix.
//
// `feedPcm` is thread-safe (any thread), so we forward straight from the sample-handler queue.
final class ScreenAudioCapture: NSObject, SCStreamOutput, SCStreamDelegate {

    /// Receives a full 20 ms frame: (interleaved int16 PCM, samplesPerChannel = 960, channels).
    typealias PcmHandler = (UnsafePointer<Int16>, Int, UInt32) -> Void

    enum CaptureError: Error { case noDisplay }

    private static let frameSamplesPerChannel = 960  // 20 ms @ 48 kHz

    private let onPcm: PcmHandler
    private let channels: Int  // 1 (mono) or 2 (stereo interleaved), matches the stream's mode
    private let sampleQueue = DispatchQueue(label: "cat.voice.screenaudio.samples")
    private var stream: SCStream?

    /// Interleaved int16 carry-over between callbacks (ScreenCaptureKit buffers don't align to
    /// 20 ms), drained in whole `frameSamplesPerChannel * channels` chunks. Only touched on
    /// `sampleQueue`.
    private var pending: [Int16] = []

    init(channels: UInt32, onPcm: @escaping PcmHandler) {
        self.channels = max(1, min(2, Int(channels)))
        self.onPcm = onPcm
        super.init()
    }

    /// Begin capture. Throws if Screen Recording permission is denied (the first
    /// `SCShareableContent.current` access is what surfaces the TCC prompt) or no display exists.
    func start() async throws {
        let content = try await SCShareableContent.current
        guard let display = content.displays.first else { throw CaptureError.noDisplay }

        let filter = SCContentFilter(display: display, excludingWindows: [])

        let config = SCStreamConfiguration()
        config.capturesAudio = true
        config.excludesCurrentProcessAudio = true
        config.sampleRate = 48000
        config.channelCount = channels
        // SCStream requires a video config even when we only consume audio — keep it minimal.
        config.width = 2
        config.height = 2
        config.minimumFrameInterval = CMTime(value: 1, timescale: 1)  // ~1 fps
        config.queueDepth = 6

        let stream = SCStream(filter: filter, configuration: config, delegate: self)
        try stream.addStreamOutput(self, type: .audio, sampleHandlerQueue: sampleQueue)
        try await stream.startCapture()
        self.stream = stream
    }

    /// Stop capture and release the stream. Safe to call multiple times.
    func stop() {
        guard let stream else { return }
        self.stream = nil
        Task { try? await stream.stopCapture() }
    }

    // MARK: - SCStreamOutput

    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer,
                of type: SCStreamOutputType) {
        guard type == .audio, CMSampleBufferDataIsReady(sampleBuffer) else { return }
        guard let fmt = sampleBuffer.formatDescription,
              let asbd = fmt.audioStreamBasicDescription else { return }
        // ScreenCaptureKit always delivers Float32 PCM; bail on anything unexpected.
        guard asbd.mFormatFlags & kAudioFormatFlagIsFloat != 0 else { return }

        let nonInterleaved = asbd.mFormatFlags & kAudioFormatFlagIsNonInterleaved != 0
        let srcChannels = max(1, Int(asbd.mChannelsPerFrame))

        try? sampleBuffer.withAudioBufferList { ablPtr, _ in
            convert(ablPtr, nonInterleaved: nonInterleaved, srcChannels: srcChannels)
        }
    }

    // MARK: - Conversion (called on sampleQueue)

    private func convert(_ abl: UnsafeMutableAudioBufferListPointer,
                         nonInterleaved: Bool, srcChannels: Int) {
        guard let first = abl.first, first.mData != nil else { return }

        let out = channels
        var interleaved: [Int16]

        if nonInterleaved {
            // One buffer per source channel; each is `frames` Float32 samples.
            let frames = Int(first.mDataByteSize) / MemoryLayout<Float>.size
            if frames == 0 { return }
            let ch0 = first.mData!.assumingMemoryBound(to: Float.self)
            let ch1: UnsafePointer<Float>? = (abl.count > 1)
                ? UnsafePointer(abl[1].mData!.assumingMemoryBound(to: Float.self)) : nil
            interleaved = [Int16](repeating: 0, count: frames * out)
            for i in 0..<frames {
                let l = ch0[i]
                let r = ch1?[i] ?? l
                if out == 2 {
                    interleaved[i * 2]     = Self.f2i(l)
                    interleaved[i * 2 + 1] = Self.f2i(srcChannels >= 2 ? r : l)
                } else {
                    interleaved[i] = Self.f2i(srcChannels >= 2 ? (l + r) * 0.5 : l)
                }
            }
        } else {
            // Single interleaved Float32 buffer, srcChannels wide.
            let total = Int(first.mDataByteSize) / MemoryLayout<Float>.size
            let frames = total / srcChannels
            if frames == 0 { return }
            let src = first.mData!.assumingMemoryBound(to: Float.self)
            interleaved = [Int16](repeating: 0, count: frames * out)
            for i in 0..<frames {
                let l = src[i * srcChannels]
                let r = srcChannels >= 2 ? src[i * srcChannels + 1] : l
                if out == 2 {
                    interleaved[i * 2]     = Self.f2i(l)
                    interleaved[i * 2 + 1] = Self.f2i(r)
                } else {
                    interleaved[i] = Self.f2i(srcChannels >= 2 ? (l + r) * 0.5 : l)
                }
            }
        }

        emit(interleaved)
    }

    /// Accumulate interleaved int16 and fire `onPcm` for every whole 20 ms frame.
    private func emit(_ interleaved: [Int16]) {
        pending.append(contentsOf: interleaved)
        let full = Self.frameSamplesPerChannel * channels
        while pending.count >= full {
            pending.withUnsafeBufferPointer { buf in
                onPcm(buf.baseAddress!, Self.frameSamplesPerChannel, UInt32(channels))
            }
            pending.removeFirst(full)
        }
    }

    private static func f2i(_ f: Float) -> Int16 {
        let v = max(-1.0, min(1.0, f)) * 32767.0
        return Int16(v.rounded())
    }
}