import CommonCrypto
import Foundation

enum AesCtr {
    /// AES-128-CTR (big-endian counter), zoals CENC `cenc` scheme.
    static func crypt(key: Data, iv: Data, data: Data) -> Data? {
        guard key.count == 16, data.count > 0 else {
            return data.isEmpty ? data : nil
        }
        var counter = Data(count: 16)
        let copyLen = min(iv.count, 16)
        counter.replaceSubrange(0..<copyLen, with: iv.prefix(copyLen))

        var cryptor: CCCryptorRef?
        let status = counter.withUnsafeBytes { ivBytes in
            key.withUnsafeBytes { keyBytes in
                CCCryptorCreateWithMode(
                    CCOperation(kCCEncrypt),
                    CCMode(kCCModeCTR),
                    CCAlgorithm(kCCAlgorithmAES),
                    CCPadding(ccNoPadding),
                    ivBytes.baseAddress,
                    keyBytes.baseAddress,
                    key.count,
                    nil,
                    0,
                    0,
                    CCModeOptions(kCCModeOptionCTR_BE),
                    &cryptor
                )
            }
        }
        guard status == kCCSuccess, let cryptor else { return nil }
        defer { CCCryptorRelease(cryptor) }

        var out = Data(count: data.count)
        var moved = 0
        let cryptStatus = out.withUnsafeMutableBytes { outBytes in
            data.withUnsafeBytes { inBytes in
                CCCryptorUpdate(
                    cryptor,
                    inBytes.baseAddress,
                    data.count,
                    outBytes.baseAddress,
                    data.count,
                    &moved
                )
            }
        }
        guard cryptStatus == kCCSuccess, moved == data.count else { return nil }
        return out
    }

    /// CBCS pattern decrypt — gelijk aan Chromium `DecryptWithPattern` (ISO 23001-7 §10.4):
    /// - IV wordt één keer gezet per protected range (subsample), niet per pattern.
    /// - Crypt-blokken vormen één CBC-keten; skip-blokken blijven clear en updaten de IV niet.
    /// - Pattern 0:0 = alle volle blokken encrypted (geen pattern).
    /// - Partiële trailing bytes blijven clear.
    static func cryptCbcs(
        key: Data,
        iv: Data,
        data: Data,
        cryptByteBlock: Int,
        skipByteBlock: Int
    ) -> Data? {
        guard key.count == 16 else { return nil }
        if data.isEmpty { return data }

        let blockSize = 16
        let totalBlocks = data.count / blockSize
        guard totalBlocks > 0 else { return data }

        var cryptBlocks = cryptByteBlock
        var skipBlocks = skipByteBlock
        // ISO/Chromium: 0:0 → alles encrypted als één CBC-run.
        if cryptBlocks == 0 && skipBlocks == 0 {
            cryptBlocks = totalBlocks
            skipBlocks = 0
        }
        if cryptBlocks <= 0 { return data }

        var out = data
        var blockIv = Data(count: 16)
        let ivn = min(iv.count, 16)
        blockIv.replaceSubrange(0..<ivn, with: iv.prefix(ivn))

        var blocksProcessed = 0
        var offset = 0
        // Chromium: start met encrypted run (toggle van false → true).
        var isCryptRun = false

        while blocksProcessed < totalBlocks {
            isCryptRun = !isCryptRun
            let runSize = isCryptRun ? cryptBlocks : skipBlocks
            if runSize == 0 { continue }

            let blocksToProcess = min(runSize, totalBlocks - blocksProcessed)
            if blocksToProcess == 0 { continue }

            let byteCount = blocksToProcess * blockSize
            let range = offset..<(offset + byteCount)

            if isCryptRun {
                let cipher = out.subdata(in: range)
                guard let plain = aesCbcDecryptNoPad(key: key, iv: blockIv, data: cipher) else {
                    return nil
                }
                out.replaceSubrange(range, with: plain)
                // CBC-keten over crypt-runs: IV = laatste ciphertext-blok.
                if cipher.count >= blockSize {
                    blockIv.replaceSubrange(0..<blockSize, with: cipher.suffix(blockSize))
                }
            }
            // Skip-run: clear laten; IV niet aanpassen.

            offset += byteCount
            blocksProcessed += blocksToProcess
        }
        return out
    }

    private static func aesCbcDecryptNoPad(key: Data, iv: Data, data: Data) -> Data? {
        guard data.count % 16 == 0, data.count > 0 else { return data.isEmpty ? data : nil }
        var out = Data(count: data.count)
        var moved = 0
        var ivMut = iv
        if ivMut.count < 16 { ivMut.append(Data(count: 16 - ivMut.count)) }
        let status = out.withUnsafeMutableBytes { outBytes in
            data.withUnsafeBytes { inBytes in
                key.withUnsafeBytes { keyBytes in
                    ivMut.withUnsafeMutableBytes { ivBytes in
                        CCCrypt(
                            CCOperation(kCCDecrypt),
                            CCAlgorithm(kCCAlgorithmAES),
                            CCOptions(0),
                            keyBytes.baseAddress,
                            key.count,
                            ivBytes.baseAddress,
                            inBytes.baseAddress,
                            data.count,
                            outBytes.baseAddress,
                            data.count,
                            &moved
                        )
                    }
                }
            }
        }
        guard status == kCCSuccess, moved == data.count else { return nil }
        return out
    }
}