// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (C) 2026 Allen Sarkisyan

/// Combines two numbers in GF(2^8) (XOR)
#[inline(always)]
pub(crate) const fn add(a: u8, b: u8) -> u8 {
    a ^ b
}

/// Multiplies two numbers in GF(2^8) - Branchless GF(2^8) multiplication
#[inline(always)]
pub(crate) const fn mult(a: u8, b: u8) -> u8 {
    let mut r = 0u8;
    let mut i = 8u8;

    // Iterate from MSB (7) down to LSB (0)
    while i > 0 {
        i -= 1;

        // 0x00 or 0xFF
        let mask = 0u8.wrapping_sub((b >> i) & 1);
        let reduce = 0u8.wrapping_sub((r >> 7) & 1);

        r = (mask & a) ^ (reduce & 0x1B) ^ (r << 1);
    }
    r
}

/// Calculates the inverse of a number in GF(2^8) (AES field / a^254)
/// Equivalent to a^254. Uses the optimized 11-multiplication chain.
/// The inverse of `a` in GF(2^8) is `a^254`.
#[inline(always)]
const fn inverse_11x(a: u8) -> u8 {
    let mut b = mult(a, a);
    let mut c = mult(a, b);

    b = mult(c, c);
    b = mult(b, b);
    c = mult(b, c);
    b = mult(b, b);
    b = mult(b, b);
    b = mult(b, c);
    b = mult(b, b);
    b = mult(a, b);

    mult(b, b)
}

/// Helper that generates the 256-byte inverse lookup table at compile time.
const fn generate_inverse_table() -> [u8; 256] {
    let mut table = [0u8; 256];
    let mut i = 0usize;
    while i < 256 {
        table[i] = inverse_11x(i as u8);
        i += 1;
    }
    table
}

/// 256-byte lookup table for multiplicative inverse in GF(2^8) (a^254).
const INV_TABLE: [u8; 256] = generate_inverse_table();

/// Calculates the inverse of a number in GF(2^8) (a^254)
///
/// # Performance & Security
///
/// - **Default (enabled)**: Uses a compile-time 256-byte lookup table (`fast-inverse` feature).
///   This is the recommended and fastest option.
/// - **Without `fast-inverse`**: Falls back to the pure-arithmetic 11-multiplication chain
///   (`inverse_11x`). Use `--no-default-features` only if you have an extremely strict
///   "no lookup tables" policy.
///
/// Both implementations are constant-time and side-channel resistant because the
/// input to `inverse` is always derived from **public** share IDs (x-coordinates).
#[inline(always)]
pub(crate) const fn inverse(a: u8) -> u8 {
    #[cfg(feature = "fast-inverse")]
    {
        INV_TABLE[a as usize]
    }
    #[cfg(not(feature = "fast-inverse"))]
    {
        inverse_11x(a)
    }
}

/// Divides two numbers in GF(2^8)
#[inline(always)]
pub(crate) fn div(a: u8, b: u8) -> Option<u8> {
    if b == 0 {
        return None;
    }
    Some(mult(a, inverse(b)))
}

/// Precomputes the Lagrange basis evaluated at x=0
/// Zero-copy internally: writes directly to the provided mutable slice.
#[inline(always)]
pub(crate) fn compute_lagrange_basis_at_zero(x_samples: &[u8], out: &mut [u8]) {
    let n = x_samples.len();
    if n != 0 {
        let mut prod_all = 1u8;
        for &x in x_samples {
            prod_all = mult(prod_all, x);
        }

        for i in 0..n {
            let xi = x_samples[i];

            if xi == 0 {
                out[i] = 1;
            } else {
                let num = mult(prod_all, inverse(xi));
                let mut denom = 1u8;

                #[allow(clippy::needless_range_loop)]
                // Skip j == i because xi ^ xi = 0 in GF(256), which would zero the denominator product
                for j in 0..n {
                    if i != j {
                        denom = mult(denom, xi ^ x_samples[j]);
                    }
                }

                out[i] = mult(num, inverse(denom));
            }
        }
    }
}

#[cfg(test)]
mod shamir_math_tests {
    use super::*;

    #[test]
    fn test_field_add() {
        assert_eq!(add(16, 16), 0);
        assert_eq!(add(3, 4), 7);
    }

    #[test]
    fn test_field_mult() {
        assert_eq!(mult(3, 7), 9);
        assert_eq!(mult(3, 0), 0);
        assert_eq!(mult(0, 3), 0);
    }

    #[test]
    fn test_field_divide() {
        assert_eq!(div(0, 7), Some(0));
        assert_eq!(div(3, 3), Some(1));
        assert_eq!(div(6, 3), Some(2));
        assert_eq!(div(12, 0), None);
    }

