import Foundation

// MARK: - Poly1305Engine

/// Poly1305 message authentication code implementation.
/// Used by XChaCha20-Poly1305 AEAD construction.
enum Poly1305Engine {

    /// Computes a Poly1305 MAC matching the AEAD construction.
    nonisolated static func mac(data: Data, key: Data) -> Data {
        // Clamp r (first 16 bytes of key)
        var r = [UInt8](key[0..<16])
        r[3] &= 15; r[7] &= 15; r[11] &= 15; r[15] &= 15
        r[4] &= 252; r[8] &= 252; r[12] &= 252

        // s = last 16 bytes of key
        let s = [UInt8](key[16..<32])

        var accumulator = [UInt64](repeating: 0, count: 5)
        let rLimbs = toLimbs26(r)

        // Build padded data: data + padding to 16-byte boundary + lengths
        var macInput = Data(data)
        let padding = (16 - (data.count % 16)) % 16
        if padding > 0 {
            macInput.append(Data(repeating: 0, count: padding))
        }
        // AAD length (0 — no associated data)
        macInput.append(Data(repeating: 0, count: 8))
        // Ciphertext length (little-endian 64-bit)
        var ctLen = UInt64(data.count).littleEndian
        macInput.append(Data(bytes: &ctLen, count: 8))

        // Process in 16-byte blocks
        for offset in stride(from: 0, to: macInput.count, by: 16) {
            let blockEnd = min(offset + 16, macInput.count)
            var block = [UInt8](macInput[offset..<blockEnd])
            while block.count < 16 { block.append(0) }

            // Poly1305 appends a 1-bit to each complete 16-byte block.
            // In AEAD mode we process a fully padded stream
            // (ciphertext + pad16 + aadLen + ctLen), so every block here is complete.
            let bit: UInt64 = 1 << 24

            let n = toLimbs26(block)
            accumulator[0] = accumulator[0] &+ n[0]
            accumulator[1] = accumulator[1] &+ n[1]
            accumulator[2] = accumulator[2] &+ n[2]
            accumulator[3] = accumulator[3] &+ n[3]
            accumulator[4] = accumulator[4] &+ n[4] &+ bit

            accumulator = multiply(accumulator, rLimbs)
        }

        return Data(freeze(accumulator, s: s))
    }
}

// MARK: - Limb Arithmetic

private extension Poly1305Engine {

    /// Convert 16 bytes to 5 limbs of 26 bits each (little-endian).
    nonisolated static func toLimbs26(_ bytes: [UInt8]) -> [UInt64] {
        let b = bytes.count >= 16 ? bytes : bytes + [UInt8](repeating: 0, count: 16 - bytes.count)

        var full = [UInt8](repeating: 0, count: 17)
        for i in 0..<min(16, b.count) { full[i] = b[i] }

        let lo = UInt64(full[0]) | UInt64(full[1]) << 8 | UInt64(full[2]) << 16 | UInt64(full[3]) << 24 |
                 UInt64(full[4]) << 32 | UInt64(full[5]) << 40 | UInt64(full[6]) << 48 | UInt64(full[7]) << 56
        let hi = UInt64(full[8]) | UInt64(full[9]) << 8 | UInt64(full[10]) << 16 | UInt64(full[11]) << 24 |
                 UInt64(full[12]) << 32 | UInt64(full[13]) << 40 | UInt64(full[14]) << 48 | UInt64(full[15]) << 56

        return [
            lo & 0x3FFFFFF,
            (lo >> 26) & 0x3FFFFFF,
            ((lo >> 52) | (hi << 12)) & 0x3FFFFFF,
            (hi >> 14) & 0x3FFFFFF,
            (hi >> 40) & 0x3FFFFFF
        ]
    }

