use core::{
    cell::{Cell, RefCell},
    ffi::{c_char, c_int, c_size_t, c_uchar, c_ulonglong},
};

use embassy_sync::blocking_mutex::{self, raw::CriticalSectionRawMutex};
use hmac::digest::{FixedOutput, KeyInit, Update};
use password_hash::Key;
use sha2::{
    Digest,
    digest::{consts::U32, generic_array::GenericArray},
};
use spectre_api_sys::{SpectreKeyPurpose, spectre_purpose_scope};

use crate::PSRAM_ALLOCATOR;

// Just a way to mark non-null pointers.
// Rust's `NonNull` is for mutable pointers only.
pub type NonNullPtr<P> = P;

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_crypto_hash_sha256(
    data_out: NonNullPtr<*mut c_uchar>,
    data_in: *const c_uchar,
    data_in_len: c_ulonglong,
) -> c_int {
    unsafe {
        let data_out = &mut *(data_out as *mut GenericArray<u8, U32>);
        let data_in = core::slice::from_raw_parts(data_in, data_in_len as usize);
        // TODO: Use SHA peripheral for acceleration
        let mut digest = sha2::Sha256::new();
        Digest::update(&mut digest, data_in);
        digest.finalize_into_reset(data_out);
    }

    0
}

/// This is the encrypted user key currently being used in the key derivation function of spectre.
/// It decrypts using the user's password into the key that would be derived with the original password hashing function.
pub static ACTIVE_ENCRYPTED_USER_KEY: blocking_mutex::Mutex<CriticalSectionRawMutex, Cell<Key>> =
    blocking_mutex::Mutex::new(Cell::new([0; _]));

#[unsafe(no_mangle)]
#[must_use]
unsafe extern "C" fn __spre_crypto_pwhash_scryptsalsa208sha256_ll(
    password: NonNullPtr<*const u8>,
    password_len: c_size_t,
    salt: NonNullPtr<*const u8>,
    salt_len: c_size_t,
    n: u64,
    r: u32,
    p: u32,
    output: NonNullPtr<*mut u8>,
    output_len: c_size_t,
) -> c_int {
    assert_eq!(output_len, 64);

    let encryption_key = unsafe {
        let password: &[u8] = core::slice::from_raw_parts(password, password_len);
        let salt: &[u8] = core::slice::from_raw_parts(salt, salt_len);
        let purpose = spectre_purpose_scope(SpectreKeyPurpose::Authentication);

        password_hash::derive_encryption_key(salt, password, &PSRAM_ALLOCATOR)
    };

    let output: &mut [u8] = unsafe { core::slice::from_raw_parts_mut(output, output_len) };
    let mut user_key = ACTIVE_ENCRYPTED_USER_KEY.lock(|user_key| user_key.get());

    password_hash::decrypt_with(&mut user_key, &encryption_key);
    output.copy_from_slice(&user_key);

    0
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_crypto_generichash_blake2b_salt_personal(
    data_out: NonNullPtr<*mut c_uchar>,
    data_out_len: c_size_t,
    data_in: *const c_uchar,
    data_in_len: c_ulonglong,
    key: *const c_uchar,
    keylen: c_size_t,
    salt: *const c_uchar,
    personal: *const c_uchar,
) -> c_int {
    todo!()
}

#[allow(non_camel_case_types)]
#[repr(C)]
struct crypto_hash_sha256_state {
    state: [u32; 8],
    count: u64,
    buf: [u8; 64],
}

#[allow(non_camel_case_types, unused)]
struct crypto_auth_hmacsha256_state {
    ictx: crypto_hash_sha256_state,
    octx: crypto_hash_sha256_state,
}

struct UnsafeSend<T>(pub T);

unsafe impl<T> Send for UnsafeSend<T> {}

// TODO: The software implementation is currently faster.
//       My branch `sha-traits` of `esp-hal` implements the necessary traits
//       to make `Sha256Context` work with `SimpleHmac`, but it would likely
//       have to be optimized for `Hmac` instead to be faster.
// type HmacImpl = hmac::SimpleHmac<esp_hal::sha::Sha256Context>;
type HmacImpl = hmac::Hmac<sha2::Sha256>;

static HMAC: blocking_mutex::Mutex<CriticalSectionRawMutex, RefCell<Option<UnsafeSend<HmacImpl>>>> =
    blocking_mutex::Mutex::new(RefCell::new(None));

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_crypto_auth_hmacsha256_init(
    state: NonNullPtr<*mut crypto_auth_hmacsha256_state>,
    key: NonNullPtr<*const c_uchar>,
    key_len: c_size_t,
) -> c_int {
    let key: &[u8] = unsafe { core::slice::from_raw_parts(key, key_len) };
    // TODO: Hardware-accelerated hashing via the SHA peripheral.
    // SHA.lock(|sha| {
    //     let sha = sha
    //         .borrow_mut()
    //         .as_mut()
    //         .expect("HMAC peripheral not initialized.");

    //     sha.
    // });
    HMAC.lock(|hmac| {
        let mut hmac = hmac.borrow_mut();

        if hmac.is_some() {
            panic!("HMAC already initialized. Cannot handle multiple HMAC's at once.");
        }

        // TODO: Implement a lower level `Sha256BlockContext` that would work with `Hmac`.
        //       This requires the implementation of a few `core_api` traits.
        *hmac = Some(UnsafeSend(HmacImpl::new_from_slice(key).unwrap()));
    });

    0
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_crypto_auth_hmacsha256_update(
    state: NonNullPtr<*mut crypto_auth_hmacsha256_state>,
    data_in: *const c_uchar,
    data_in_len: c_ulonglong,
) -> c_int {
    let data_in: &[u8] = unsafe { core::slice::from_raw_parts(data_in, data_in_len as usize) };

    HMAC.lock(|hmac| {
        let mut hmac = hmac.borrow_mut();
        let UnsafeSend(hmac) = hmac
            .as_mut()
            .expect("HMAC must first be initialized before it is updated.");

        hmac.update(data_in);
    });

    0
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_crypto_auth_hmacsha256_final(
    state: NonNullPtr<*mut crypto_auth_hmacsha256_state>,
    out: NonNullPtr<*mut c_uchar>,
) -> c_int {
    let out = unsafe { &mut *(out as *mut GenericArray<u8, U32>) };

    HMAC.lock(|hmac| {
        let mut hmac = hmac.borrow_mut();
        let UnsafeSend(hmac) = hmac
            .take()
            .expect("HMAC must first be initialized before it is updated.");

        hmac.finalize_into(out);
    });

    0
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_sodium_bin2base64(
    b64: NonNullPtr<*mut c_char>,
    b64_maxlen: c_size_t,
    bin: *const c_uchar,
    bin_len: c_size_t,
    variant: c_int,
) -> *mut c_char {
    todo!()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __spre_sodium_base642bin(
    bin: NonNullPtr<*mut c_uchar>,
    bin_maxlen: c_size_t,
    b64: *const c_char,
    b64_len: c_size_t,
    ignore: *const c_char,
    bin_len: *mut c_size_t,
    b64_end: *mut *const c_char,
    variant: c_int,
) -> c_int {
    todo!()
}
