// stdlib/crypto/kpke.sy — K-PKE: Kyber's inner public-key encryption
//
// This implements the IND-CPA-secure public-key encryption scheme that
// ML-KEM's Fujisaki-Okamoto transform wraps for CCA security.
//
// K-PKE.KeyGen(d):     (d is 32-byte seed) -> (ek, dk)
// K-PKE.Encrypt(ek, m, r): encrypt 32-byte message m with randomness r
// K-PKE.Decrypt(dk, ct): decrypt ciphertext to recover m
//
// For Kyber-768: k=3, eta1=2, eta2=2, du=10, dv=4
//
// References:
//   - FIPS 203 Algorithms 4, 5, 6
//   - Kyber reference: indcpa.c
//
// Aislinn McCubbins, Apr 13 2026

use crypto::fieldq
use crypto::ntt
use crypto::sample

// ── Kyber-768 parameters ─────────────────────────────────────────
fn kyber_k() -> Int { return 3 }
fn kyber_eta1() -> Int { return 2 }
fn kyber_eta2() -> Int { return 2 }
fn kyber_du() -> Int { return 10 }
fn kyber_dv() -> Int { return 4 }

// ── Polynomial compression/decompression ─────────────────────────
// Compress: round(2^d / q * x) mod 2^d
// Decompress: round(q / 2^d * x)

fn compress(x: Int, d: Int) -> Int {
    // round((2^d * x + q/2) / q) mod 2^d
    let two_d = shl(1, d)
    let val = (two_d * x + fq_q() / 2) / fq_q()
    return bit_and(val, two_d - 1)
}

fn decompress(x: Int, d: Int) -> Int {
    // round(q * x / 2^d) = (q * x + 2^(d-1)) / 2^d
    return (fq_q() * x + shl(1, d - 1)) / shl(1, d)
}

// ── Encode/Decode polynomials to byte tensors ────────────────────
// Encode a polynomial with d-bit coefficients into a byte tensor.
// Each coefficient is d bits; pack into bytes LSB-first.

fn poly_encode(poly: Tensor, d: Int) -> Tensor {
    let num_bytes = 256 * d / 8
    let mut out = t_zeros(num_bytes)
    let mut bit_pos = 0
    let mut i = 0
    while i < 256 {
        let coeff = t_index_get(poly, i)
        let mut j = 0
        while j < d {
            let bit_val = bit_and(shr(coeff, j), 1)
            let byte_idx = bit_pos / 8
            let bit_idx = bit_pos % 8
            let old_byte = t_index_get(out, byte_idx)
            out = t_index_set(out, byte_idx, old_byte + shl(bit_val, bit_idx) + 0.0)
            bit_pos = bit_pos + 1
            j = j + 1
        }
        i = i + 1
    }
    return out
}

fn poly_decode(bytes: Tensor, d: Int) -> Tensor {
    let mut poly = t_zeros(256)
    let mut bit_pos = 0
    let mut i = 0
    while i < 256 {
        let mut coeff = 0
        let mut j = 0
        while j < d {
            let byte_idx = bit_pos / 8
            let bit_idx = bit_pos % 8
            let byte_val = t_index_get(bytes, byte_idx)
            let bit_val = bit_and(shr(byte_val, bit_idx), 1)
            coeff = coeff + shl(bit_val, j)
            bit_pos = bit_pos + 1
            j = j + 1
        }
        poly = t_index_set(poly, i, coeff + 0.0)
        i = i + 1
    }
    return poly
}

// ── Encode 12-bit NTT coefficients (standard Kyber format) ───────
fn poly_encode12(poly: Tensor) -> Tensor {
    return poly_encode(poly, 12)
}

fn poly_decode12(bytes: Tensor) -> Tensor {
    return poly_decode(bytes, 12)
}

// ── Vector operations (k=3 polynomials) ──────────────────────────
// A "vector" is an Array of k Tensors (polynomials).

fn poly_add(a: Tensor, b: Tensor) -> Tensor {
    let mut r = t_zeros(256)
    let mut i = 0
    while i < 256 {
        let sum = t_index_get(a, i) + t_index_get(b, i)
        r = t_index_set(r, i, fq_barrett(sum) + 0.0)
        i = i + 1
    }
    return r
}

