import Foundation

public enum ProtocolManager {

    // MARK: - Checksum

    /// 计算校验和：所有字节求和后取二进制补码的低 8 位
    /// checksum = (~sum + 1) & 0xFF
    public static func calculateChecksum(_ data: Data) -> UInt8 {
        var sum: UInt32 = 0
        for byte in data {
            sum += UInt32(byte)
        }
        return UInt8((~sum &+ 1) & 0xFF)
    }

    /// 校验给定数据的校验和是否匹配。
    /// - Parameters:
    ///   - data: 用于计算校验的字节序列（不含最后的校验位）。
    ///   - expected: 期望的校验和。
    /// - Returns: 是否匹配。
    public static func verifyChecksum(_ data: Data, expected: UInt8) -> Bool {
        calculateChecksum(data) == expected
    }

    // MARK: - Frame Creation

    /// 根据数据大小自动选择单帧或分片
    /// - Throws: 当 payload 超过 UInt16 分页能力时（> 65535 个分片）抛出 invalidFrame
    public static func createFrames(
        type: UInt8,
        data: Data,
        head: FrameHead = .appToDevice,
        maxDataSize: Int = FrameConstants.maxDataSize
    ) -> [Data] {
        // subpageTotal/curPage are encoded as UInt16, so the absolute ceiling is 65535 pages.
        // Returning [] keeps the call non-throwing for backwards compatibility; callers (the
        // protocol service / file transfer) detect the empty array as a payload-too-large signal.
        let maxChunk = max(1, maxDataSize)
        let totalPages = (data.count + maxChunk - 1) / maxChunk
        if totalPages > Int(UInt16.max) {
            return []
        }

        if data.count <= maxDataSize {
            return [createSingleFrame(type: type, data: data, head: head)]
        } else {
            return createFragmentedFrames(type: type, data: data, head: head, maxDataSize: maxChunk)
        }
    }

    private static func createSingleFrame(type: UInt8, data: Data, head: FrameHead) -> Data {
        let totalSize = FrameConstants.headerSize + data.count + FrameConstants.footerSize
        var buffer = Data(capacity: totalSize)

        // Header: head(1) + type(1) + subpageTotal(2) + curPage(2) + dataLen(2)
        buffer.append(head.rawValue)
        buffer.append(type)
        buffer.append(0) // subpageTotal high
        buffer.append(0) // subpageTotal low
        buffer.append(0) // curPage high
        buffer.append(0) // curPage low
        buffer.append(UInt8((data.count >> 8) & 0xFF)) // dataLen high
        buffer.append(UInt8(data.count & 0xFF))         // dataLen low

        // Data
        buffer.append(data)

        // Checksum (calculated on everything before checksum)
        let checksum = calculateChecksum(buffer)
        buffer.append(checksum)

        return buffer
    }

    private static func createFragmentedFrames(
        type: UInt8,
        data: Data,
        head: FrameHead,
        maxDataSize: Int
    ) -> [Data] {
        let totalPages = (data.count + maxDataSize - 1) / maxDataSize
        var frames: [Data] = []
        frames.reserveCapacity(totalPages)

        for i in 0..<totalPages {
            let start = i * maxDataSize
            let end = min(start + maxDataSize, data.count)
            let chunk = data[start..<end]

            var buffer = Data(capacity: FrameConstants.headerSize + chunk.count + FrameConstants.footerSize)

            buffer.append(head.rawValue)
            buffer.append(type)

            // subpageTotal (big-endian UInt16)
            buffer.append(UInt8((totalPages >> 8) & 0xFF))
            buffer.append(UInt8(totalPages & 0xFF))

            // curPage: 页号从 (total-1) 倒数到 0
            let curPageVal = totalPages - 1 - i
            buffer.append(UInt8((curPageVal >> 8) & 0xFF))
            buffer.append(UInt8(curPageVal & 0xFF))

            // dataLen
            buffer.append(UInt8((chunk.count >> 8) & 0xFF))
            buffer.append(UInt8(chunk.count & 0xFF))

            // data
            buffer.append(contentsOf: chunk)

            // checksum
            let checksum = calculateChecksum(buffer)
            buffer.append(checksum)

            frames.append(buffer)
        }

        return frames
    }

    // MARK: - Frame Parsing

    /// 解析原始二进制为协议帧。
    /// - Parameter data: 含有完整帧内容的字节序列。
    /// - Returns: 若校验与长度正确，返回 `ProtocolFrame`，否则返回 `nil`。
    public static func parseFrame(_ data: Data) -> ProtocolFrame? {
        let minSize = FrameConstants.headerSize + FrameConstants.footerSize
        guard data.count >= minSize else { return nil }

        let head = data[data.startIndex]
        guard head == FrameHead.deviceToApp.rawValue || head == FrameHead.appToDevice.rawValue else {
            return nil
        }

        let type = data[data.startIndex + 1]
        let subpageTotal = UInt16(data[data.startIndex + 2]) << 8 | UInt16(data[data.startIndex + 3])
        let curPage = UInt16(data[data.startIndex + 4]) << 8 | UInt16(data[data.startIndex + 5])
        let dataLen = UInt16(data[data.startIndex + 6]) << 8 | UInt16(data[data.startIndex + 7])

        let expectedSize = FrameConstants.headerSize + Int(dataLen) + FrameConstants.footerSize
        guard data.count >= expectedSize else { return nil }

        let dataStart = data.startIndex + FrameConstants.headerSize
        let dataEnd = dataStart + Int(dataLen)
        let frameData = data[dataStart..<dataEnd]
        let checksum = data[dataEnd]

        // Verify checksum on everything before checksum byte
        let dataToCheck = data[data.startIndex..<dataEnd]
        guard verifyChecksum(dataToCheck, expected: checksum) else {
            return nil
        }

        return ProtocolFrame(
            head: head,
            type: type,
            subpageTotal: subpageTotal,
            curPage: curPage,
            dataLen: dataLen,
            data: Data(frameData),
            checksum: checksum
        )
    }
}