video_decode_plugin/ios/Classes/VideoDecoder.swift

import Foundation
import VideoToolbox
import AVFoundation

/// 视频解码器，基于VideoToolbox实现H264/H265硬件解码，输出CVPixelBuffer
class VideoDecoder {
    enum CodecType: String {
        case h264 = "h264"
        case h265 = "h265"
        var codecType: CMVideoCodecType {
            switch self {
            case .h264: return kCMVideoCodecType_H264
            case .h265: return kCMVideoCodecType_HEVC
            }
        }
    }

    // ====== 关键成员变量注释 ======
    /// 解码会话对象
    private var decompressionSession: VTDecompressionSession?
    /// 视频格式描述
    private var formatDesc: CMVideoFormatDescription?
    /// 视频宽度
    private let width: Int
    /// 视频高度
    private let height: Int
    /// 编码类型（H264/H265）
    private let codecType: CodecType
    /// 解码线程队列
    private let decodeQueue = DispatchQueue(label: "video_decode_plugin.decode.queue")
    /// 解码会话是否已准备好
    private var isSessionReady = false
    /// 最近一次I帧序号
    private var lastIFrameSeq: Int?
    /// 已处理帧序号集合
    private var frameSeqSet = Set<Int>()
    /// 最大允许延迟（毫秒）
    private let maxAllowedDelayMs: Int64 = 750 // 与Android端保持一致
    /// 时间戳基准
    private var timestampBaseMs: Int64?
    /// 首帧相对时间戳
    private var firstFrameRelativeTimestamp: Int64?

    // ====== 缓冲区与渲染相关成员 ======
    /// 输入缓冲区（待解码帧队列），线程安全
    private let inputQueue = DispatchQueue(label: "video_decode_plugin.input.queue", attributes: .concurrent)
    private var inputBuffer: [(frameData: Data, frameType: Int, timestamp: Int64, frameSeq: Int, refIFrameSeq: Int?, sps: Data?, pps: Data?)] = []
    private let inputBufferSemaphore = DispatchSemaphore(value: 1)
    private let inputBufferMaxCount = 100 // 与Android端保持一致
    /// 输出缓冲区（解码后帧队列），线程安全
    private let outputQueue = DispatchQueue(label: "video_decode_plugin.output.queue", attributes: .concurrent)
    private var outputBuffer: [(pixelBuffer: CVPixelBuffer, timestamp: Int64)] = []
    private let outputBufferSemaphore = DispatchSemaphore(value: 1)
    private let outputBufferMaxCount = 100 // 与Android端保持一致
    /// 渲染定时器
    private var renderTimer: DispatchSourceTimer?
    /// 渲染线程运行标志
    private var renderThreadRunning = false
    /// 首次渲染回调标志
    private var hasNotifiedFlutter = false
    /// 当前渲染帧率
    private var renderFps: Int = 20 // 与Android端保持一致
    /// 渲染间隔（毫秒）
    private var renderIntervalMs: Int64 = 0
    /// 允许的渲染时间抖动范围（毫秒）
    private let renderJitterMs: Int64 = 2
    /// 上次渲染时间
    private var lastRenderTimeMs: Int64 = 0
    /// 渲染启动标志（低水位控制）
    private var renderStarted = false
    
    // ====== 新增：帧重排序与丢帧相关成员 ======
    /// 帧重排序缓冲区
    private var reorderBuffer: [Int: (frameData: Data, frameType: Int, timestamp: Int64, frameSeq: Int, refIFrameSeq: Int?, sps: Data?, pps: Data?)] = [:]
    /// 已收到的I帧序号集合
    private var receivedIFrames = Set<Int>()
    /// 重排序缓冲区锁
    private let reorderLock = NSLock()
    /// 重排序缓冲区最大容量
    private let maxReorderBufferSize = 100 // 与输入缓冲区大小一致

    /// 解码回调，输出CVPixelBuffer和时间戳
    var onFrameDecoded: ((CVPixelBuffer, Int64) -> Void)? = { _, _ in }

    /// 初始化解码器，启动渲染线程
    init(width: Int, height: Int, codecType: String) {
        self.width = width
        self.height = height
        self.codecType = CodecType(rawValue: codecType.lowercased()) ?? .h264
        self.renderIntervalMs = Int64(1000.0 / Double(renderFps))
        startRenderTimer()
        print("[VideoDecoder] 初始化解码器: width=\(width), height=\(height)")
    }

