import Foundation

/// Pure Swift GF(2^255-19) arithmetic for Ed25519 → X25519 public key conversion.
///
/// The conversion formula is: u = (1 + y) / (1 - y) mod p
/// where p = 2^255 - 19, and y is the Ed25519 public key's y-coordinate.
///
/// Uses 4-limb UInt64 representation (little-endian).
enum FieldArithmetic {

    // p = 2^255 - 19
    static let p: [UInt64] = [
        0xFFFF_FFFF_FFFF_FFED,  // limb 0 (least significant)
        0xFFFF_FFFF_FFFF_FFFF,  // limb 1
        0xFFFF_FFFF_FFFF_FFFF,  // limb 2
        0x7FFF_FFFF_FFFF_FFFF,  // limb 3 (most significant, 2^63 - 1 accounting for -19)
    ]

    /// Load a 256-bit little-endian byte array into 4 UInt64 limbs
    static func load(_ bytes: Data) -> [UInt64] {
        precondition(bytes.count == 32)
        var limbs = [UInt64](repeating: 0, count: 4)
        for i in 0..<4 {
            var val: UInt64 = 0
            for j in 0..<8 {
                val |= UInt64(bytes[i * 8 + j]) << (j * 8)
            }
            limbs[i] = val
        }
        return limbs
    }

    /// Store 4 UInt64 limbs as 32 little-endian bytes
    static func store(_ limbs: [UInt64]) -> Data {
        var bytes = Data(count: 32)
        for i in 0..<4 {
            for j in 0..<8 {
                bytes[i * 8 + j] = UInt8((limbs[i] >> (j * 8)) & 0xFF)
            }
        }
        return bytes
    }

    /// a + b mod p
    static func add(_ a: [UInt64], _ b: [UInt64]) -> [UInt64] {
        var result = [UInt64](repeating: 0, count: 4)
        var carry: UInt64 = 0
        for i in 0..<4 {
            let (sum1, c1) = a[i].addingReportingOverflow(b[i])
            let (sum2, c2) = sum1.addingReportingOverflow(carry)
            result[i] = sum2
            carry = (c1 ? 1 : 0) + (c2 ? 1 : 0)
        }
        // Reduce mod p
        return reduceOnce(result, carry: carry)
    }

    /// a - b mod p
    static func sub(_ a: [UInt64], _ b: [UInt64]) -> [UInt64] {
        var result = [UInt64](repeating: 0, count: 4)
        var borrow: UInt64 = 0
        for i in 0..<4 {
            let (diff1, b1) = a[i].subtractingReportingOverflow(b[i])
            let (diff2, b2) = diff1.subtractingReportingOverflow(borrow)
            result[i] = diff2
            borrow = (b1 ? 1 : 0) + (b2 ? 1 : 0)
        }
        if borrow > 0 {
            // Add p back
            var c: UInt64 = 0
            for i in 0..<4 {
                let (s1, c1) = result[i].addingReportingOverflow(p[i])
                let (s2, c2) = s1.addingReportingOverflow(c)
                result[i] = s2
                c = (c1 ? 1 : 0) + (c2 ? 1 : 0)
            }
        }
        return result
    }

    /// Multiply two 256-bit numbers mod p using schoolbook multiplication
    static func mul(_ a: [UInt64], _ b: [UInt64]) -> [UInt64] {
        // Full 512-bit product in 8 limbs
        var product = [UInt64](repeating: 0, count: 8)

        for i in 0..<4 {
            var carry: UInt64 = 0
            for j in 0..<4 {
                let (hi, lo) = a[i].multipliedFullWidth(by: b[j])
                let (sum1, c1) = product[i + j].addingReportingOverflow(lo)
                let (sum2, c2) = sum1.addingReportingOverflow(carry)
                product[i + j] = sum2
                carry = hi + (c1 ? 1 : 0) + (c2 ? 1 : 0)
            }
            product[i + 4] = carry
        }

        // Reduce mod p using Barrett-like reduction
        // Since p = 2^255 - 19, for a 512-bit number we can use:
        // x mod p = (x_low + x_high * 2^256) mod p
        // Since 2^255 ≡ 19 (mod p), 2^256 ≡ 38 (mod p)
        return reduceFull(product)
    }

