use aes_gcm_siv::{
    aead::{Aead, OsRng}, AeadCore, Aes256GcmSiv, KeyInit
};
use log::{info, warn};
use tokio::sync::RwLock;

use super::DbPool;

#[derive(PartialEq)]
enum KeyEnum {
    /// Final key
    MainKey(Vec<u8>),
    Uninitialized,
    // Portion(Vec<String>),
}

static KEY_MAPS: RwLock<KeyEnum> = RwLock::const_new(KeyEnum::Uninitialized);

pub async fn get_or_init_root_key(pool: &DbPool) {
    let mut value = KEY_MAPS.write().await;

    match *value {
        KeyEnum::MainKey(_) => panic!("Root key already initialized!"),
        KeyEnum::Uninitialized => get_root_key_db(pool, &mut value).await,
    }
    // Write lock on KEY_MAPS is hold throughout the process and set at last
    // to avoid secrets being sealed with a new but discarded key
}

async fn get_root_key_db(pool: &DbPool, value: &mut KeyEnum) {
    let rk = sqlx::query!("SELECT encrypted_key, type FROM root_key ORDER BY version LIMIT 1")
        .fetch_optional(pool)
        .await
        .expect("Failed to optionally read root key from the database");

    let v = match rk {
        Some(v) => v,
        None => {
            warn!("No root key was found in the database!");
            init_new_root_key(pool, value).await;
            return;
        },
    };
    info!("Root key of type {} found in the database", v.r#type);
    match &*v.r#type {
        "dev_only" => {
            warn!("Root key is of type {}. This is INSECURE and must only be used for development purposes!", v.r#type);
            *value = KeyEnum::MainKey(v.encrypted_key);
            return;
        },
        _ => panic!("Unknown root key type in database"),
    }
}

async fn init_new_root_key(pool: &DbPool, value: &mut KeyEnum) {
    warn!("Initializing new root key!");
    let key = Aes256GcmSiv::generate_key(&mut OsRng);

    let key = key.as_slice().to_owned();

    let _ = sqlx::query!(
        "
        INSERT INTO root_key (encrypted_key, type, version)
        VALUES ($1, 'dev_only', 1)
        ",
        key
    )
    .execute(pool)
    .await
    .expect("Failed to write new root key to the database");

    *value = KeyEnum::MainKey(key);
    info!("Initialized new root key!");
}

pub async fn seal(data: &String) -> ([u8; 12], Vec<u8>) {
    let cipher = match &*KEY_MAPS.read().await {
        KeyEnum::MainKey(key) => Aes256GcmSiv::new_from_slice(key),
        KeyEnum::Uninitialized => panic!("Cannot seal secret since the vault is not unsealed"),
    }
    .expect("Failed to create new AesGcmSiv cipher from variable size key");

    let nonce: aes_gcm_siv::aead::generic_array::GenericArray<u8, <Aes256GcmSiv as aes_gcm_siv::AeadCore>::NonceSize> =
        Aes256GcmSiv::generate_nonce(&mut OsRng); // 96-bits; unique per message
    let enc = cipher.encrypt(&nonce, data.as_bytes()).unwrap();
    debug_assert!(nonce.len() == 12);
    let nonce = nonce.as_slice().try_into().expect("Nonce should be exactly 12 bytes");
    (nonce, enc)
}

pub async fn unseal(nonce: Vec<u8>, data: impl AsRef<[u8]>) -> String {
    assert!(nonce.len() == 12);
    let cipher = match &*KEY_MAPS.read().await {
        KeyEnum::MainKey(key) => Aes256GcmSiv::new_from_slice(key),
        KeyEnum::Uninitialized => panic!("Cannot seal secret since the vault is not unsealed"),
    }
    .expect("Failed to create new AesGcmSiv cipher from variable size key");

    let nonce = aes_gcm_siv::aead::generic_array::GenericArray::from_slice(&nonce);
    let enc = cipher.decrypt(nonce, data.as_ref()).unwrap();
    String::from_utf8(enc).expect("Failed to parse as utf8")
}