    // ====== 输入缓冲区操作 ======
    /// 入队待解码帧
    private func enqueueInput(_ item: (Data, Int, Int64, Int, Int?, Data?, Data?)) {
        inputQueue.async(flags: .barrier) {
            if self.inputBuffer.count >= self.inputBufferMaxCount {
                self.inputBuffer.removeFirst() // 缓冲区满时丢弃最旧帧
                print("[VideoDecoder][警告] 输入缓冲区满，丢弃最旧帧")
            }
            self.inputBuffer.append(item)
        }
    }
    /// 出队待解码帧
    private func dequeueInput() -> (Data, Int, Int64, Int, Int?, Data?, Data?)? {
        var item: (Data, Int, Int64, Int, Int?, Data?, Data?)?
        inputQueue.sync {
            if !self.inputBuffer.isEmpty {
                item = self.inputBuffer.removeFirst()
            }
        }
        return item
    }
    // ====== 输出缓冲区操作 ======
    /// 入队解码后帧
    private func enqueueOutput(_ item: (CVPixelBuffer, Int64)) {
        outputQueue.async(flags: .barrier) {
            if self.outputBuffer.count >= self.outputBufferMaxCount {
                self.outputBuffer.removeFirst() // 缓冲区满时丢弃最旧帧
                print("[VideoDecoder][警告] 输出缓冲区满，丢弃最旧帧")
            }
            self.outputBuffer.append(item)
        }
    }
    /// 出队解码后帧
    private func dequeueOutput() -> (CVPixelBuffer, Int64)? {
        var item: (CVPixelBuffer, Int64)?
        outputQueue.sync {
            if !self.outputBuffer.isEmpty {
                item = self.outputBuffer.removeFirst()
            }
        }
        return item
    }

    // ====== 帧重排序与依赖管理 ======
    /// 处理帧重排序和依赖关系
    private func handleFrameReordering(frameData: Data, frameType: Int, timestamp: Int64, frameSeq: Int, refIFrameSeq: Int?, sps: Data?, pps: Data?) -> Bool {
        reorderLock.lock()
        defer { reorderLock.unlock() }
        
        // 1. 延迟丢帧检查
        let now = Int64(Date().timeIntervalSince1970 * 1000)
        let base = timestampBaseMs ?? 0
        let firstRel = firstFrameRelativeTimestamp ?? 0
        let absTimestamp = base + (timestamp - firstRel)
        
        if absTimestamp < now - maxAllowedDelayMs {
            print("[VideoDecoder][警告] 丢弃延迟帧: type=\(frameType), seq=\(frameSeq), delay=\(now - absTimestamp)ms")
            return false
        }
        
        // 2. 帧重排序处理
        if frameType == 0 { // I帧
            receivedIFrames.insert(frameSeq)
            lastIFrameSeq = frameSeq
            
            // I帧直接解码
            enqueueInput((frameData, frameType, timestamp, frameSeq, refIFrameSeq, sps, pps))
            
            // 检查并解码所有依赖该I帧的P帧
            let readyPFrames = reorderBuffer.values
                .filter { $0.refIFrameSeq == frameSeq }
                .sorted { $0.frameSeq < $1.frameSeq }
            
            for pFrame in readyPFrames {
                enqueueInput(pFrame)
                reorderBuffer.removeValue(forKey: pFrame.frameSeq)
            }
            
            if !readyPFrames.isEmpty {
                print("[VideoDecoder] I帧\(frameSeq)释放\(readyPFrames.count)个P帧")
            }
            
            // 清理过期缓存
            if reorderBuffer.count > maxReorderBufferSize {
                let toRemove = reorderBuffer.keys.sorted().prefix(reorderBuffer.count - maxReorderBufferSize)
                for seq in toRemove {
                    reorderBuffer.removeValue(forKey: seq)
                }
                print("[VideoDecoder][警告] 重排序缓冲区溢出，清理\(toRemove.count)个P帧")
            }
            
            return true
            
        } else { // P帧
            // 检查P帧依赖
            if let refIFrameSeq = refIFrameSeq, receivedIFrames.contains(refIFrameSeq) {
                // 依赖的I帧已收到，直接解码
                enqueueInput((frameData, frameType, timestamp, frameSeq, refIFrameSeq, sps, pps))
                return true
            } else {
                // 依赖的I帧未到，缓存P帧
                reorderBuffer[frameSeq] = (frameData, frameType, timestamp, frameSeq, refIFrameSeq, sps, pps)
                print("[VideoDecoder] P帧\(frameSeq)缓存，等待I帧\(refIFrameSeq ?? -1)")
                
                // 控制缓冲区大小
                if reorderBuffer.count > maxReorderBufferSize {
                    let toRemove = reorderBuffer.keys.sorted().prefix(reorderBuffer.count - maxReorderBufferSize)
                    for seq in toRemove {
                        reorderBuffer.removeValue(forKey: seq)
                    }
                    print("[VideoDecoder][警告] 重排序缓冲区溢出，清理\(toRemove.count)个P帧")
                }
                return false
            }
        }
    }

