use crypto::kpke
use crypto::fieldq

fn main() {
    println("=== K-PKE Full Pipeline ===")
    let seed = secure_random_bytes(32)
    let keys = kpke_keygen(seed)
    let ek = array_get(keys, 0)
    let dk = array_get(keys, 1)

    // Encrypt all-zero message
    let msg = t_zeros(32)
    let r = secure_random_bytes(32)
    let ct = kpke_encrypt(ek, msg, r)

    // Manually run decrypt steps to trace where it breaks
    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))

    let c1_0_enc = poly_decode(tensor_slice(ct, 0, 320), 10)
    let c2_enc = poly_decode(tensor_slice(ct, 960, 1088), 4)

    println("Encoded c1_0[0..4]: " + float_to_str(t_index_get(c1_0_enc, 0)) + " " + float_to_str(t_index_get(c1_0_enc, 1)) + " " + float_to_str(t_index_get(c1_0_enc, 2)) + " " + float_to_str(t_index_get(c1_0_enc, 3)))
    println("Encoded c2[0..4]: " + float_to_str(t_index_get(c2_enc, 0)) + " " + float_to_str(t_index_get(c2_enc, 1)) + " " + float_to_str(t_index_get(c2_enc, 2)) + " " + float_to_str(t_index_get(c2_enc, 3)))

    // Decompress
    let u0_val = decompress(t_index_get(c1_0_enc, 0), 10)
    let v_val = decompress(t_index_get(c2_enc, 0), 4)
    println("Decompressed u0[0]: " + int_to_str(u0_val) + " (should be in [0, q))")
    println("Decompressed v[0]: " + int_to_str(v_val) + " (should be in [0, q))")

    // Full decrypt
    let recovered = kpke_decrypt(dk, ct)
    println("")
    println("Decrypted[0..8]: " + float_to_str(t_index_get(recovered, 0)) + " " + float_to_str(t_index_get(recovered, 1)) + " " + float_to_str(t_index_get(recovered, 2)) + " " + float_to_str(t_index_get(recovered, 3)) + " " + float_to_str(t_index_get(recovered, 4)) + " " + float_to_str(t_index_get(recovered, 5)) + " " + float_to_str(t_index_get(recovered, 6)) + " " + float_to_str(t_index_get(recovered, 7)))
    println("(should be all 0 for zero message)")
}
