import Foundation
import AVFoundation
import CoreMedia

/// “自编码”串流模式解码器：
/// 从 video_channel DataChannel 接收被控端 MediaCodec 编码的 H.264 裸流，
/// 重组分片 -> 解析单元 -> 转换为 AVCC 格式 -> 通过 AVSampleBufferDisplayLayer 硬解渲染。
///
/// 二进制协议（与被控端 SelfCodecEncoder 保持一致，均为大端序）：
/// 分片: MAGIC(0xAB,1B) + seq(int32) + total(int16) + idx(int16) + len(int32) + payload
/// 单元: type(1B: 1=CONFIG 2=FRAME) + pts(uint32,ms) + isKey(1B) + len(int32) + H.264 Annex-B 数据
final class SelfCodecDecoder {

    // ---- 必须与被控端 SelfCodecEncoder 保持一致 ----
    private static let magic: UInt8 = 0xAB
    private static let unitTypeConfig: UInt8 = 1
    private static let unitTypeFrame: UInt8 = 2

    /// 渲染图层（由 SelfCodecDisplayView 提供）
    weak var displayLayer: AVSampleBufferDisplayLayer?
    /// 分辨率变化回调（主线程）
    var onResolutionUpdate: ((Int, Int) -> Void)?
    /// 自编码模式录制回调：解码出的 AVCC CMSampleBuffer 会回传（由 VideoRecorder 写盘）
    var onDecodedSampleBuffer: ((CMSampleBuffer, Bool) -> Void)?

    private(set) var videoWidth = 0
    private(set) var videoHeight = 0

    private let queue = DispatchQueue(label: "com.ttstd.selfcodec.decoder")
    private var enabled = false

    // 分片重组
    private final class ChunkBuffer {
        var total = 0
        var received = 0
        var chunks: [Data?] = []
    }
    private var assembling: [Int32: ChunkBuffer] = [:]
    private var lastSeq: Int32 = -1

    // H.264 参数集与格式
    private var sps: Data?
    private var pps: Data?
    private var formatDescription: CMVideoFormatDescription?

    // MARK: - 开关

    func setEnabled(_ on: Bool) {
        queue.async {
            self.enabled = on
            if !on {
                self.resetLocked()
            }
        }
    }

    func release() {
        queue.async {
            self.enabled = false
            self.resetLocked()
        }
    }

    private func resetLocked() {
        assembling.removeAll()
        lastSeq = -1
        sps = nil
        pps = nil
        formatDescription = nil
        videoWidth = 0
        videoHeight = 0
        displayLayer?.flushAndRemoveImage()
    }

    // MARK: - 入口：DataChannel 二进制消息

    func onBinaryMessage(_ raw: Data) {
        queue.async {
            self.handleChunk(raw)
        }
    }

    private func handleChunk(_ raw: Data) {
        guard enabled, raw.count >= 13 else { return }
        var r = BigEndianReader(data: raw)
        guard r.readByte() == Self.magic else { return }
        guard let seq = r.readInt32() else { return }

        // 过期检查：同一单元的所有分片 seq 相同，只丢弃明显更旧的单元
        if lastSeq != -1 && seq < lastSeq && (lastSeq - seq) < 1000 {
            return
        }
        guard let total = r.readInt16(), let idx = r.readInt16(),
              let len = r.readInt32(), let chunk = r.readData(Int(len)),
              total > 0 else { return }

        let cb: ChunkBuffer
        if let existing = assembling[seq] {
            cb = existing
        } else {
            cb = ChunkBuffer()
            cb.total = Int(total)
            cb.chunks = Array(repeating: nil, count: Int(total))
            assembling[seq] = cb
            if assembling.count > 24, let minKey = assembling.keys.min() {
                assembling.removeValue(forKey: minKey)
            }
        }

        let i = Int(idx)
        if i >= 0 && i < cb.total && cb.chunks[i] == nil {
            cb.chunks[i] = chunk
            cb.received += 1
        }
        if cb.received == cb.total {
            assembling.removeValue(forKey: seq)
            if seq > lastSeq { lastSeq = seq }
            var unit = Data()
            for c in cb.chunks where c != nil { unit.append(c!) }
            handleUnit(unit)
        }
    }

