import Foundation

// BroadcastAudioPump — host side of iOS screen-audio sharing (docs/voice.md §9, iOS detail).
//
// Drains the App Group shared-memory ring written by the broadcast upload extension and hands
// whole 20 ms frames to a feed closure (which calls VoiceCatClient.feedPcm). The host app owns
// the SCREEN_AUDIO stream, so screen audio appears as a second stream of the SAME user — exactly
// like the Windows/macOS desktop-audio share, and a single session (no separate connection).
//
// It reacts to the extension's Darwin notifications for prompt start/stop, and the drain timer
// also watches the ring's active flag as a safety net if a notification is missed. The ring
// always carries stereo; we downmix to mono when the stream's effective config is mono.
//
// Not @MainActor: the drain runs on a background queue (feedPcm is thread-safe). The start/stop
// callbacks are dispatched to main so they can drive @MainActor SessionState.
final class BroadcastAudioPump {

    /// The host should start the SCREEN_AUDIO stream (a broadcast became active).
    var onBroadcastStarted: (() -> Void)?
    /// The host should stop the SCREEN_AUDIO stream (the broadcast ended).
    var onBroadcastFinished: (() -> Void)?

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

    private let queue = DispatchQueue(label: "cat.voice.broadcast.pump")
    private var ring: BroadcastAudioRing?
    private var timer: DispatchSourceTimer?
    private var feed: ((UnsafePointer<Int16>, Int, UInt32) -> Void)?
    private var streamChannels = 1
    private var ringChannels = 2
    private var pending: [Int16] = []                 // interleaved at ringChannels width
    private var scratch = [Int16](repeating: 0, count: 8192)
    private var started = false

    // MARK: - Lifecycle (host connect/disconnect)

    func start() {
        guard !started else { return }
        ring = try? BroadcastAudioRing()
        started = true
        registerDarwin()
        // Host reconnected while a broadcast is still running — pick it up.
        if ring?.isActive == true { onBroadcastStarted?() }
    }

    func stop() {
        guard started else { return }
        started = false
        unregisterDarwin()
        endFeeding()
        ring = nil
    }

    // MARK: - Feeding (driven by the host once the stream is live)

    /// Begin draining the ring into `feed`. Called after the SCREEN_AUDIO stream's
    /// `.streamStarted` event, when its effective channel count is known.
    func beginFeeding(streamChannels: UInt32,
                      feed: @escaping (UnsafePointer<Int16>, Int, UInt32) -> Void) {
        queue.async {
            self.streamChannels = max(1, min(2, Int(streamChannels)))
            self.ringChannels = max(1, Int(self.ring?.channels ?? 2))
            self.feed = feed
            self.pending.removeAll(keepingCapacity: true)
            self.ring?.drainStale()  // discard pre-roll buffered before we were ready
            self.startTimer()
        }
    }

    func endFeeding() {
        queue.async {
            self.timer?.cancel()
            self.timer = nil
            self.feed = nil
            self.pending.removeAll(keepingCapacity: true)
        }
    }

    // MARK: - Drain

    private func startTimer() {
        let t = DispatchSource.makeTimerSource(queue: queue)
        t.schedule(deadline: .now(), repeating: .milliseconds(10), leeway: .milliseconds(2))
        t.setEventHandler { [weak self] in self?.drain() }
        timer = t
        t.resume()
    }

    private func drain() {
        guard let ring, let feed else { return }
        // Safety net: broadcast ended but we missed the Darwin note.
        if !ring.isActive {
            DispatchQueue.main.async { [weak self] in self?.onBroadcastFinished?() }
            return
        }
        while true {
            let got = scratch.withUnsafeMutableBufferPointer { ring.read(into: $0) }
            if got == 0 { break }
            pending.append(contentsOf: scratch[0..<got])
            if got < scratch.count { break }
        }
        emitFrames(feed)
    }

    private func emitFrames(_ feed: (UnsafePointer<Int16>, Int, UInt32) -> Void) {
        let n = Self.frameSamplesPerChannel
        let rc = ringChannels
        let sc = streamChannels
        let inFrame = n * rc
        while pending.count >= inFrame {
            if sc == rc {
                pending.withUnsafeBufferPointer { feed($0.baseAddress!, n, UInt32(sc)) }
            } else if sc == 1 && rc == 2 {
                var mono = [Int16](repeating: 0, count: n)
                pending.withUnsafeBufferPointer { buf in
                    let p = buf.baseAddress!
                    for i in 0..<n { mono[i] = Int16((Int(p[i * 2]) + Int(p[i * 2 + 1])) / 2) }
                }
                mono.withUnsafeBufferPointer { feed($0.baseAddress!, n, 1) }
            } else if sc == 2 && rc == 1 {
                var stereo = [Int16](repeating: 0, count: n * 2)
                pending.withUnsafeBufferPointer { buf in
                    let p = buf.baseAddress!
                    for i in 0..<n { stereo[i * 2] = p[i]; stereo[i * 2 + 1] = p[i] }
                }
                stereo.withUnsafeBufferPointer { feed($0.baseAddress!, n, 2) }
            }
            pending.removeFirst(inFrame)
        }
    }

    // MARK: - Darwin notifications

    private func registerDarwin() {
        let observer = Unmanaged.passUnretained(self).toOpaque()
        let center = CFNotificationCenterGetDarwinNotifyCenter()
        let callback: CFNotificationCallback = { _, observer, name, _, _ in
            guard let observer, let name else { return }
            let pump = Unmanaged<BroadcastAudioPump>.fromOpaque(observer).takeUnretainedValue()
            let raw = name.rawValue as String
            DispatchQueue.main.async {
                if raw == BroadcastNotification.started { pump.onBroadcastStarted?() }
                else if raw == BroadcastNotification.finished { pump.onBroadcastFinished?() }
            }
        }
        CFNotificationCenterAddObserver(center, observer, callback,
            BroadcastNotification.started as CFString, nil, .deliverImmediately)
        CFNotificationCenterAddObserver(center, observer, callback,
            BroadcastNotification.finished as CFString, nil, .deliverImmediately)
    }

    private func unregisterDarwin() {
        CFNotificationCenterRemoveEveryObserver(
            CFNotificationCenterGetDarwinNotifyCenter(),
            Unmanaged.passUnretained(self).toOpaque())
    }

    deinit { if started { unregisterDarwin() } }
}