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 `crypt` 16-byte blocks, skip `skip` blocks; leftover &lt;16 clear.
    static func cryptCbcs(
        key: Data,
        iv: Data,
        data: Data,
        cryptByteBlock: Int,
        skipByteBlock: Int
    ) -> Data? {
        guard key.count == 16 else { return nil }
        if cryptByteBlock <= 0 {
            return data
        }
        let cryptLen = 16 * cryptByteBlock
        let skipLen = 16 * max(skipByteBlock, 0)
        var result = data
        var offset = 0
        var blockIv = Data(count: 16)
        let ivCopy = min(iv.count, 16)
        blockIv.replaceSubrange(0..<ivCopy, with: iv.prefix(ivCopy))

        while offset < result.count {
            let remaining = result.count - offset
            if remaining < cryptLen { break }
            let slice = result.subdata(in: offset..<(offset + cryptLen))
            guard let plain = aesCbcDecryptNoPad(key: key, iv: blockIv, data: slice) else { return nil }
            result.replaceSubrange(offset..<(offset + cryptLen), with: plain)
            // Next IV = last ciphertext block (CBC chain within crypt run).
            blockIv = slice.suffix(16)
            offset += cryptLen
            let skip = min(skipLen, result.count - offset)
            offset += skip
            // Pattern: IV resets per subsample in caller; within subsample after skip, continue with updated IV.
        }
        return result
    }

    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
    }
}