    /// Multiply two numbers in 26-bit limb form, reduce mod 2^130 - 5.
    nonisolated static func multiply(_ a: [UInt64], _ r: [UInt64]) -> [UInt64] {
        let r0 = r[0], r1 = r[1], r2 = r[2], r3 = r[3], r4 = r[4]
        let s1 = r1 * 5, s2 = r2 * 5, s3 = r3 * 5, s4 = r4 * 5
        let a0 = a[0], a1 = a[1], a2 = a[2], a3 = a[3], a4 = a[4]

        var h0 = a0 * r0 + a1 * s4 + a2 * s3 + a3 * s2 + a4 * s1
        var h1 = a0 * r1 + a1 * r0 + a2 * s4 + a3 * s3 + a4 * s2
        var h2 = a0 * r2 + a1 * r1 + a2 * r0 + a3 * s4 + a4 * s3
        var h3 = a0 * r3 + a1 * r2 + a2 * r1 + a3 * r0 + a4 * s4
        var h4 = a0 * r4 + a1 * r3 + a2 * r2 + a3 * r1 + a4 * r0

        var c: UInt64
        c = h0 >> 26; h1 += c; h0 &= 0x3FFFFFF
        c = h1 >> 26; h2 += c; h1 &= 0x3FFFFFF
        c = h2 >> 26; h3 += c; h2 &= 0x3FFFFFF
        c = h3 >> 26; h4 += c; h3 &= 0x3FFFFFF
        c = h4 >> 26; h0 += c * 5; h4 &= 0x3FFFFFF
        c = h0 >> 26; h1 += c; h0 &= 0x3FFFFFF

        return [h0, h1, h2, h3, h4]
    }

    /// Final reduction mod 2^130-5 and add s.
    nonisolated static func freeze(_ h: [UInt64], s: [UInt8]) -> [UInt8] {
        var h0 = h[0], h1 = h[1], h2 = h[2], h3 = h[3], h4 = h[4]

        var c: UInt64
        c = h0 >> 26; h1 += c; h0 &= 0x3FFFFFF
        c = h1 >> 26; h2 += c; h1 &= 0x3FFFFFF
        c = h2 >> 26; h3 += c; h2 &= 0x3FFFFFF
        c = h3 >> 26; h4 += c; h3 &= 0x3FFFFFF
        c = h4 >> 26; h0 += c * 5; h4 &= 0x3FFFFFF
        c = h0 >> 26; h1 += c; h0 &= 0x3FFFFFF

        // Compute h + -(2^130 - 5) = h - p
        var g0 = h0 &+ 5; c = g0 >> 26; g0 &= 0x3FFFFFF
        var g1 = h1 &+ c; c = g1 >> 26; g1 &= 0x3FFFFFF
        var g2 = h2 &+ c; c = g2 >> 26; g2 &= 0x3FFFFFF
        var g3 = h3 &+ c; c = g3 >> 26; g3 &= 0x3FFFFFF
        let g4 = h4 &+ c &- (1 << 26)

        // If g4 didn't underflow, use g (h >= p)
        let mask = (g4 >> 63) &- 1
        let nmask = ~mask
        h0 = (h0 & nmask) | (g0 & mask)
        h1 = (h1 & nmask) | (g1 & mask)
        h2 = (h2 & nmask) | (g2 & mask)
        h3 = (h3 & nmask) | (g3 & mask)
        h4 = (h4 & nmask) | (g4 & mask)

        // Convert to two 64-bit values
        let lo = h0 | (h1 << 26) | (h2 << 52)
        let hi = (h2 >> 12) | (h3 << 14) | (h4 << 40)

        // Add s (little-endian)
        var sLo: UInt64 = 0
        var sHi: UInt64 = 0
        for i in stride(from: 7, through: 0, by: -1) { sLo = sLo << 8 | UInt64(s[i]) }
        for i in stride(from: 15, through: 8, by: -1) { sHi = sHi << 8 | UInt64(s[i]) }

        let resultLo = lo &+ sLo
        let carry: UInt64 = resultLo < lo ? 1 : 0
        let resultHi = hi &+ sHi &+ carry

        var output = [UInt8](repeating: 0, count: 16)
        for i in 0..<8 { output[i] = UInt8(truncatingIfNeeded: resultLo >> (i * 8)) }
        for i in 0..<8 { output[8 + i] = UInt8(truncatingIfNeeded: resultHi >> (i * 8)) }
        return output
    }
}