// Convert each coefficient to Montgomery form: coeff -> coeff * R mod q
// Required after basemul accumulation to make it compatible with NTT'd noise.
// See Kyber reference: poly_tomont() in poly.c
fn poly_tomont(p: Tensor) -> Tensor {
    let mont = fq_mont()
    let mut r = t_zeros(256)
    let mut i = 0
    while i < 256 {
        let v = t_index_get(p, i)
        r = t_index_set(r, i, fq_mul(v, mont) + 0.0)
        i = i + 1
    }
    return r
}

// Reduce all coefficients to [0, q) — needed after INTT which leaves values in [-q, q)
fn poly_reduce(p: Tensor) -> Tensor {
    let mut r = t_zeros(256)
    let mut i = 0
    while i < 256 {
        r = t_index_set(r, i, fq_barrett(t_index_get(p, i) + fq_q()) + 0.0)
        i = i + 1
    }
    return r
}

fn poly_sub(a: Tensor, b: Tensor) -> Tensor {
    let mut r = t_zeros(256)
    let mut i = 0
    while i < 256 {
        let diff = t_index_get(a, i) - t_index_get(b, i)
        r = t_index_set(r, i, fq_barrett(diff + fq_q()) + 0.0)
        i = i + 1
    }
    return r
}

// ── K-PKE.KeyGen (FIPS 203 Algorithm 4) ──────────────────────────
// Input: 32-byte seed d
// Output: [ek_bytes, dk_bytes] as an Array of two Tensors
//
// 1. (rho, sigma) = G(d) = SHA3-512(d) split into two 32-byte halves
//    (We use SHA3-256 twice with domain separation as a stand-in)
// 2. Sample matrix A from rho (in NTT domain)
// 3. Sample secret vector s from sigma (CBD, then NTT)
// 4. Sample error vector e from sigma (CBD, then NTT)
// 5. t = A * s + e (in NTT domain)
// 6. ek = encode(t) || rho, dk = encode(s)

fn kpke_keygen(d: Tensor) -> Array {
    // Derive rho (for matrix A) and sigma (for noise) from seed
    // FIPS 203 uses SHA3-512; we approximate with two SHA3-256 calls
    let rho_input = t_cat(d, t_zeros(1))
    let rho = sha3_256(rho_input)
    let sigma_input_byte = t_index_set(t_zeros(1), 0, 1.0)
    let sigma = sha3_256(t_cat(d, sigma_input_byte))

    // Sample matrix A (k x k, each element is a polynomial in NTT domain)
    // A[i][j] = sample_ntt(rho, i, j)
    let a00 = sample_ntt(rho, 0, 0)
    let a01 = sample_ntt(rho, 0, 1)
    let a02 = sample_ntt(rho, 0, 2)
    let a10 = sample_ntt(rho, 1, 0)
    let a11 = sample_ntt(rho, 1, 1)
    let a12 = sample_ntt(rho, 1, 2)
    let a20 = sample_ntt(rho, 2, 0)
    let a21 = sample_ntt(rho, 2, 1)
    let a22 = sample_ntt(rho, 2, 2)

    // Sample secret vector s (k polynomials, CBD then NTT)
    let s0 = ntt_forward(sample_noise_cbd2(sigma, 0))
    let s1 = ntt_forward(sample_noise_cbd2(sigma, 1))
    let s2 = ntt_forward(sample_noise_cbd2(sigma, 2))

    // Sample error vector e (k polynomials, CBD then NTT)
    let e0 = ntt_forward(sample_noise_cbd2(sigma, 3))
    let e1 = ntt_forward(sample_noise_cbd2(sigma, 4))
    let e2 = ntt_forward(sample_noise_cbd2(sigma, 5))

    // t = A * s + e (matrix-vector multiply in NTT domain)
    // poly_tomont converts basemul output to Montgomery form, matching the NTT'd error.
    let t0_as = poly_tomont(poly_add(ntt_basemul(a00, s0), poly_add(ntt_basemul(a01, s1), ntt_basemul(a02, s2))))
    let t1_as = poly_tomont(poly_add(ntt_basemul(a10, s0), poly_add(ntt_basemul(a11, s1), ntt_basemul(a12, s2))))
    let t2_as = poly_tomont(poly_add(ntt_basemul(a20, s0), poly_add(ntt_basemul(a21, s1), ntt_basemul(a22, s2))))
    let t0 = poly_add(t0_as, e0)
    let t1 = poly_add(t1_as, e1)
    let t2 = poly_add(t2_as, e2)

    // Encode public key: ek = encode12(t0) || encode12(t1) || encode12(t2) || rho
    let ek = t_cat(t_cat(t_cat(poly_encode12(t0), poly_encode12(t1)), poly_encode12(t2)), rho)

    // Encode secret key: dk = encode12(s0) || encode12(s1) || encode12(s2)
    let dk = t_cat(t_cat(poly_encode12(s0), poly_encode12(s1)), poly_encode12(s2))

    let mut result = array_new()
    result = array_push(result, ek)
    result = array_push(result, dk)
    return result
}