    // ====== 渲染定时器相关 ======
    /// 启动渲染定时器，定时从输出缓冲区取帧并刷新Flutter纹理
    private func startRenderTimer() {
        renderThreadRunning = true
        let timer = DispatchSource.makeTimerSource(queue: DispatchQueue.global())
        timer.schedule(deadline: .now(), repeating: .milliseconds(Int(renderIntervalMs / 2)))
        timer.setEventHandler { [weak self] in
            guard let self = self else { return }
            
            let now = Int64(Date().timeIntervalSince1970 * 1000)
            // 控制渲染间隔，避免过快渲染
            let timeSinceLastRender = now - self.lastRenderTimeMs
            if timeSinceLastRender < self.renderIntervalMs - self.renderJitterMs {
                return
            }
            
            // 低水位启动渲染逻辑
            if !self.renderStarted {
                var outputCount = 0
                self.outputQueue.sync { outputCount = self.outputBuffer.count }
                if outputCount >= Int(Double(self.outputBufferMaxCount) * 0.15) {
                    self.renderStarted = true
                    print("[VideoDecoder] 渲染启动，outputBuffer已达低水位: \(outputCount)")
                } else {
                    // 未达到低水位前不渲染
                    return
                }
            }
            
            if let (pixelBuffer, timestamp) = self.dequeueOutput() {
                // 延迟丢帧检查
                let now = Int64(Date().timeIntervalSince1970 * 1000)
                let base = timestampBaseMs ?? 0
                let firstRel = firstFrameRelativeTimestamp ?? 0
                let absTimestamp = base + (timestamp - firstRel)
                
                if absTimestamp < now - self.maxAllowedDelayMs {
                    print("[VideoDecoder][警告] 丢弃延迟渲染帧: delay=\(now - absTimestamp)ms")
                    return
                }
                
                DispatchQueue.main.async {
                    self.onFrameDecoded?(pixelBuffer, timestamp)
                }
                self.lastRenderTimeMs = now
                if !self.hasNotifiedFlutter {
                    self.hasNotifiedFlutter = true
                }
            }
        }
        timer.resume()
        renderTimer = timer
    }
    /// 停止渲染定时器
    private func stopRenderTimer() {
        renderThreadRunning = false
        renderTimer?.cancel()
        renderTimer = nil
    }