    // MARK: - 单元处理

    private func handleUnit(_ unit: Data) {
        var r = BigEndianReader(data: unit)
        guard let type = r.readByte(),
              let ptsRaw = r.readInt32(),
              let isKey = r.readByte(),
              let len = r.readInt32(),
              let data = r.readData(Int(len)) else { return }
        let ptsMs = Int64(UInt32(bitPattern: ptsRaw))

        if type == Self.unitTypeConfig {
            handleConfig(data)
        } else if type == Self.unitTypeFrame {
            handleFrame(data, ptsMs: ptsMs, isKey: isKey != 0)
        }
    }

    /// CONFIG 单元：MediaCodec 的 csd-0(SPS)/csd-1(PPS)，Annex-B 格式
    private func handleConfig(_ data: Data) {
        for nal in Self.annexBNalUnits(in: data) {
            guard let first = nal.first else { continue }
            switch first & 0x1F {
            case 7: if sps != nal { sps = nal; formatDescription = nil }
            case 8: if pps != nal { pps = nal; formatDescription = nil }
            default: break
            }
        }
        rebuildFormatDescriptionIfNeeded()
    }

    private func rebuildFormatDescriptionIfNeeded() {
        guard formatDescription == nil, let sps, let pps else { return }
        var desc: CMVideoFormatDescription?
        let status = sps.withUnsafeBytes { spsPtr -> OSStatus in
            pps.withUnsafeBytes { ppsPtr -> OSStatus in
                let paramSets: [UnsafePointer<UInt8>] = [
                    spsPtr.bindMemory(to: UInt8.self).baseAddress!,
                    ppsPtr.bindMemory(to: UInt8.self).baseAddress!
                ]
                let sizes: [Int] = [sps.count, pps.count]
                return CMVideoFormatDescriptionCreateFromH264ParameterSets(
                    allocator: kCFAllocatorDefault,
                    parameterSetCount: 2,
                    parameterSetPointers: paramSets,
                    parameterSetSizes: sizes,
                    nalUnitHeaderLength: 4,
                    formatDescriptionOut: &desc)
            }
        }
        guard status == noErr, let desc else {
            NSLog("[SelfCodec] create format description failed: \(status)")
            return
        }
        formatDescription = desc
        displayLayer?.flush()

        let dims = CMVideoFormatDescriptionGetDimensions(desc)
        let w = Int(dims.width), h = Int(dims.height)
        if w != videoWidth || h != videoHeight {
            videoWidth = w
            videoHeight = h
            DispatchQueue.main.async {
                self.onResolutionUpdate?(w, h)
            }
        }
    }

