use antimatter::capsule::common::{CellReader, Column, KEY_SIZE, NONCE_SIZE};
use antimatter::capsule::streaming_aead::{
    streaming_decrypt_aes_256_gcm, streaming_encrypt_aes_256_gcm, DecryptingAEAD,
    EncryptingAEADReader,
};
use antimatter::session::api_helper::Session;
use antimatter::session::session::{EncapsulateConfig, SessionConf};
use antimatter::session::RUNTIME;
use antimatter_api::apis::configuration::Configuration;
use antimatter_api::apis::{authentication_api, general_api};
use antimatter_api::models::{DomainAuthenticate, NewDomain};
use criterion::{criterion_group, criterion_main, Criterion};
use p256::pkcs8::der::Encode;
use rand::Rng;
use std::collections::HashMap;
use std::env;
use std::io::{Cursor, Read};
use std::sync::{Arc, Mutex};
use std::time::Duration;

const ANTIMATTER_TEST_ADDRESS: &str = "test@antimatter.io";
const API_TARGET_VERSION: &str = "v2";

fn generate_capsule(s: &mut Session) -> Vec<u8> {
    let cell_readers = vec![vec![CellReader::new(
        vec![],
        std::io::Cursor::new("row0, col0: example data element".as_bytes().to_vec()),
    )
    .expect("failed to create reader")]];
    let columns = vec![Column {
        name: "col0".to_string(),
        tags: vec![],
        skip_classification: true,
    }];
    let cfg = EncapsulateConfig {
        write_context_name: "default".to_string(),
        extra: "some extra data".to_string(),
        subdomain: None,
        subdomain_from: None,
        create_subdomains: None,
        async_seal: false,
    };

    // encapsulate the data
    let (capsule_data, _meta) = s
        .encapsulate_to_bytes(columns, cell_readers, vec![], cfg)
        .expect("failed to encapsulate");
    return capsule_data;
}

fn open_capsule(s: &mut Session, data: Vec<u8>) {
    // open the sealed capsule
    let writer = Cursor::new(data);
    s.open("default", HashMap::new(), HashMap::new(), writer)
        .expect("failed to open capsule");
}

fn create_domain(email: &str) -> Result<(String, String, Configuration), String> {
    let mut config = Configuration {
        base_path: format!("{}/{}", antimatter_api_url(), API_TARGET_VERSION),
        user_agent: None,
        client: antimatter::session::http_client::HTTPClient::new()
            .expect("failed to create HTTP client")
            .client(),
        basic_auth: None,
        oauth_access_token: None,
        bearer_access_token: Some("TODO".to_string()),
        api_key: None,
    };

    let response = RUNTIME
        .block_on(general_api::domain_add_new(
            &config,
            NewDomain {
                admin_email: email.to_string(),
                google_jwt: None,
                display_name: None,
            },
        ))
        .map_err(|e| format!("Failed to create domain: {}", e))?;

    let auth = RUNTIME
        .block_on(authentication_api::domain_authenticate(
            &config,
            response.id.clone().as_str(),
            DomainAuthenticate {
                token: response.api_key.clone(),
            },
            None,
            None,
            None,
        ))
        .expect("failed to auth parent");
    config.bearer_access_token = Some(auth.token);
    Ok((response.id, response.api_key, config))
}
fn antimatter_api_url() -> String {
    env::var("ANTIMATTER_TEST_API_URL")
        .unwrap_or_else(|_| "https://api.dev.antimatter.io".to_string())
}

fn benchmark_streaming_encrypt_aes_256_gcm(c: &mut Criterion) {
    let plaintext_sizes = vec![1024 * 1024 * 1024];
    let mut rng = rand::thread_rng();
    for size in plaintext_sizes {
        c.benchmark_group(format!("streaming encrypt:{}", size))
            .sample_size(10)
            .measurement_time(std::time::Duration::from_secs(20))
            .throughput(criterion::Throughput::Bytes(size))
            .bench_function("best possible", |b| {
                // basic encrypt using the ring library
                let mut data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();

                b.iter(|| {
                    let key =
                        ring::aead::UnboundKey::new(&ring::aead::AES_256_GCM, &[0u8; KEY_SIZE])
                            .expect("failed to create key");
                    let cipher = ring::aead::LessSafeKey::new(key);

                    let _ = cipher
                        .seal_in_place_separate_tag(
                            ring::aead::Nonce::assume_unique_for_key([0u8; NONCE_SIZE]),
                            ring::aead::Aad::empty(),
                            data.as_mut_slice(),
                        )
                        .expect("encryption failed");
                });
            })
            .bench_function("EncryptingAEAD", |b| {
                let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();

                b.iter(|| {
                    let mut encrypted_data = Vec::<u8>::new();
                    let mut nonce_block = [0u8; NONCE_SIZE];
                    let mut input = std::io::Cursor::new(&data);
                    let mut encrypt =
                        EncryptingAEADReader::new(nonce_block, &[0u8; KEY_SIZE], &mut input)
                            .expect("failed to create EncryptingAEAD");
                    encrypt
                        .read_to_end(&mut encrypted_data)
                        .expect("encrypt failed");
                });
            })
            .bench_function("streaming_encrypt_aes_256_gcm", |b| {
                // streaming encrypt using antimatter::capsule::streaming_aead
                let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();

                b.iter(|| {
                    let mut encrypted_data = Vec::<u8>::new();
                    let _ = streaming_encrypt_aes_256_gcm(
                        &[0u8; KEY_SIZE],
                        &mut [0u8; NONCE_SIZE],
                        std::io::Cursor::new(&data),
                        &mut encrypted_data,
                    )
                    .expect("streaming_encrypt_aes_256_gcm failed");
                });
            });
    }
}