    /// 初始化解码会话（首次收到I帧时调用）
    private func setupSession(sps: Data?, pps: Data?) -> Bool {
        // 释放旧会话
        if let session = decompressionSession {
            VTDecompressionSessionInvalidate(session)
            decompressionSession = nil
        }
        formatDesc = nil
        isSessionReady = false
        guard let sps = sps, let pps = pps else {
            print("[VideoDecoder][错误] 缺少SPS/PPS，无法初始化解码会话")
            return false
        }
        
        // 校验SPS/PPS长度和类型
        let spsType: UInt8 = sps.count > 0 ? (sps[0] & 0x1F) : 0
        let ppsType: UInt8 = pps.count > 0 ? (pps[0] & 0x1F) : 0
        if sps.count < 3 || spsType != 7 {
            print("[VideoDecoder][错误] SPS内容异常，len=\(sps.count), type=\(spsType)")
            return false
        }
        if pps.count < 3 || ppsType != 8 {
            print("[VideoDecoder][错误] PPS内容异常，len=\(pps.count), type=\(ppsType)")
            return false
        }
        
        var success = false
        sps.withUnsafeBytes { spsPtr in
            pps.withUnsafeBytes { ppsPtr in
                let parameterSetPointers: [UnsafePointer<UInt8>] = [
                    spsPtr.baseAddress!.assumingMemoryBound(to: UInt8.self),
                    ppsPtr.baseAddress!.assumingMemoryBound(to: UInt8.self)
                ]
                let parameterSetSizes: [Int] = [sps.count, pps.count]
                let status = CMVideoFormatDescriptionCreateFromH264ParameterSets(
                    allocator: kCFAllocatorDefault,
                    parameterSetCount: 2,
                    parameterSetPointers: parameterSetPointers,
                    parameterSetSizes: parameterSetSizes,
                    nalUnitHeaderLength: 4,
                    formatDescriptionOut: &formatDesc
                )
                if status != noErr {
                    print("[VideoDecoder][错误] 创建FormatDescription失败: \(status)")
                    success = false
                } else {
                    success = true
                }
            }
        }
        if !success { return false }
        var callback = VTDecompressionOutputCallbackRecord(
            decompressionOutputCallback: { (decompressionOutputRefCon, _, status, _, imageBuffer, pts, _) in
                let decoder = Unmanaged<VideoDecoder>.fromOpaque(decompressionOutputRefCon!).takeUnretainedValue()
                if status == noErr, let pixelBuffer = imageBuffer {
                    // 入队到输出缓冲区，由渲染线程拉取
                    decoder.enqueueOutput((pixelBuffer, Int64(pts.seconds * 1000)))
                } else {
                    print("[VideoDecoder][错误] 解码回调失败: \(status)")
                }
            },
            decompressionOutputRefCon: UnsafeMutableRawPointer(Unmanaged.passUnretained(self).toOpaque())
        )
        let attrs: [NSString: Any] = [
            kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_420YpCbCr8BiPlanarFullRange,
            kCVPixelBufferWidthKey: width,
            kCVPixelBufferHeightKey: height,
            kCVPixelBufferOpenGLESCompatibilityKey: true,
            kCVPixelBufferIOSurfacePropertiesKey: [:]
        ]
        
        let status2 = VTDecompressionSessionCreate(
            allocator: kCFAllocatorDefault,
            formatDescription: formatDesc!,
            decoderSpecification: nil,
            imageBufferAttributes: attrs as CFDictionary,
            outputCallback: &callback,
            decompressionSessionOut: &decompressionSession
        )
        if status2 != noErr {
            print("[VideoDecoder][错误] 创建解码会话失败: \(status2)")
            return false
        }
        isSessionReady = true
        print("[VideoDecoder] 解码会话初始化成功")
        return true
    }

