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

fn main() {
    println("=== K-PKE Internal Trace ===")

    // Generate keys with known seed
    let d = secure_random_bytes(32)
    let keys = kpke_keygen(d)
    let ek = array_get(keys, 0)
    let dk = array_get(keys, 1)

    // Extract rho from ek
    let rho = tensor_slice(ek, 1152, 1184)

    // Decode t and s from keys
    let t0 = poly_decode12(tensor_slice(ek, 0, 384))
    let s0 = poly_decode12(tensor_slice(dk, 0, 384))

    // Resample A[0][0] (same as keygen did)
    let a00 = sample_ntt(rho, 0, 0)

    // Check: does basemul(a00, s0) + e0 ≈ t0?
    // We can't recover e0, but we can check basemul(a00, s0) is close to t0
    let as0 = ntt_basemul(a00, s0)
    println("Check A*s vs t (first 4 NTT coefficients):")
    let mut i = 0
    while i < 4 {
        let as_val = t_index_get(as0, i)
        let t_val = t_index_get(t0, i)
        let diff = t_val - as_val
        println("  [" + int_to_str(i) + "] A*s=" + float_to_str(as_val) + " t=" + float_to_str(t_val) + " diff=" + float_to_str(diff))
        i = i + 1
    }

    // Now test encrypt path: encrypt zero message
    let msg = t_zeros(32)
    let r_seed = secure_random_bytes(32)
    let r0 = ntt_forward(sample_noise_cbd2(r_seed, 0))

    // u0 ≈ A^T[0][0] * r0 + e1_0 (simplified, ignoring other rows)
    let at00 = sample_ntt(rho, 0, 0)
    let at_r = ntt_basemul(at00, r0)
    let u0_approx = ntt_inverse(at_r)

    // v ≈ t0^T * r0 + e2 (simplified)
    let t_r = ntt_basemul(t0, r0)
    let v_approx = ntt_inverse(t_r)

    // Decrypt check: s0^T * NTT(u0_approx) ≈ v_approx?
    let su = ntt_basemul(s0, ntt_forward(u0_approx))
    let su_plain = ntt_inverse(su)

    println("")
    println("Decrypt check (first 4 coefficients):")
    let mut j = 0
    while j < 4 {
        let v_val = fq_normalize(t_index_get(v_approx, j))
        let su_val = fq_normalize(t_index_get(su_plain, j))
        let w = fq_normalize(v_val - su_val + fq_q())
        println("  [" + int_to_str(j) + "] v=" + int_to_str(v_val) + " s*u=" + int_to_str(su_val) + " w=v-su=" + int_to_str(w))
        j = j + 1
    }
    println("  (w should be near 0 or q for zero message)")
}