    /// FRAME 单元：Annex-B -> AVCC（4 字节大端长度前缀），封装 CMSampleBuffer 后送显
    private func handleFrame(_ data: Data, ptsMs: Int64, isKey: Bool) {
        // 帧内可能携带 SPS/PPS（部分编码器关键帧前重发参数集）
        var vclData = Data()
        for nal in Self.annexBNalUnits(in: data) {
            guard let first = nal.first else { continue }
            let nalType = first & 0x1F
            if nalType == 7 {
                if sps != nal { sps = nal; formatDescription = nil }
                continue
            }
            if nalType == 8 {
                if pps != nal { pps = nal; formatDescription = nil }
                continue
            }
            var lenBE = UInt32(nal.count).bigEndian
            withUnsafeBytes(of: &lenBE) { vclData.append(contentsOf: $0) }
            vclData.append(nal)
        }
        rebuildFormatDescriptionIfNeeded()
        guard enabled, !vclData.isEmpty, let formatDescription, let layer = displayLayer else { return }

        // 构造 CMBlockBuffer
        var blockBuffer: CMBlockBuffer?
        var status = CMBlockBufferCreateWithMemoryBlock(
            allocator: kCFAllocatorDefault,
            memoryBlock: nil,
            blockLength: vclData.count,
            blockAllocator: kCFAllocatorDefault,
            customBlockSource: nil,
            offsetToData: 0,
            dataLength: vclData.count,
            flags: 0,
            blockBufferOut: &blockBuffer)
        guard status == kCMBlockBufferNoErr, let blockBuffer else { return }
        status = vclData.withUnsafeBytes { ptr in
            CMBlockBufferReplaceDataBytes(
                with: ptr.baseAddress!,
                blockBuffer: blockBuffer,
                offsetIntoDestination: 0,
                dataLength: vclData.count)
        }
        guard status == kCMBlockBufferNoErr else { return }

        // 构造 CMSampleBuffer（立即显示，不依赖时间轴）
        var sampleBuffer: CMSampleBuffer?
        var timing = CMSampleTimingInfo(
            duration: .invalid,
            presentationTimeStamp: CMTime(value: ptsMs, timescale: 1000),
            decodeTimeStamp: .invalid)
        var sampleSize = vclData.count
        status = CMSampleBufferCreateReady(
            allocator: kCFAllocatorDefault,
            dataBuffer: blockBuffer,
            formatDescription: formatDescription,
            sampleCount: 1,
            sampleTimingEntryCount: 1,
            sampleTimingArray: &timing,
            sampleSizeEntryCount: 1,
            sampleSizeArray: &sampleSize,
            sampleBufferOut: &sampleBuffer)
        guard status == noErr, let sampleBuffer else { return }

        if let attachments = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, createIfNecessary: true) as? [CFMutableDictionary],
           let dict = attachments.first {
            CFDictionarySetValue(
                dict,
                Unmanaged.passUnretained(kCMSampleAttachmentKey_DisplayImmediately).toOpaque(),
                Unmanaged.passUnretained(kCFBooleanTrue).toOpaque())
            if !isKey {
                CFDictionarySetValue(
                    dict,
                    Unmanaged.passUnretained(kCMSampleAttachmentKey_NotSync).toOpaque(),
                    Unmanaged.passUnretained(kCFBooleanTrue).toOpaque())
            }
        }

        if layer.status == .failed || layer.requiresFlushToResumeDecoding {
            layer.flush()
        }
        layer.enqueue(sampleBuffer)

        // 自编码模式录制：把解码后的 AVCC 帧回传给录制器写盘
        onDecodedSampleBuffer?(sampleBuffer, isKey)
    }

    // MARK: - Annex-B 拆分

    /// 拆分 Annex-B 码流（00 00 01 / 00 00 00 01 起始码）为 NAL 单元数组（不含起始码）
    static func annexBNalUnits(in data: Data) -> [Data] {
        var result: [Data] = []
        let bytes = [UInt8](data)
        let count = bytes.count
        var starts: [Int] = []
        var i = 0
        while i + 2 < count {
            if bytes[i] == 0 && bytes[i + 1] == 0 && bytes[i + 2] == 1 {
                starts.append(i + 3)
                i += 3
            } else {
                i += 1
            }
        }
        guard !starts.isEmpty else { return data.isEmpty ? [] : [data] }
        for (idx, start) in starts.enumerated() {
            var end = count
            if idx + 1 < starts.count {
                end = starts[idx + 1] - 3
                // 兼容 4 字节起始码（00 00 00 01）
                if end > start && bytes[end - 1] == 0 { end -= 1 }
            }
            if end > start {
                result.append(Data(bytes[start..<end]))
            }
        }
        return result
    }
}

// MARK: - 大端序读取器

private struct BigEndianReader {
    let data: Data
    private var offset: Int

    init(data: Data) {
        self.data = data
        self.offset = data.startIndex
    }

    mutating func readByte() -> UInt8? {
        guard offset < data.endIndex else { return nil }
        defer { offset += 1 }
        return data[offset]
    }

    mutating func readInt16() -> Int16? {
        guard offset + 2 <= data.endIndex else { return nil }
        let v = (UInt16(data[offset]) << 8) | UInt16(data[offset + 1])
        offset += 2
        return Int16(bitPattern: v)
    }

    mutating func readInt32() -> Int32? {
        guard offset + 4 <= data.endIndex else { return nil }
        var v: UInt32 = 0
        for i in 0..<4 { v = (v << 8) | UInt32(data[offset + i]) }
        offset += 4
        return Int32(bitPattern: v)
    }

    mutating func readData(_ length: Int) -> Data? {
        guard length >= 0, offset + length <= data.endIndex else { return nil }
        defer { offset += length }
        return data.subdata(in: offset..<(offset + length))
    }
}