    /// Reduce 512-bit product mod p using 2^256 ≡ 38 (mod p)
    private static func reduceFull(_ product: [UInt64]) -> [UInt64] {
        // Split: low = product[0..3], high = product[4..7]
        // result = low + high * 38
        var result = [UInt64](repeating: 0, count: 5)

        // Start with low part
        for i in 0..<4 {
            result[i] = product[i]
        }

        // Add high * 38
        var carry: UInt64 = 0
        for i in 0..<4 {
            let (hi, lo) = product[i + 4].multipliedFullWidth(by: 38)
            let (sum1, c1) = result[i].addingReportingOverflow(lo)
            let (sum2, c2) = sum1.addingReportingOverflow(carry)
            result[i] = sum2
            carry = hi + (c1 ? 1 : 0) + (c2 ? 1 : 0)
        }
        result[4] = carry

        // The result might still be >= p, so reduce once more
        // result[4] * 2^256 ≡ result[4] * 38 (mod p)
        var extra: UInt64 = result[4]
        result[4] = 0
        if extra > 0 {
            let (hi, lo) = extra.multipliedFullWidth(by: 38)
            let (sum1, c1) = result[0].addingReportingOverflow(lo)
            result[0] = sum1
            var c = hi + (c1 ? 1 : 0)
            for i in 1..<4 {
                let (s, cf) = result[i].addingReportingOverflow(c)
                result[i] = s
                c = cf ? 1 : 0
            }
            // One more round if carry
            if c > 0 {
                let (s, _) = result[0].addingReportingOverflow(c * 38)
                result[0] = s
            }
        }

        var out = Array(result[0..<4])
        // Final reduction: if >= p, subtract p
        out = reduceOnce(out, carry: 0)
        return out
    }

    /// If the number >= p, subtract p
    private static func reduceOnce(_ val: [UInt64], carry: UInt64) -> [UInt64] {
        if carry > 0 || isGreaterOrEqual(val, p) {
            var result = [UInt64](repeating: 0, count: 4)
            var borrow: UInt64 = 0
            for i in 0..<4 {
                let (diff1, b1) = val[i].subtractingReportingOverflow(p[i])
                let (diff2, b2) = diff1.subtractingReportingOverflow(borrow)
                result[i] = diff2
                borrow = (b1 ? 1 : 0) + (b2 ? 1 : 0)
            }
            // If borrow after subtracting p, the original was fine (shouldn't happen with carry)
            if borrow > 0 && carry == 0 {
                return val
            }
            return result
        }
        return val
    }

    /// Compare a >= b
    private static func isGreaterOrEqual(_ a: [UInt64], _ b: [UInt64]) -> Bool {
        for i in stride(from: 3, through: 0, by: -1) {
            if a[i] > b[i] { return true }
            if a[i] < b[i] { return false }
        }
        return true  // equal
    }

    /// Modular inverse using Fermat's little theorem: a^(-1) = a^(p-2) mod p
    static func inverse(_ a: [UInt64]) -> [UInt64] {
        // p - 2 = 2^255 - 21
        let pMinus2 = sub(p, [2, 0, 0, 0])
        return power(a, pMinus2)
    }

    /// Modular exponentiation using square-and-multiply
    static func power(_ base: [UInt64], _ exp: [UInt64]) -> [UInt64] {
        var result: [UInt64] = [1, 0, 0, 0]  // 1
        var b = base

        for i in 0..<4 {
            var limb = exp[i]
            let bits = (i == 3) ? 63 : 64  // top limb has 63 bits for p-2
            for _ in 0..<bits {
                if limb & 1 == 1 {
                    result = mul(result, b)
                }
                b = mul(b, b)
                limb >>= 1
            }
        }
        return result
    }

    // MARK: - Ed25519 → X25519 Public Key Conversion

    /// Convert Ed25519 public key (32 bytes) to X25519 public key (32 bytes).
    /// Formula: u = (1 + y) * inverse(1 - y) mod p
    static func ed25519PublicToX25519(_ ed25519Pub: Data) -> Data {
        precondition(ed25519Pub.count == 32)

        // Ed25519 public key is the y-coordinate with sign bit in the top bit of byte 31
        var keyBytes = ed25519Pub
        // Clear the sign bit
        keyBytes[31] &= 0x7F

        let y = load(keyBytes)
        let one: [UInt64] = [1, 0, 0, 0]

        let onePlusY = add(one, y)
        let oneMinusY = sub(one, y)
        let inv = inverse(oneMinusY)
        let u = mul(onePlusY, inv)

        return store(u)
    }
}
