import AVFoundation
import ScreenCaptureKit

// Which apps' audio the SCREEN_AUDIO stream captures. ScreenCaptureKit filters audio at the
// *application* level (not per-window), so the selection is expressed as bundle IDs. The
// picker UI (ScreenSharePickerSheet) produces a `ScreenAudioSelection`; `start()` turns it
// into the matching `SCContentFilter`.
enum ScreenAudioScope: Equatable {
    case entireDesktop          // whole display — the original behaviour
    case onlyApps([String])     // capture only these bundle IDs
    case allExcept([String])    // capture everything except these bundle IDs
}

struct ScreenAudioSelection: Equatable {
    var scope: ScreenAudioScope = .entireDesktop
    /// Drop the macOS screen-reader (VoiceOver) speech from the shared mix. Meaningful for
    /// `.entireDesktop`/`.allExcept`; for `.onlyApps` the screen reader is already excluded.
    var excludeScreenReader: Bool = false

    static let `default` = ScreenAudioSelection()
}

// 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 selection: ScreenAudioSelection
    private let sampleQueue = DispatchQueue(label: "cat.voice.screenaudio.samples")
    private var stream: SCStream?

    /// Bundle IDs whose audio carries the macOS screen-reader speech. VoiceOver itself plus the
    /// speech-synthesis daemon that actually renders the spoken audio — the speech is usually
    /// emitted by the daemon, not the VoiceOver app, so we exclude whichever are running.
    static let screenReaderBundleIDs: Set<String> = [
        "com.apple.VoiceOver",
        "com.apple.VoiceOver4",
        "com.apple.speech.speechsynthesisd",
    ]

    /// 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, selection: ScreenAudioSelection, onPcm: @escaping PcmHandler) {
        self.channels = max(1, min(2, Int(channels)))
        self.selection = selection
        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 = Self.makeFilter(selection: selection, display: display,
                                     apps: content.applications)

        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
    }

    /// Turn a `ScreenAudioSelection` into an `SCContentFilter` against the running apps.
    /// ScreenCaptureKit filters audio per application, so we map bundle IDs → SCRunningApplication.
    private static func makeFilter(selection: ScreenAudioSelection, display: SCDisplay,
                                   apps: [SCRunningApplication]) -> SCContentFilter {
        func appsMatching(_ ids: Set<String>) -> [SCRunningApplication] {
            apps.filter { ids.contains($0.bundleIdentifier) }
        }

        switch selection.scope {
        case .onlyApps(let bundleIDs):
            // Include-only already excludes everything else (the screen reader included), so the
            // excludeScreenReader flag is moot in this mode.
            return SCContentFilter(display: display,
                                   including: appsMatching(Set(bundleIDs)),
                                   exceptingWindows: [])

        case .allExcept(let bundleIDs):
            var ids = Set(bundleIDs)
            if selection.excludeScreenReader { ids.formUnion(screenReaderBundleIDs) }
            return SCContentFilter(display: display,
                                   excludingApplications: appsMatching(ids),
                                   exceptingWindows: [])

        case .entireDesktop:
            if selection.excludeScreenReader {
                return SCContentFilter(display: display,
                                       excludingApplications: appsMatching(screenReaderBundleIDs),
                                       exceptingWindows: [])
            }
            return SCContentFilter(display: display, excludingWindows: [])
        }
    }

    /// 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())
    }
}