    /// 解码一帧数据
    func decodeFrame(frameData: Data, frameType: Int, timestamp: Int64, frameSeq: Int, refIFrameSeq: Int?, sps: Data? = nil, pps: Data? = nil) {
        // 1. 初始化时间戳基准
        if timestampBaseMs == nil {
            timestampBaseMs = Int64(Date().timeIntervalSince1970 * 1000)
            firstFrameRelativeTimestamp = timestamp
            print("[VideoDecoder] 设置时间戳基准: base=\(timestampBaseMs!), firstRel=\(firstFrameRelativeTimestamp!)")
        }
        
        // 2. 帧重排序和依赖管理
        if !handleFrameReordering(frameData: frameData, frameType: frameType, timestamp: timestamp, frameSeq: frameSeq, refIFrameSeq: refIFrameSeq, sps: sps, pps: pps) {
            return
        }
        
        // 3. 解码处理（由inputQueue触发）
        decodeQueue.async { [weak self] in
            guard let self = self else { return }
            guard let (frameData, frameType, timestamp, frameSeq, refIFrameSeq, sps, pps) = self.dequeueInput() else { return }
            
            if !self.isSessionReady, let sps = sps, let pps = pps {
                guard self.setupSession(sps: sps, pps: pps) else { return }
            }
            guard let session = self.decompressionSession else { return }
            guard frameData.count > 3 else { return }

            // 检查并移除AnnexB起始码
            var startCodeSize = 0
            var naluData: Data
            
            if frameData.count >= 4 && frameData[0] == 0x00 && frameData[1] == 0x00 && frameData[2] == 0x00 && frameData[3] == 0x01 {
                startCodeSize = 4
                naluData = frameData.subdata(in: 4..<frameData.count)
            } else if frameData.count >= 3 && frameData[0] == 0x00 && frameData[1] == 0x00 && frameData[2] == 0x01 {
                startCodeSize = 3
                naluData = frameData.subdata(in: 3..<frameData.count)
            } else {
                print("[VideoDecoder][警告] 未找到起始码")
                naluData = frameData
            }
            
            // 创建AVCC格式数据
            let naluLength = UInt32(naluData.count).bigEndian
            var avccData = Data(capacity: naluData.count + 4)
            withUnsafeBytes(of: naluLength) { ptr in
                avccData.append(ptr.baseAddress!.assumingMemoryBound(to: UInt8.self), count: 4)
            }
            avccData.append(naluData)

            var blockBuffer: CMBlockBuffer?
            let status = CMBlockBufferCreateWithMemoryBlock(
                allocator: kCFAllocatorDefault,
                memoryBlock: nil,
                blockLength: avccData.count,
                blockAllocator: nil,
                customBlockSource: nil,
                offsetToData: 0,
                dataLength: avccData.count,
                flags: kCMBlockBufferAssureMemoryNowFlag,
                blockBufferOut: &blockBuffer
            )
            if status != kCMBlockBufferNoErr { 
                print("[VideoDecoder][错误] 创建BlockBuffer失败: \(status)")
                return 
            }
            
            // 复制数据到BlockBuffer
            if let blockBuffer = blockBuffer {
                let status2 = avccData.withUnsafeBytes { ptr in
                    CMBlockBufferReplaceDataBytes(
                        with: ptr.baseAddress!,
                        blockBuffer: blockBuffer,
                        offsetIntoDestination: 0,
                        dataLength: avccData.count
                    )
                }
                if status2 != kCMBlockBufferNoErr {
                    print("[VideoDecoder][错误] 复制数据到BlockBuffer失败: \(status2)")
                    return
                }
            }

            var sampleBuffer: CMSampleBuffer?
            var timing = CMSampleTimingInfo(
                duration: .invalid,
                presentationTimeStamp: CMTime(value: timestamp, timescale: 1000),
                decodeTimeStamp: .invalid
            )
            let status2 = CMSampleBufferCreate(
                allocator: kCFAllocatorDefault,
                dataBuffer: blockBuffer,
                dataReady: true,
                makeDataReadyCallback: nil,
                refcon: nil,
                formatDescription: self.formatDesc,
                sampleCount: 1,
                sampleTimingEntryCount: 1,
                sampleTimingArray: &timing,
                sampleSizeEntryCount: 1,
                sampleSizeArray: [avccData.count],
                sampleBufferOut: &sampleBuffer
            )
            if status2 != noErr { 
                print("[VideoDecoder][错误] 创建SampleBuffer失败: \(status2)")
                return 
            }

            let decodeFlags: VTDecodeFrameFlags = [._EnableAsynchronousDecompression]
            var infoFlags = VTDecodeInfoFlags()
            let status3 = VTDecompressionSessionDecodeFrame(
                session,
                sampleBuffer: sampleBuffer!,
                flags: decodeFlags,
                frameRefcon: nil,
                infoFlagsOut: &infoFlags
            )
            if status3 != noErr {
                print("[VideoDecoder][错误] 解码失败: \(status3)")
                if status3 == -6661 {
                    print("  - 错误类型: kVTInvalidSessionErr (解码会话无效)")
                    print("  - 会话状态: \(self.isSessionReady ? "就绪" : "未就绪")")
                    print("  - formatDesc: \(self.formatDesc != nil ? "有效" : "无效")")
                }
            }
        }
    }

    /// 释放解码器资源
    func release() {
        stopRenderTimer()
        decodeQueue.sync {
            if let session = decompressionSession {
                VTDecompressionSessionInvalidate(session)
            }
            decompressionSession = nil
            formatDesc = nil
            isSessionReady = false
            frameSeqSet.removeAll()
            lastIFrameSeq = nil
            
            // 清理重排序相关资源
            reorderLock.lock()
            reorderBuffer.removeAll()
            receivedIFrames.removeAll()
            reorderLock.unlock()
        }
        inputQueue.async(flags: .barrier) { self.inputBuffer.removeAll() }
        outputQueue.async(flags: .barrier) { self.outputBuffer.removeAll() }
        print("[VideoDecoder] 解码器已释放")
    }
}