//! ML-DSA scheme parameters (FIPS 204 §4 + Table 2). The per-set
//! `const` blocks ([`ML_DSA_44`], [`ML_DSA_65`], [`ML_DSA_87`]) carry
//! the matrix shape, eta, gamma1/gamma2, tau, omega, and derived
//! encoding byte lengths.

use subtle::{Choice, ConditionallySelectable};

/// Modulus q = 2^23 - 2^13 + 1.
pub const Q: u32 = 8_380_417;

/// `-Q^-1 mod 2^32`, used by `wide::mul` Montgomery REDC.
pub const Q_N_PRIME: u32 = 4_236_238_847;

/// `R^2 mod Q` = `2^64 mod Q`, used by canonical-to-Mont conversion.
pub const Q_R2_MOD_Q: u32 = 2_365_951;

/// Polynomial degree.
pub const N: usize = 256;

/// 512-th principal root of unity mod q; only needed by NTT test oracles.
#[cfg(test)]
pub const ZETA: u32 = 1753;

/// 2^d split used in Power2Round.
pub const D: u32 = 13;

// ML-DSA stores canonical `[0, q)` reps but reasons about norms in
// centered `(-q/2, q/2]` form.

/// Centered `(-q/2, q/2]` → canonical `[0, q)`. Correct for
/// `|x| < 2*q`; all in-tree callers pass `|x| < q`.
#[inline]
pub fn from_signed(x: i32, q: u32) -> u32 {
    // Arithmetic right shift fills with the sign bit; & 1 isolates it.
    let is_neg = Choice::from(((x >> 31) as u32 & 1) as u8);
    let abs = x.wrapping_abs() as u32;
    let geq_q = Choice::from((((abs.wrapping_sub(q)) >> 31) ^ 1) as u8 & 1);
    let abs_red = u32::conditional_select(&abs, &abs.wrapping_sub(q), geq_q);
    // x | wrapping_neg(x) has the high bit set for any nonzero x.
    let is_zero = Choice::from((((abs_red | abs_red.wrapping_neg()) >> 31) ^ 1) as u8);
    let neg_val = u32::conditional_select(&q.wrapping_sub(abs_red), &0u32, is_zero);
    u32::conditional_select(&abs_red, &neg_val, is_neg)
}

/// Canonical `[0, q)` → centered `(-q/2, q/2]`.
#[inline]
pub fn to_signed(x: u32, q: u32) -> i32 {
    let in_high = Choice::from((((q / 2).wrapping_sub(x) >> 31) & 1) as u8);
    u32::conditional_select(&x, &x.wrapping_sub(q), in_high) as i32
}

/// Absolute value of the centered representative: `min(x, q - x)` for
/// `x < q`.
#[inline]
pub fn abs_centered(x: u32, q: u32) -> u32 {
    let q_minus_x = q.wrapping_sub(x);
    let take_q_minus_x = Choice::from(((q_minus_x.wrapping_sub(x) >> 31) & 1) as u8);
    u32::conditional_select(&x, &q_minus_x, take_q_minus_x)
}

/// FIPS 204 `eta`. Closed enum so an unsupported value can't reach
/// the per-byte rejection samplers, where it would otherwise be a
/// runtime panic.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Eta {
    /// `eta = 2`, used by ML-DSA-44 and ML-DSA-87.
    Eta2,
    /// `eta = 4`, used by ML-DSA-65.
    Eta4,
}

impl Eta {
    #[inline]
    pub const fn value(self) -> u32 {
        match self {
            Eta::Eta2 => 2,
            Eta::Eta4 => 4,
        }
    }
}

/// FIPS 204 `gamma2`. Closed enum so the Barrett reduction inside
/// the `decompose` / `use_hint` rounding helpers never silently falls
/// through to the wrong denominator.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Gamma2 {
    /// `(Q - 1) / 88`, used by ML-DSA-44.
    G88,
    /// `(Q - 1) / 32`, used by ML-DSA-65 and ML-DSA-87.
    G32,
}

impl Gamma2 {
    #[inline]
    pub const fn value(self) -> u32 {
        match self {
            Gamma2::G88 => (Q - 1) / 88,
            Gamma2::G32 => (Q - 1) / 32,
        }
    }

    /// `floor(2^32 / (2 * gamma2))` — the Barrett factor for the
    /// `% (2 * gamma2)` reduction in `Decompose`.
    #[inline]
    pub const fn barrett_factor(self) -> u32 {
        ((1u64 << 32) / (2 * self.value() as u64)) as u32
    }

    /// `(q - 1) / (2 * gamma2)` — the `m` in `UseHint`'s modular
    /// adjustments.
    #[inline]
    pub const fn use_hint_m(self) -> u32 {
        (Q - 1) / (2 * self.value())
    }
}