    // Exhaustively covers every possible value in GF(2^8) (0..255)
    // Tests all possible pairs of distinct x-coordinates.
    #[test]
    fn test_compute_lagrange_basis_at_zero_all_gf28_values() {
        let mut out = [0u8; 256];
        compute_lagrange_basis_at_zero(&[], &mut out[..0]);

        assert_eq!(out, [0u8; 256]);

        for i in 0u8..255 {
            for j in (i + 1)..=255 {
                let x_samples = [i, j];
                compute_lagrange_basis_at_zero(&x_samples, &mut out[..2]);

                let basis_0 = out[0];
                let basis_1 = out[1];

                // Direct mathematical computation for n=2 at x=0:
                let denom = inverse(i ^ j);
                let expected_0 = mult(j, denom);
                let expected_1 = mult(i, denom);

                assert_eq!(basis_0, expected_0, "Failed for x_samples=[{}, {}]", i, j);
                assert_eq!(basis_1, expected_1, "Failed for x_samples=[{}, {}]", i, j);

                assert_eq!(
                    basis_0 ^ basis_1,
                    1,
                    "Basis sum invariant failed for x_samples=[{}, {}]",
                    i,
                    j
                );
            }
        }
    }

    /// Comprehensive unit test for the reference inverse implementation (`inverse_11x`).
    /// This verifies the mathematical properties of the GF(2^8) inverse (a^254):
    #[test]
    fn test_inverse_11x() {
        // 1. Zero and identity edge cases
        assert_eq!(inverse_11x(0), 0, "inverse(0) must be 0");
        assert_eq!(inverse_11x(1), 1, "inverse(1) must be 1");

        // 2. Exhaustive correctness test over the entire field (256 values)
        for a in 0u8..=255 {
            let inv = inverse_11x(a);

            // a * inv(a) == 1 (when a != 0)
            if a != 0 {
                let product = mult(a, inv);
                assert_eq!(
                    product, 1,
                    "Round-trip failure for a = 0x{:02x}: a * inv(a) = 0x{:02x} (expected 1)",
                    a, product
                );
            }

            // involution: inv(inv(a)) == a
            let inv_inv = inverse_11x(inv);
            assert_eq!(
                inv_inv, a,
                "Involution failure for a = 0x{:02x}: inv(inv(a)) = 0x{:02x} (expected 0x{:02x})",
                a, inv_inv, a
            );
        }

        // 3. Must be identical to the fast lookup table (single source of truth)
        for a in 0u8..=255 {
            assert_eq!(
                inverse_11x(a),
                INV_TABLE[a as usize],
                "inverse_11x(0x{:02x}) does not match INV_TABLE entry",
                a
            );
        }

        // 4. Well-known test vectors from the GF(2^8) field (AES S-box related)
        let known_cases = [
            (0x00, 0x00),
            (0x01, 0x01),
            (0x02, 0x8d),
            (0x03, 0xf6),
            (0x53, 0xca),
            (0xff, 0x1c),
        ];

        for (a, expected) in known_cases {
            let actual = inverse_11x(a);
            assert_eq!(
                actual, expected,
                "Known value failed: inverse(0x{:02x}) = 0x{:02x} (expected 0x{:02x})",
                a, actual, expected
            );
        }

        println!("✅ All tests passed for inverse_11x (256 values fully verified)");
    }

    #[test]
    fn generates_inverse_lut() {
        let gen_lut = generate_inverse_table();
        let inverse_table: [u8; 256] = [
            0, 1, 141, 246, 203, 82, 123, 209, 232, 79, 41, 192, 176, 225, 229, 199, 116, 180, 170,
            75, 153, 43, 96, 95, 88, 63, 253, 204, 255, 64, 238, 178, 58, 110, 90, 241, 85, 77,
            168, 201, 193, 10, 152, 21, 48, 68, 162, 194, 44, 69, 146, 108, 243, 57, 102, 66, 242,
            53, 32, 111, 119, 187, 89, 25, 29, 254, 55, 103, 45, 49, 245, 105, 167, 100, 171, 19,
            84, 37, 233, 9, 237, 92, 5, 202, 76, 36, 135, 191, 24, 62, 34, 240, 81, 236, 97, 23,
            22, 94, 175, 211, 73, 166, 54, 67, 244, 71, 145, 223, 51, 147, 33, 59, 121, 183, 151,
            133, 16, 181, 186, 60, 182, 112, 208, 6, 161, 250, 129, 130, 131, 126, 127, 128, 150,
            115, 190, 86, 155, 158, 149, 217, 247, 2, 185, 164, 222, 106, 50, 109, 216, 138, 132,
            114, 42, 20, 159, 136, 249, 220, 137, 154, 251, 124, 46, 195, 143, 184, 101, 72, 38,
            200, 18, 74, 206, 231, 210, 98, 12, 224, 31, 239, 17, 117, 120, 113, 165, 142, 118, 61,
            189, 188, 134, 87, 11, 40, 47, 163, 218, 212, 228, 15, 169, 39, 83, 4, 27, 252, 172,
            230, 122, 7, 174, 99, 197, 219, 226, 234, 148, 139, 196, 213, 157, 248, 144, 107, 177,
            13, 214, 235, 198, 14, 207, 173, 8, 78, 215, 227, 93, 80, 30, 179, 91, 35, 56, 52, 104,
            70, 3, 140, 221, 156, 125, 160, 205, 26, 65, 28,
        ];

        assert_eq!(INV_TABLE, inverse_table);
        assert_eq!(gen_lut, inverse_table);
    }
}