fn benchmark_streaming_decrypt_aes_256_gcm(c: &mut Criterion) {
    let plaintext_sizes = vec![1024 * 1024 * 1024];
    let mut rng = rand::thread_rng();
    for size in plaintext_sizes {
        c.benchmark_group(format!("streaming decrypt:{}", size))
            .sample_size(10)
            .measurement_time(std::time::Duration::from_secs(20))
            .throughput(criterion::Throughput::Bytes(size))
            .bench_function("best possible", |b| {
                // basic decrypt using the aes_gcm library
                let mut data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();
                let key = ring::aead::UnboundKey::new(&ring::aead::AES_256_GCM, &[0u8; KEY_SIZE])
                    .expect("failed to create key");
                let cipher = ring::aead::LessSafeKey::new(key);

                let tag = cipher
                    .seal_in_place_separate_tag(
                        ring::aead::Nonce::assume_unique_for_key([0u8; NONCE_SIZE]),
                        ring::aead::Aad::empty(),
                        data.as_mut_slice(),
                    )
                    .expect("encryption failed");

                data.extend_from_slice(tag.as_ref());

                b.iter(|| {
                    cipher
                        .open_in_place(
                            ring::aead::Nonce::assume_unique_for_key([0u8; NONCE_SIZE]),
                            ring::aead::Aad::empty(),
                            data.clone().as_mut_slice(),
                        )
                        .expect("decryption failed");
                });
            })
            .bench_function("DecryptingAEAD", |b| {
                let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();
                let mut encrypted_data = Vec::<u8>::new();
                let nonce_block = [0u8; NONCE_SIZE];
                let mut input = std::io::Cursor::new(&data);
                let mut encrypt =
                    EncryptingAEADReader::new(nonce_block, &[0u8; KEY_SIZE], &mut input)
                        .expect("failed to create EncryptingAEAD");
                encrypt
                    .read_to_end(&mut encrypted_data)
                    .expect("encrypt failed");

                b.iter(|| {
                    let input = Arc::new(Mutex::new(Cursor::new(&encrypted_data)));
                    let mut decrypt = DecryptingAEAD::new(&[0u8; KEY_SIZE], input)
                        .expect("failed to creating DecryptingAEAD");

                    let mut decrypted_data: Vec<u8> = Vec::new();
                    decrypt
                        .read_to_end(&mut decrypted_data)
                        .expect("decrypt failed");
                });
            })
            .bench_function("streaming_decrypt_aes_256_gcm", |b| {
                // streaming decrypt using antimatter::capsule::streaming_aead
                let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();
                let mut encrypted_data = Vec::<u8>::new();
                let _ = streaming_encrypt_aes_256_gcm(
                    &[0u8; KEY_SIZE],
                    &mut [0u8; NONCE_SIZE],
                    std::io::Cursor::new(&data),
                    &mut encrypted_data,
                )
                .expect("streaming_encrypt_aes_256_gcm failed");

                b.iter(|| {
                    let mut decrypted_data = Vec::<u8>::new();
                    let _ = streaming_decrypt_aes_256_gcm(
                        &[0u8; KEY_SIZE],
                        std::io::Cursor::new(&encrypted_data),
                        &mut decrypted_data,
                    )
                    .expect("streaming_encrypt_aes_256_gcm failed");
                });
            });
    }
}

