use std::env;
use std::time::Instant;

use rand::RngCore;

use spiral_spir::client::Client;
use spiral_spir::init;
use spiral_spir::server::Server;
use spiral_spir::DbEntry;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 || args.len() > 5 {
        println!(
            "Usage: {} r [num_threads [num_preproc [num_pirs]]]\nr = log_2(num_records)",
            args[0]
        );
        return;
    }
    let r: usize = args[1].parse().unwrap();
    let mut num_threads = 1usize;
    let mut num_preproc = 1usize;
    let num_pirs: usize;
    if args.len() > 2 {
        num_threads = args[2].parse().unwrap();
    }
    if args.len() > 3 {
        num_preproc = args[3].parse().unwrap();
    }
    if args.len() > 4 {
        num_pirs = args[4].parse().unwrap();
    } else {
        num_pirs = num_preproc;
    }
    let num_records = 1 << r;
    let num_records_mask = num_records - 1;

    println!("===== ONE-TIME SETUP =====\n");

    let otsetup_start = Instant::now();

    init(num_threads);
    let (client, pub_params) = Client::new(r);
    let pub_params_len = pub_params.len();
    let server = Server::new(r, pub_params);

    let otsetup_us = otsetup_start.elapsed().as_micros();
    println!("One-time setup: {} µs", otsetup_us);
    println!("pub_params len = {}", pub_params_len);

    println!("\n===== PREPROCESSING =====\n");

    println!("num_preproc = {}", num_preproc);

    let preproc_client_start = Instant::now();
    let preproc_msg = client.preproc(num_preproc);
    let preproc_client_us = preproc_client_start.elapsed().as_micros();

    println!("Preprocessing client: {} µs", preproc_client_us);
    println!("preproc_msg len = {}", preproc_msg.len());

    let preproc_server_start = Instant::now();
    let preproc_resp = server.preproc_process(&preproc_msg);
    let preproc_server_us = preproc_server_start.elapsed().as_micros();

    println!("Preprocessing server: {} µs", preproc_server_us);
    println!("preproc_resp len = {}", preproc_resp.len());

    let preproc_finish_start = Instant::now();
    client.preproc_finish(&preproc_resp);
    let preproc_finish_us = preproc_finish_start.elapsed().as_micros();

    println!("Preprocessing client finish: {} µs", preproc_finish_us);

    // Create a database with recognizable contents
    let db: Vec<DbEntry> = ((0 as DbEntry)..(num_records as DbEntry))
        .map(|x| 10000001 * x)
        .collect();
    let dbptr = db.as_ptr();

    let mut rng = rand::thread_rng();
    for i in 1..num_pirs + 1 {
        println!("\n===== SPIR QUERY {} =====\n", i);

        let idx = (rng.next_u64() as usize) & num_records_mask;
        let query_client_start = Instant::now();
        let query_msg = client.query(idx);
        let query_client_us = query_client_start.elapsed().as_micros();

        println!("Query client: {} µs", query_client_us);
        println!("query_msg len = {}", query_msg.len());

        let query_server_start = Instant::now();
        let query_resp = server.query_process(&query_msg, dbptr, 100, 20);
        let query_server_us = query_server_start.elapsed().as_micros();

        println!("Query server: {} µs", query_server_us);
        println!("query_resp len = {}", query_resp.len());

        let query_finish_start = Instant::now();
        let entry = client.query_finish(&query_resp);
        let query_finish_us = query_finish_start.elapsed().as_micros();

        println!("Query client finish: {} µs", query_finish_us);
        println!("idx = {}; entry = {}", idx, entry);
    }
}