// ── K-PKE.Encrypt (FIPS 203 Algorithm 5) ──────────────────────────
// Input: ek (public key bytes), m (32-byte message), r (32-byte randomness)
// Output: ciphertext bytes

fn kpke_encrypt(ek: Tensor, m: Tensor, r: Tensor) -> Tensor {
    // Decode public key
    let t0 = poly_decode12(tensor_slice(ek, 0, 384))
    let t1 = poly_decode12(tensor_slice(ek, 384, 768))
    let t2 = poly_decode12(tensor_slice(ek, 768, 1152))
    let rho = tensor_slice(ek, 1152, 1184)

    // Re-sample matrix A^T (transpose: swap i,j)
    let at00 = sample_ntt(rho, 0, 0)
    let at01 = sample_ntt(rho, 1, 0)
    let at02 = sample_ntt(rho, 2, 0)
    let at10 = sample_ntt(rho, 0, 1)
    let at11 = sample_ntt(rho, 1, 1)
    let at12 = sample_ntt(rho, 2, 1)
    let at20 = sample_ntt(rho, 0, 2)
    let at21 = sample_ntt(rho, 1, 2)
    let at22 = sample_ntt(rho, 2, 2)

    // Sample randomness vector r_vec and errors
    let r0 = ntt_forward(sample_noise_cbd2(r, 0))
    let r1 = ntt_forward(sample_noise_cbd2(r, 1))
    let r2 = ntt_forward(sample_noise_cbd2(r, 2))
    let e1_0 = sample_noise_cbd2(r, 3)
    let e1_1 = sample_noise_cbd2(r, 4)
    let e1_2 = sample_noise_cbd2(r, 5)
    let e2 = sample_noise_cbd2(r, 6)

    // u = NTT^-1(A^T * r_vec) + e1
    let u0_ntt = poly_add(ntt_basemul(at00, r0), poly_add(ntt_basemul(at01, r1), ntt_basemul(at02, r2)))
    let u1_ntt = poly_add(ntt_basemul(at10, r0), poly_add(ntt_basemul(at11, r1), ntt_basemul(at12, r2)))
    let u2_ntt = poly_add(ntt_basemul(at20, r0), poly_add(ntt_basemul(at21, r1), ntt_basemul(at22, r2)))
    let u0 = poly_reduce(poly_add(ntt_inverse(u0_ntt), e1_0))
    let u1 = poly_reduce(poly_add(ntt_inverse(u1_ntt), e1_1))
    let u2 = poly_reduce(poly_add(ntt_inverse(u2_ntt), e1_2))

    // v = NTT^-1(t^T * r_vec) + e2 + decompress(m, 1)
    let v_ntt = poly_add(ntt_basemul(t0, r0), poly_add(ntt_basemul(t1, r1), ntt_basemul(t2, r2)))
    let v_no_msg = poly_reduce(poly_add(ntt_inverse(v_ntt), e2))

    // Decode message: each bit of m becomes a coefficient via decompress(bit, 1)
    let mut msg_poly = t_zeros(256)
    let mut mi = 0
    while mi < 256 {
        let byte_idx = mi / 8
        let bit_idx = mi % 8
        let byte_val = t_index_get(m, byte_idx)
        let bit_val = bit_and(shr(byte_val, bit_idx), 1)
        let decompressed = decompress(bit_val, 1)
        msg_poly = t_index_set(msg_poly, mi, decompressed + 0.0)
        mi = mi + 1
    }
    let v = poly_add(v_no_msg, msg_poly)

    // Compress and encode ciphertext
    // c1 = encode(compress(u, du)) for each component
    let mut c1_0 = t_zeros(256)
    let mut c1_1 = t_zeros(256)
    let mut c1_2 = t_zeros(256)
    let mut ci = 0
    while ci < 256 {
        c1_0 = t_index_set(c1_0, ci, compress(fq_normalize(t_index_get(u0, ci)), kyber_du()) + 0.0)
        c1_1 = t_index_set(c1_1, ci, compress(fq_normalize(t_index_get(u1, ci)), kyber_du()) + 0.0)
        c1_2 = t_index_set(c1_2, ci, compress(fq_normalize(t_index_get(u2, ci)), kyber_du()) + 0.0)
        ci = ci + 1
    }

    // c2 = encode(compress(v, dv))
    let mut c2_poly = t_zeros(256)
    let mut vi = 0
    while vi < 256 {
        c2_poly = t_index_set(c2_poly, vi, compress(fq_normalize(t_index_get(v, vi)), kyber_dv()) + 0.0)
        vi = vi + 1
    }

    // Pack: c1 (3 * 320 bytes for du=10) || c2 (128 bytes for dv=4)
    let c1_bytes = t_cat(t_cat(poly_encode(c1_0, kyber_du()), poly_encode(c1_1, kyber_du())), poly_encode(c1_2, kyber_du()))
    let c2_bytes = poly_encode(c2_poly, kyber_dv())
    return t_cat(c1_bytes, c2_bytes)
}

