video_decode_plugin/ios/Classes/VideoDecodePlugin.swift

import Flutter
import UIKit
import AVFoundation

public class VideoDecodePlugin: NSObject, FlutterPlugin, FlutterTexture {
    private var channel: FlutterMethodChannel?
    private var registrar: FlutterPluginRegistrar?
    private var decoder: VideoDecoder?
    private var textureId: Int64?
    private var textureRegistry: FlutterTextureRegistry?
    private var latestPixelBuffer: CVPixelBuffer?
    private let textureQueue = DispatchQueue(label: "video_decode_plugin.texture.queue")
    private var cachedSps: Data?
    private var cachedPps: Data?
    private var hasNotifiedFlutter = false

    public static func register(with registrar: FlutterPluginRegistrar) {
        let channel = FlutterMethodChannel(name: "video_decode_plugin", binaryMessenger: registrar.messenger())
        let instance = VideoDecodePlugin()
        instance.channel = channel
        instance.registrar = registrar
        instance.textureRegistry = registrar.textures()
        registrar.addMethodCallDelegate(instance, channel: channel)
    }

    public func handle(_ call: FlutterMethodCall, result: @escaping FlutterResult) {
        switch call.method {
        case "initDecoder":
            handleInitDecoder(call: call, result: result)
        case "decodeFrame":
            handleDecodeFrame(call: call, result: result)
        case "releaseDecoder":
            handleReleaseDecoder(call: call, result: result)
        default:
            result(FlutterMethodNotImplemented)
        }
    }

    /// 初始化解码器
    private func handleInitDecoder(call: FlutterMethodCall, result: @escaping FlutterResult) {
        guard let args = call.arguments as? [String: Any],
              let width = args["width"] as? Int,
              let height = args["height"] as? Int,
              let codecType = args["codecType"] as? String else {
            print("[VideoDecodePlugin][错误] 参数解析失败：\(String(describing: call.arguments))")
            result(FlutterError(code: "INVALID_ARGS", message: "参数错误", details: nil))
            return
        }
        
        // 释放旧解码器和纹理
        decoder?.release()
        decoder = nil
        if let tid = textureId {
            textureRegistry?.unregisterTexture(tid)
            textureId = nil
        }
        
        // 注册Flutter纹理
        guard let registry = textureRegistry else {
            print("[VideoDecodePlugin][错误] 无法获取纹理注册表")
            result(FlutterError(code: "NO_TEXTURE_REGISTRY", message: "无法获取纹理注册表", details: nil))
            return
        }
        
        let textureId = registry.register(self)
        self.textureId = textureId
        
        // 创建解码器
        let decoder = VideoDecoder(width: width, height: height, codecType: codecType)
        self.decoder = decoder
        decoder.onFrameDecoded = { [weak self] pixelBuffer, _ in
            guard let self = self else { return }
            self.textureQueue.async {
                self.latestPixelBuffer = pixelBuffer
                self.textureRegistry?.textureFrameAvailable(self.textureId ?? 0)
                if !self.hasNotifiedFlutter {
                    self.hasNotifiedFlutter = true
                    DispatchQueue.main.async {
                        self.channel?.invokeMethod("onFrameRendered", arguments: ["textureId": self.textureId ?? 0])
                    }
                }
            }
        }
        print("[VideoDecodePlugin] 解码器初始化成功，textureId=\(textureId)")
        result(textureId)
    }

    /// 去除NALU起始码的工具方法
    private func stripStartCode(_ data: Data) -> Data {
        var stripped: Data = data
        if data.count > 4 && data[0] == 0x00 && data[1] == 0x00 && data[2] == 0x00 && data[3] == 0x01 {
            stripped = data.subdata(in: 4..<data.count)
        } else if data.count > 3 && data[0] == 0x00 && data[1] == 0x00 && data[2] == 0x01 {
            stripped = data.subdata(in: 3..<data.count)
        }
        return stripped
    }

    // 修改：更合理的side_data检测工具
    private func checkNaluForSideData(_ nalu: Data, naluType: UInt8) -> Bool {
        let maxSize = naluType == 5 ? 150_000 : 30_000 // I帧最大150KB，P帧最大30KB
        if nalu.count > maxSize {
            print("[VideoDecodePlugin][警告] NALU长度异常，可能包含side_data，type=\(naluType)，len=\(nalu.count)")
            return true
        }
        return false
    }

