import ReplayKit
import AVFoundation

/// 系统录屏广播的采样处理器 —— 多媒体翻译在 iOS 上拿到「其他 App 正在播放的
/// 声音」的**唯一合规途径**。
///
/// ## 为什么必须是一个独立的 extension
/// iOS 不允许一个 App 读另一个 App 的音频。ReplayKit 的 Broadcast Upload
/// Extension 是系统留的唯一口子：用户从系统录屏面板里选中本扩展、开始广播，
/// 系统才会把全局的 `audioApp`（各 App 播放的混音）和 `audioMic`（麦克风）
/// 送进这个扩展进程。主 App 自己是拿不到的。
///
/// ## ⚠️ 只取 audioApp，丢掉 audioMic 和 video
/// - `video` 完全不要：我们只做翻译，收视频帧纯属浪费内存，而广播扩展的
///   内存上限是 **50MB**，超了系统直接杀进程（表现为广播莫名其妙停掉）。
/// - `audioMic` 也不要：混进去会把用户和旁人的说话声一起送去识别，
///   翻译结果会串成一团。
///
/// ## ⚠️ 采样率必须在这里转
/// 系统给的是 44.1k/48k、立体声、Float32；Azure ASR 要 16k/单声道/Int16。
/// 不转的话推进去的是一堆噪声，识别恒为空，且不报任何错。
class SampleHandler: RPBroadcastSampleHandler {
    private let sender = BroadcastAudioSender()

    /// 输出格式：16kHz / 单声道 / Int16 交错，与 azureAsrHelper.audioStream 的要求一致
    private let outputFormat = AVAudioFormat(
        commonFormat: .pcmFormatInt16,
        sampleRate: 16000,
        channels: 1,
        interleaved: true
    )!

    /// 输入格式会变（不同 App 的输出不一样），换了就重建 converter
    private var converter: AVAudioConverter?
    private var converterInputFormat: AVAudioFormat?

    /// 主 App 收工时自动结束广播用的心跳。
    ///
    /// ⚠️ 不能只靠 processSampleBuffer 里那次检查：视频停了或者没有任何 App
    /// 在出声时，`audioApp` 帧可能长时间不来，那样就永远发现不了主 App 已经走了，
    /// 屏幕上那条红色录制条会一直挂着。
    private var watchdog: Timer?

    override func broadcastStarted(withSetupInfo setupInfo: [String: NSObject]?) {
        NSLog("[BroadcastExtension] broadcastStarted")
        sender.onPeerGone = { [weak self] in
            self?.finishBecausePeerGone()
        }
        DispatchQueue.main.async { [weak self] in
            self?.watchdog = Timer.scheduledTimer(withTimeInterval: 1.0, repeats: true) { [weak self] _ in
                self?.sender.checkPeerAlive()
            }
        }
    }

    override func broadcastFinished() {
        NSLog("[BroadcastExtension] broadcastFinished")
        watchdog?.invalidate()
        watchdog = nil
    }

    /// 主 App 已经停止多媒体翻译 → 结束这次广播。
    ///
    /// ⚠️ `finishBroadcastWithError` 是扩展**唯一**能终止广播的入口，而且必须
    /// 传一个 NSError；系统会把 localizedDescription 弹给用户看。所以这里给的
    /// 不是错误信息而是一句正常的说明，否则用户会看到一个莫名其妙的报错。
    private func finishBecausePeerGone() {
        watchdog?.invalidate()
        watchdog = nil
        NSLog("[BroadcastExtension] 主 App 已收工，结束广播")
        let reason = NSError(
            domain: "com.yimai.eaimar.broadcast",
            code: 0,
            userInfo: [NSLocalizedDescriptionKey: "多媒体翻译已结束，屏幕直播已自动关闭"]
        )
        DispatchQueue.main.async { [weak self] in
            self?.finishBroadcastWithError(reason)
        }
    }

    override func processSampleBuffer(_ sampleBuffer: CMSampleBuffer,
                                      with sampleBufferType: RPSampleBufferType) {
        guard sampleBufferType == .audioApp else { return }
        guard CMSampleBufferDataIsReady(sampleBuffer) else { return }
        guard let pcm = convertToTargetFormat(sampleBuffer) else { return }
        sender.send(pcm)
    }

    // MARK: - 格式转换