// ── K-PKE.Decrypt (FIPS 203 Algorithm 6) ──────────────────────────
// Input: dk (secret key bytes), ct (ciphertext bytes)
// Output: 32-byte message

fn kpke_decrypt(dk: Tensor, ct: Tensor) -> Tensor {
    // Decode secret key
    let s0 = poly_decode12(tensor_slice(dk, 0, 384))
    let s1 = poly_decode12(tensor_slice(dk, 384, 768))
    let s2 = poly_decode12(tensor_slice(dk, 768, 1152))

    // Decode ciphertext
    // c1: 3 * 320 bytes (du=10), c2: 128 bytes (dv=4)
    let c1_0 = poly_decode(tensor_slice(ct, 0, 320), kyber_du())
    let c1_1 = poly_decode(tensor_slice(ct, 320, 640), kyber_du())
    let c1_2 = poly_decode(tensor_slice(ct, 640, 960), kyber_du())
    let c2 = poly_decode(tensor_slice(ct, 960, 1088), kyber_dv())

    // Decompress u and v
    let mut u0 = t_zeros(256)
    let mut u1 = t_zeros(256)
    let mut u2 = t_zeros(256)
    let mut v = t_zeros(256)
    let mut i = 0
    while i < 256 {
        u0 = t_index_set(u0, i, decompress(t_index_get(c1_0, i), kyber_du()) + 0.0)
        u1 = t_index_set(u1, i, decompress(t_index_get(c1_1, i), kyber_du()) + 0.0)
        u2 = t_index_set(u2, i, decompress(t_index_get(c1_2, i), kyber_du()) + 0.0)
        v = t_index_set(v, i, decompress(t_index_get(c2, i), kyber_dv()) + 0.0)
        i = i + 1
    }

    // w = v - NTT^-1(s^T * NTT(u))
    let u0_ntt = ntt_forward(u0)
    let u1_ntt = ntt_forward(u1)
    let u2_ntt = ntt_forward(u2)
    let su_ntt = poly_add(ntt_basemul(s0, u0_ntt), poly_add(ntt_basemul(s1, u1_ntt), ntt_basemul(s2, u2_ntt)))
    let su = poly_reduce(ntt_inverse(su_ntt))
    let w = poly_sub(v, su)

    // Recover message: compress each coefficient to 1 bit
    let mut msg = t_zeros(32)
    let mut mi = 0
    while mi < 256 {
        let coeff = fq_normalize(t_index_get(w, mi))
        let bit = compress(coeff, 1)
        let byte_idx = mi / 8
        let bit_idx = mi % 8
        let old_byte = t_index_get(msg, byte_idx)
        msg = t_index_set(msg, byte_idx, old_byte + shl(bit, bit_idx) + 0.0)
        mi = mi + 1
    }
    return msg
}