    // 简化：提取第一个合法NALU工具
    private func extractFirstValidNalu(_ nalu: Data) -> Data {
        // 查找第一个起始码（支持3字节和4字节格式）
        var start = -1
        var startCodeLength = 0
        
        // 检查4字节起始码
        if nalu.count >= 4 && nalu[0] == 0x00 && nalu[1] == 0x00 && nalu[2] == 0x00 && nalu[3] == 0x01 {
            start = 0
            startCodeLength = 4
        }
        // 检查3字节起始码
        else if nalu.count >= 3 && nalu[0] == 0x00 && nalu[1] == 0x00 && nalu[2] == 0x01 {
            start = 0
            startCodeLength = 3
        }
        // 在数据中间查找起始码
        else {
            for i in 0..<(nalu.count - 4) {
                if nalu[i] == 0x00 && nalu[i + 1] == 0x00 {
                    if nalu[i + 2] == 0x00 && nalu[i + 3] == 0x01 {
                        start = i
                        startCodeLength = 4
                        break
                    } else if nalu[i + 2] == 0x01 {
                        start = i
                        startCodeLength = 3
                        break
                    }
                }
            }
        }
        
        if start == -1 {
            print("[VideoDecodePlugin][警告] NALU无AnnexB起始码，丢弃该帧")
            return Data()
        }
        
        let searchRange = (start + startCodeLength)..<nalu.count
        var end = nalu.count
        
        // 查找下一个起始码（同时支持3字节和4字节格式）
        for i in searchRange.lowerBound..<(nalu.count - 3) {
            if nalu[i] == 0x00 && nalu[i + 1] == 0x00 {
                if i + 3 < nalu.count && nalu[i + 2] == 0x00 && nalu[i + 3] == 0x01 {
                    end = i
                    break
                } else if nalu[i + 2] == 0x01 {
                    end = i
                    break
                }
            }
        }
        
        let extractedNalu = nalu[start..<end]
        
        // 验证提取的NALU
        if extractedNalu.count >= startCodeLength + 1 {
            let naluType = extractedNalu[startCodeLength] & 0x1F
            if checkNaluForSideData(extractedNalu, naluType: naluType) {
                return Data()
            }
        }
        
        return extractedNalu
    }

    /// 解码视频帧
    private func handleDecodeFrame(call: FlutterMethodCall, result: @escaping FlutterResult) {
        guard let args = call.arguments as? [String: Any],
              let frameData = args["frameData"] as? FlutterStandardTypedData,
              let frameType = args["frameType"] as? Int,
              let timestamp = args["timestamp"] as? Int,
              let frameSeq = args["frameSeq"] as? Int else {
            result(FlutterError(code: "INVALID_ARGS", message: "参数错误", details: nil))
            return
        }
        let refIFrameSeq = args["refIFrameSeq"] as? Int
        let data = frameData.data

        // 解析NALU类型
        let naluType: UInt8 = {
            if data.count > 4 && data[0] == 0x00 && data[1] == 0x00 && data[2] == 0x00 && data[3] == 0x01 {
                return data[4] & 0x1F
            } else if data.count > 3 && data[0] == 0x00 && data[1] == 0x00 && data[2] == 0x01 {
                return data[3] & 0x1F
            }
            return 0
        }()

        // 缓存SPS/PPS（去除起始码）
        if naluType == 7 { // SPS
            cachedSps = stripStartCode(data)
            result(true)
            return
        } else if naluType == 8 { // PPS
            cachedPps = stripStartCode(data)
            result(true)
            return
        } else if naluType == 5 { // IDR/I帧
            // 提取第一个合法NALU
            let firstNalu = extractFirstValidNalu(data)
            if firstNalu.isEmpty { 
                result(false)
                return 
            }
            
            decoder?.decodeFrame(frameData: firstNalu, frameType: frameType, timestamp: Int64(timestamp), frameSeq: frameSeq, refIFrameSeq: refIFrameSeq, sps: cachedSps, pps: cachedPps)
        } else {
            // 提取第一个合法NALU
            let firstNalu = extractFirstValidNalu(data)
            if firstNalu.isEmpty { 
                result(false)
                return 
            }
            
            decoder?.decodeFrame(frameData: firstNalu, frameType: frameType, timestamp: Int64(timestamp), frameSeq: frameSeq, refIFrameSeq: refIFrameSeq)
        }
        result(true)
    }

    /// 释放解码器资源
    private func handleReleaseDecoder(call: FlutterMethodCall, result: @escaping FlutterResult) {
        decoder?.release()
        decoder = nil
        if let tid = textureId {
            textureRegistry?.unregisterTexture(tid)
            textureId = nil
        }
        latestPixelBuffer = nil
        hasNotifiedFlutter = false
        print("[VideoDecodePlugin] 解码器和纹理已释放")
        result(true)
    }

    // MARK: - FlutterTexture协议实现
    public func copyPixelBuffer() -> Unmanaged<CVPixelBuffer>? {
        var pixelBuffer: CVPixelBuffer?
        textureQueue.sync {
            pixelBuffer = self.latestPixelBuffer
        }
        if let pb = pixelBuffer {
            return Unmanaged.passRetained(pb)
        }
        return nil
    }
}