/// Per-set ML-DSA constants. `K` (rows) and `L` (columns) are const
/// generics so `PolyVec` / `PolyMatrix` sizes specialize at compile
/// time.
pub struct Params<const K: usize, const L: usize> {
    pub eta: Eta,
    pub tau: usize,
    pub beta: u32,
    pub gamma1: u32,
    pub gamma1_bits: usize,
    pub gamma2: Gamma2,
    pub omega: usize,
    pub w1_bits: usize,
    pub pk_bytes: usize,
    pub sk_bytes: usize,
    pub sig_bytes: usize,
    pub ctilde_bytes: usize,
}

/// SHAKE-256 expansion buffer for keygen's seed bundle
/// `(rho | rho_prime | big_k)`.
pub const SEED_EXPAND_BYTES: usize = 128;

/// Worst-case ctilde length across all three sets (`lambda/4` at
/// `lambda = 256`).
pub const MAX_CTILDE_BYTES: usize = 64;

/// Worst-case packed-w1 length across all three sets.
pub const MAX_W1_PACKED_BYTES: usize = 1024;

const fn eta_bits(eta: u32) -> usize {
    let mut a = 2 * eta;
    let mut n = 0;
    while a > 0 {
        n += 1;
        a >>= 1;
    }
    n
}

const fn pk_bytes(k: usize) -> usize {
    32 + 32 * k * 10
}
const fn sk_bytes(k: usize, l: usize, eta: u32) -> usize {
    32 + 32 + 64 + 32 * ((k + l) * eta_bits(eta) + (D as usize) * k)
}
const fn sig_bytes(k: usize, l: usize, lambda: usize, gamma1_bits: usize, omega: usize) -> usize {
    lambda / 4 + l * 32 * (1 + gamma1_bits) + omega + k
}

/// ML-DSA-44 (k=4, l=4, eta=2, gamma1=2^17, gamma2=(q-1)/88).
pub const ML_DSA_44: Params<4, 4> = Params {
    eta: Eta::Eta2,
    tau: 39,
    beta: 39 * 2,
    gamma1: 1 << 17,
    gamma1_bits: 17,
    gamma2: Gamma2::G88,
    omega: 80,
    w1_bits: 6,
    pk_bytes: pk_bytes(4),
    sk_bytes: sk_bytes(4, 4, 2),
    sig_bytes: sig_bytes(4, 4, 128, 17, 80),
    ctilde_bytes: 128 / 4,
};

/// ML-DSA-65 (k=6, l=5, eta=4, gamma1=2^19, gamma2=(q-1)/32).
pub const ML_DSA_65: Params<6, 5> = Params {
    eta: Eta::Eta4,
    tau: 49,
    beta: 49 * 4,
    gamma1: 1 << 19,
    gamma1_bits: 19,
    gamma2: Gamma2::G32,
    omega: 55,
    w1_bits: 4,
    pk_bytes: pk_bytes(6),
    sk_bytes: sk_bytes(6, 5, 4),
    sig_bytes: sig_bytes(6, 5, 192, 19, 55),
    ctilde_bytes: 192 / 4,
};

/// ML-DSA-87 (k=8, l=7, eta=2, gamma1=2^19, gamma2=(q-1)/32).
pub const ML_DSA_87: Params<8, 7> = Params {
    eta: Eta::Eta2,
    tau: 60,
    beta: 60 * 2,
    gamma1: 1 << 19,
    gamma1_bits: 19,
    gamma2: Gamma2::G32,
    omega: 75,
    w1_bits: 4,
    pk_bytes: pk_bytes(8),
    sk_bytes: sk_bytes(8, 7, 2),
    sig_bytes: sig_bytes(8, 7, 256, 19, 75),
    ctilde_bytes: 256 / 4,
};

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

    #[test]
    fn sizes_44() {
        assert_eq!(ML_DSA_44.pk_bytes, 1312);
        assert_eq!(ML_DSA_44.sk_bytes, 2560);
        assert_eq!(ML_DSA_44.sig_bytes, 2420);
        assert_eq!(ML_DSA_44.ctilde_bytes, 32);
    }

    #[test]
    fn sizes_65() {
        assert_eq!(ML_DSA_65.pk_bytes, 1952);
        assert_eq!(ML_DSA_65.sk_bytes, 4032);
        assert_eq!(ML_DSA_65.sig_bytes, 3309);
        assert_eq!(ML_DSA_65.ctilde_bytes, 48);
    }

    #[test]
    fn sizes_87() {
        assert_eq!(ML_DSA_87.pk_bytes, 2592);
        assert_eq!(ML_DSA_87.sk_bytes, 4896);
        assert_eq!(ML_DSA_87.sig_bytes, 4627);
        assert_eq!(ML_DSA_87.ctilde_bytes, 64);
    }

    #[test]
    fn w1_packed_bound() {
        for &(k, w1) in &[(4, 6), (6, 4), (8, 4)] {
            assert!(k * 32 * w1 <= MAX_W1_PACKED_BYTES);
        }
    }

    #[test]
    fn signed_roundtrip() {
        for x in [0, 1, 100, Q / 2 - 1, Q / 2, Q / 2 + 1, Q - 1] {
            let s = to_signed(x, Q);
            let r = from_signed(s, Q);
            assert_eq!(r, x);
        }
    }
}