fn benchmark_sync_streaming_encrypt_decrypt(c: &mut Criterion) {
    use serde_tuple::{Deserialize_tuple, Serialize_tuple};

    #[derive(Serialize_tuple, Deserialize_tuple)]
    struct V2CapsuleBody {
        disable_read_logging: bool,
    }

    let plaintext_sizes = vec![1024 * 1024 * 1024];
    let mut rng = rand::thread_rng();
    for size in plaintext_sizes {
        c.benchmark_group(format!("streaming_encode_encrypt_decrypt_decode:{}", size))
            .sample_size(10)
            .measurement_time(std::time::Duration::from_secs(40))
            .throughput(criterion::Throughput::Bytes(2 * size))
            .bench_function("actual performance", |b| {
                let data: Vec<u8> = (0..size).map(|_| rng.gen()).collect();

                b.iter(|| {
                    let mut encrypted_data = Vec::<u8>::new();
                    let mut decrypted_data = Vec::<u8>::new();
                    let mut buffer_data = Vec::<u8>::new();

                    let _ = streaming_encrypt_aes_256_gcm(
                        &[0u8; KEY_SIZE],
                        &mut [0u8; NONCE_SIZE],
                        Cursor::new(&data),
                        &mut encrypted_data,
                    )
                    .expect("encryption failed");

                    let _ = streaming_decrypt_aes_256_gcm(
                        &[0u8; KEY_SIZE],
                        Cursor::new(&encrypted_data),
                        &mut buffer_data,
                    )
                    .expect("decryption failed");

                    let mut decrypted_reader = Cursor::new(buffer_data);
                    let _ = decrypted_reader
                        .read_to_end(&mut decrypted_data)
                        .expect("failed reading decrypted cell");
                })
            });
    }
}

fn benchmark_open_capsule(c: &mut Criterion) {
    env::set_var("ANTIMATTER_API_URL", antimatter_api_url());

    let (domain_id, api_key, _) =
        create_domain(ANTIMATTER_TEST_ADDRESS).expect("failed to create domain");

    let config_a = SessionConf {
        domain_id: domain_id.clone(),
        bearer_access_token: None,
        api_key: Some(api_key.clone()),
        read_cache_size: 0,
        engine_cache_size: 0,
        write_cache_size: 0,
        subdomain_cache_size: 0,
        buffered_seal: 0,
        buffered_seal_enabled: false,
        use_direct_address: false,
        current_base_path: None,
        act_for_domain: None,
    };

    let config_b = SessionConf {
        domain_id: domain_id.clone(),
        bearer_access_token: None,
        api_key: Some(api_key.clone()),
        read_cache_size: 0,
        engine_cache_size: 0,
        write_cache_size: 0,
        subdomain_cache_size: 0,
        buffered_seal: 0,
        buffered_seal_enabled: false,
        use_direct_address: true,
        current_base_path: None,
        act_for_domain: None,
    };

    let (mut session_a, _) = Session::from_config(config_a).expect("failed to create session");
    let (mut session_b, _) = Session::from_config(config_b).expect("failed to create session");
    let mut group = c.benchmark_group("direct vs indirect addressing");

    group
        .sample_size(10)
        .warm_up_time(Duration::from_secs(3))
        .measurement_time(Duration::from_secs(70));
    let capsule = generate_capsule(&mut session_a);
    group.bench_function("open capsules via global", |b| {
        b.iter(|| open_capsule(&mut session_a, capsule.clone()))
    });
    group.bench_function("open capsules vai cell  ", |b| {
        b.iter(|| open_capsule(&mut session_b, capsule.clone()))
    });
    group.finish();
}

fn create_large_capsule(session: &mut Session) -> Vec<u8> {
    // Construct a really large table of small cells to test with:
    let row_count = 10;
    let col_count = 10;
    let cell_data = " John Smith ".to_string();
    let mut cell_readers = Vec::new();
    let mut columns = Vec::new();
    for _ in 0..row_count {
        let mut readers_row = Vec::new();
        for _ in 0..col_count {
            readers_row.push(
                CellReader::new(vec![], Cursor::new(cell_data.clone()))
                    .expect("failed to make reader"),
            );
        }
        cell_readers.push(readers_row);
    }
    // Make the columns:
    for _ in 0..col_count {
        columns.push(Column {
            name: "xx".to_string(),
            tags: vec![],
            skip_classification: false,
        })
    }

    let cfg = EncapsulateConfig {
        write_context_name: "sensitive".to_string(),
        extra: "some extra data".to_string(),
        subdomain: None,
        subdomain_from: None,
        create_subdomains: None,
        async_seal: false,
    };
    let (data, _) = session
        .encapsulate_to_bytes(columns, cell_readers, vec![], cfg)
        .expect("failed to encapsulate");
    data
}

fn benchmark_creating_large_capsule(c: &mut Criterion) {
    env::set_var("ANTIMATTER_API_URL", antimatter_api_url());

    let (domain_id, api_key, _) =
        create_domain(ANTIMATTER_TEST_ADDRESS).expect("failed to create domain");

    let mut session =
        Session::new(domain_id.clone(), api_key.clone()).expect("failed to create session");
    let mut group = c.benchmark_group("classify many small cells");
    group
        .sample_size(10)
        .warm_up_time(Duration::from_secs(3))
        .measurement_time(Duration::from_secs(180));

    group.bench_function("create large capsules", |b| {
        b.iter(|| create_large_capsule(&mut session))
    });
    group.finish();
}

criterion_group!(
    benches,
    // benchmark_streaming_encrypt_aes_256_gcm,
    // benchmark_streaming_decrypt_aes_256_gcm,
    //    benchmark_sync_streaming_encrypt_decrypt,
    benchmark_open_capsule,
    benchmark_creating_large_capsule,
);
criterion_main!(benches);