//! DudeCT-based timing verification for constant-time operations.
//!
//! A |t-value| > 5 indicates a timing leak with high confidence.
//! Functions should show |t-value| < 5 after sufficient samples.

use dudect_bencher::{BenchRng, Class, CtRunner, ctbench_main, rand::Rng};
use opaque_lattice::oprf::fast_oprf::{
    PublicParams, Q, RING_N, ReconciliationHelper, RingElement, ServerKey, client_blind,
    client_finalize, server_evaluate,
};

fn coin_flip(rng: &mut BenchRng) -> bool {
    rng.gen_range(0u8..2) == 0
}

const NUM_SAMPLES: usize = 10_000;

fn timing_ring_mul(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<(RingElement, RingElement)> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let mut coeffs_a = [0i32; RING_N];
            let mut coeffs_b = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs_a[i] = 0;
                coeffs_b[i] = 0;
            }
            inputs.push((
                RingElement { coeffs: coeffs_a },
                RingElement { coeffs: coeffs_b },
            ));
            classes.push(Class::Left);
        } else {
            let mut coeffs_a = [0i32; RING_N];
            let mut coeffs_b = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs_a[i] = rng.gen_range(0..Q);
                coeffs_b[i] = rng.gen_range(0..Q);
            }
            inputs.push((
                RingElement { coeffs: coeffs_a },
                RingElement { coeffs: coeffs_b },
            ));
            classes.push(Class::Right);
        }
    }

    for (class, (a, b)) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || a.mul(&b));
    }
}

fn timing_linf_norm(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<RingElement> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let coeffs = [0i32; RING_N];
            inputs.push(RingElement { coeffs });
            classes.push(Class::Left);
        } else {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = rng.gen_range(0..Q);
            }
            inputs.push(RingElement { coeffs });
            classes.push(Class::Right);
        }
    }

    for (class, elem) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || elem.linf_norm());
    }
}

fn timing_round_to_binary(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<RingElement> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = Q / 4;
            }
            inputs.push(RingElement { coeffs });
            classes.push(Class::Left);
        } else {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = (3 * Q) / 4;
            }
            inputs.push(RingElement { coeffs });
            classes.push(Class::Right);
        }
    }

    for (class, elem) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || elem.round_to_binary());
    }
}

fn timing_ring_eq(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<(RingElement, RingElement)> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let coeffs = [0i32; RING_N];
            inputs.push((RingElement { coeffs }, RingElement { coeffs }));
            classes.push(Class::Left);
        } else {
            let mut coeffs_a = [0i32; RING_N];
            let mut coeffs_b = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs_a[i] = rng.gen_range(0..Q);
                coeffs_b[i] = rng.gen_range(0..Q);
            }
            inputs.push((
                RingElement { coeffs: coeffs_a },
                RingElement { coeffs: coeffs_b },
            ));
            classes.push(Class::Right);
        }
    }

    for (class, (a, b)) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || a.eq(&b));
    }
}

fn timing_extract_bits(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<(ReconciliationHelper, RingElement)> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = Q / 8;
            }
            let elem = RingElement { coeffs };
            let helper = ReconciliationHelper::from_ring(&elem);
            inputs.push((helper, elem));
            classes.push(Class::Left);
        } else {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = (7 * Q) / 8;
            }
            let elem = RingElement { coeffs };
            let helper = ReconciliationHelper::from_ring(&elem);
            inputs.push((helper, elem));
            classes.push(Class::Right);
        }
    }

    for (class, (helper, elem)) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || helper.extract_bits(&elem));
    }
}

fn timing_server_bits(runner: &mut CtRunner, rng: &mut BenchRng) {
    let mut inputs: Vec<RingElement> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = Q / 4;
            }
            inputs.push(RingElement { coeffs });
            classes.push(Class::Left);
        } else {
            let mut coeffs = [0i32; RING_N];
            for i in 0..RING_N {
                coeffs[i] = (3 * Q) / 4;
            }
            inputs.push(RingElement { coeffs });
            classes.push(Class::Right);
        }
    }

    for (class, elem) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || ReconciliationHelper::server_bits(&elem));
    }
}

fn timing_client_blind(runner: &mut CtRunner, rng: &mut BenchRng) {
    let pp = PublicParams::generate(b"timing-test-seed");
    let mut inputs: Vec<Vec<u8>> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            inputs.push(vec![0u8; 16]);
            classes.push(Class::Left);
        } else {
            let mut pwd = vec![0u8; 16];
            rng.fill(&mut pwd[..]);
            inputs.push(pwd);
            classes.push(Class::Right);
        }
    }

    for (class, password) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || client_blind(&pp, &password));
    }
}

fn timing_server_evaluate(runner: &mut CtRunner, rng: &mut BenchRng) {
    let pp = PublicParams::generate(b"timing-test-seed");
    let key = ServerKey::generate(&pp, b"timing-test-key");
    let mut inputs = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            let (_, blinded) = client_blind(&pp, &[0u8; 16]);
            inputs.push(blinded);
            classes.push(Class::Left);
        } else {
            let mut pwd = [0u8; 16];
            rng.fill(&mut pwd[..]);
            let (_, blinded) = client_blind(&pp, &pwd);
            inputs.push(blinded);
            classes.push(Class::Right);
        }
    }

    for (class, blinded) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || server_evaluate(&key, &blinded));
    }
}

fn timing_full_protocol(runner: &mut CtRunner, rng: &mut BenchRng) {
    let pp = PublicParams::generate(b"timing-test-seed");
    let key = ServerKey::generate(&pp, b"timing-test-key");
    let mut inputs: Vec<Vec<u8>> = Vec::new();
    let mut classes = Vec::new();

    for _ in 0..NUM_SAMPLES {
        if coin_flip(rng) {
            inputs.push(vec![0u8; 16]);
            classes.push(Class::Left);
        } else {
            let mut pwd = vec![0u8; 16];
            rng.fill(&mut pwd[..]);
            inputs.push(pwd);
            classes.push(Class::Right);
        }
    }

    for (class, password) in classes.into_iter().zip(inputs.into_iter()) {
        runner.run_one(class, || {
            let (state, blinded) = client_blind(&pp, &password);
            let response = server_evaluate(&key, &blinded);
            client_finalize(&state, &key.b, &response)
        });
    }
}

ctbench_main!(
    timing_ring_mul,
    timing_linf_norm,
    timing_round_to_binary,
    timing_ring_eq,
    timing_extract_bits,
    timing_server_bits,
    timing_client_blind,
    timing_server_evaluate,
    timing_full_protocol
);