    private func convertToTargetFormat(_ sampleBuffer: CMSampleBuffer) -> Data? {
        guard let formatDesc = CMSampleBufferGetFormatDescription(sampleBuffer),
              let asbdPtr = CMAudioFormatDescriptionGetStreamBasicDescription(formatDesc) else {
            return nil
        }
        let inFormat = AVAudioFormat(streamDescription: asbdPtr)
        guard let inFormat = inFormat else { return nil }

        if converter == nil || converterInputFormat != inFormat {
            converter = AVAudioConverter(from: inFormat, to: outputFormat)
            converterInputFormat = inFormat
        }
        guard let converter = converter else { return nil }

        guard let inBuffer = makeInputBuffer(sampleBuffer, format: inFormat) else { return nil }

        // 输出容量按采样率比例算，再留一点余量：转换器允许一次多吐一点。
        // 算小了会得到 .endOfStream 并丢样本，表现是识别断断续续。
        let ratio = outputFormat.sampleRate / inFormat.sampleRate
        let capacity = AVAudioFrameCount(Double(inBuffer.frameLength) * ratio) + 1024
        guard let outBuffer = AVAudioPCMBuffer(pcmFormat: outputFormat,
                                               frameCapacity: capacity) else {
            return nil
        }

        var consumed = false
        var error: NSError?
        converter.convert(to: outBuffer, error: &error) { _, status in
            if consumed {
                status.pointee = .noDataNow
                return nil
            }
            consumed = true
            status.pointee = .haveData
            return inBuffer
        }
        if let error = error {
            NSLog("[BroadcastExtension] 音频转换失败: %@", error.localizedDescription)
            return nil
        }
        guard outBuffer.frameLength > 0,
              let channelData = outBuffer.int16ChannelData else { return nil }
        let byteCount = Int(outBuffer.frameLength) * MemoryLayout<Int16>.size
        return Data(bytes: channelData[0], count: byteCount)
    }

    /// CMSampleBuffer → AVAudioPCMBuffer（**深拷贝**）。
    ///
    /// ⚠️ `CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer` 并不会把样本
    /// 拷进你给的 buffer，它只是把 AudioBufferList 里的 `mData` **指向 blockBuffer
    /// 持有的内存**。blockBuffer 是个局部变量，函数一返回就被 ARC 释放，
    /// 把那份指针交出去就是野指针 —— 表现是随机崩溃，而且多半崩在转换器里，
    /// 看上去像 AVAudioConverter 的锅。所以这里必须在 blockBuffer 活着的时候
    /// 就把数据 memcpy 出来。
    ///
    /// ⚠️ AudioBufferList 的 buffer 个数要按格式算：非交错多声道是「每声道一个
    /// buffer」，交错才是一个。写死成 1 时，ReplayKit 给的非交错立体声会让
    /// 上面那个调用直接返回错误码，一帧都拿不到。
    private func makeInputBuffer(_ sampleBuffer: CMSampleBuffer,
                                 format: AVAudioFormat) -> AVAudioPCMBuffer? {
        let frameCount = AVAudioFrameCount(CMSampleBufferGetNumSamples(sampleBuffer))
        guard frameCount > 0 else { return nil }

        let maxBuffers = format.isInterleaved ? 1 : Int(format.channelCount)
        let srcList = AudioBufferList.allocate(maximumBuffers: maxBuffers)
        defer { free(srcList.unsafeMutablePointer) }

        var blockBuffer: CMBlockBuffer?
        let status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(
            sampleBuffer,
            bufferListSizeNeededOut: nil,
            bufferListOut: srcList.unsafeMutablePointer,
            bufferListSize: AudioBufferList.sizeInBytes(maximumBuffers: maxBuffers),
            blockBufferAllocator: kCFAllocatorDefault,
            blockBufferMemoryAllocator: kCFAllocatorDefault,
            flags: kCMSampleBufferFlag_AudioBufferList_Assure16ByteAlignment,
            blockBufferOut: &blockBuffer
        )
        guard status == noErr, blockBuffer != nil else { return nil }

        guard let out = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else {
            return nil
        }
        out.frameLength = frameCount

        // withExtendedLifetime：把 blockBuffer 的生命周期显式压到拷贝之后。
        // 光靠一个 `_ = blockBuffer` 是不保证的，优化器可以提前释放。
        let copied: Bool = withExtendedLifetime(blockBuffer) { () -> Bool in
            let dstList = UnsafeMutableAudioBufferListPointer(out.mutableAudioBufferList)
            guard dstList.count == srcList.count else { return false }
            for i in 0..<srcList.count {
                guard let src = srcList[i].mData, let dst = dstList[i].mData else { return false }
                let n = min(Int(srcList[i].mDataByteSize), Int(dstList[i].mDataByteSize))
                memcpy(dst, src, n)
                dstList[i].mDataByteSize = UInt32(n)
            }
            return true
        }
        return copied ? out : nil
    }
}