use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;

use crate::error::{CryptError, CryptResult};
use crate::key::{RsaDecryptor, RsaEncryptor, random_bytes};

pub const ENVELOPE_PREFIX: &str = "app1.";
pub const ENVELOPE_VERSION: u8 = 1;
pub const ALGORITHM_RSA_OAEP_256_A256GCM: u8 = 1;
pub const AES_KEY_LEN: usize = 32;
pub const NONCE_LEN: usize = 12;
const KEY_ID_LEN_BYTES: usize = 1;
const WRAPPED_LEN_BYTES: usize = 2;

#[derive(Debug, Clone)]
pub struct EnvelopeParts {
    pub version: u8,
    pub algorithm: u8,
    pub key_id: String,
    pub wrapped_key: Vec<u8>,
    pub nonce: Vec<u8>,
    pub ciphertext: Vec<u8>,
}

pub fn encrypt(encryptor: &RsaEncryptor, plaintext: &[u8]) -> CryptResult<String> {
    let aes_key = random_bytes(AES_KEY_LEN);
    let nonce_bytes = random_bytes(NONCE_LEN);
    let key_id = encryptor.key_id();
    let aad = associated_data(key_id);
    let cipher = Aes256Gcm::new_from_slice(&aes_key).map_err(|_| CryptError::Aead)?;
    let nonce = Nonce::try_from(nonce_bytes.as_slice()).map_err(|_| CryptError::Aead)?;
    let ciphertext = cipher
        .encrypt(
            &nonce,
            Payload {
                msg: plaintext,
                aad: &aad,
            },
        )
        .map_err(|_| CryptError::Aead)?;
    let wrapped_key = encryptor.encrypt(&aes_key)?;
    Ok(assemble(
        ENVELOPE_VERSION,
        ALGORITHM_RSA_OAEP_256_A256GCM,
        key_id,
        &wrapped_key,
        &nonce_bytes,
        &ciphertext,
    ))
}

pub fn decrypt(decryptor: &RsaDecryptor, envelope: &str) -> CryptResult<Vec<u8>> {
    let parts = decode(envelope)?;
    let aes_key = decryptor.decrypt(&parts.key_id, &parts.wrapped_key)?;
    let aad = associated_data(&parts.key_id);
    let cipher = Aes256Gcm::new_from_slice(&aes_key).map_err(|_| CryptError::Aead)?;
    let nonce = Nonce::try_from(parts.nonce.as_slice()).map_err(|_| CryptError::Aead)?;
    cipher
        .decrypt(
            &nonce,
            Payload {
                msg: parts.ciphertext.as_slice(),
                aad: &aad,
            },
        )
        .map_err(|_| CryptError::Aead)
}

pub fn encrypt_string(encryptor: &RsaEncryptor, plaintext: &str) -> CryptResult<String> {
    encrypt(encryptor, plaintext.as_bytes())
}

pub fn decrypt_string(decryptor: &RsaDecryptor, envelope: &str) -> CryptResult<String> {
    String::from_utf8(decrypt(decryptor, envelope)?).map_err(CryptError::from)
}

pub fn envelope_key_id(envelope: &str) -> CryptResult<String> {
    Ok(decode(envelope)?.key_id)
}

pub fn decode(envelope: &str) -> CryptResult<EnvelopeParts> {
    let trimmed = envelope.trim();
    let body = trimmed.strip_prefix(ENVELOPE_PREFIX).ok_or_else(|| {
        CryptError::EnvelopePrefix({
            let value = trimmed;
            value.chars().take(16).collect()
        })
    })?;
    let raw = URL_SAFE_NO_PAD
        .decode(body)
        .map_err(|error| CryptError::Base64(error.to_string()))?;
    let mut cursor = Cursor::new(&raw);
    let version = cursor.take_u8()?;
    if version != ENVELOPE_VERSION {
        return Err(CryptError::UnsupportedVersion(version));
    }
    let algorithm = cursor.take_u8()?;
    if algorithm != ALGORITHM_RSA_OAEP_256_A256GCM {
        return Err(CryptError::UnsupportedAlgorithm(algorithm));
    }
    let key_id_len = cursor.take_u8()? as usize;
    let key_id = String::from_utf8(cursor.take(key_id_len)?.to_vec())
        .map_err(|_| CryptError::KeyIdEncoding)?;
    let wrapped_len = cursor.take_u16()? as usize;
    let wrapped_key = cursor.take(wrapped_len)?.to_vec();
    let nonce = cursor.take(NONCE_LEN)?.to_vec();
    let ciphertext = cursor.remaining().to_vec();
    if ciphertext.len() < 16 {
        return Err(CryptError::EnvelopeTruncated);
    }
    Ok(EnvelopeParts {
        version,
        algorithm,
        key_id,
        wrapped_key,
        nonce,
        ciphertext,
    })
}

fn assemble(
    version: u8,
    algorithm: u8,
    key_id: &str,
    wrapped_key: &[u8],
    nonce: &[u8],
    ciphertext: &[u8],
) -> String {
    let mut raw = Vec::with_capacity(
        2 + KEY_ID_LEN_BYTES + key_id.len() + WRAPPED_LEN_BYTES + wrapped_key.len() + NONCE_LEN,
    );
    raw.push(version);
    raw.push(algorithm);
    raw.push(key_id.len() as u8);
    raw.extend_from_slice(key_id.as_bytes());
    raw.extend_from_slice(&(wrapped_key.len() as u16).to_be_bytes());
    raw.extend_from_slice(wrapped_key);
    raw.extend_from_slice(nonce);
    raw.extend_from_slice(ciphertext);
    format!("{ENVELOPE_PREFIX}{}", URL_SAFE_NO_PAD.encode(raw))
}

fn associated_data(key_id: &str) -> Vec<u8> {
    let mut aad = Vec::with_capacity(3 + key_id.len());
    aad.push(ENVELOPE_VERSION);
    aad.push(ALGORITHM_RSA_OAEP_256_A256GCM);
    aad.extend_from_slice(key_id.as_bytes());
    aad
}

struct Cursor<'a> {
    raw: &'a [u8],
    offset: usize,
}

impl<'a> Cursor<'a> {
    fn new(raw: &'a [u8]) -> Self {
        Self { raw, offset: 0 }
    }

    fn take(&mut self, len: usize) -> CryptResult<&'a [u8]> {
        let end = self
            .offset
            .checked_add(len)
            .ok_or(CryptError::EnvelopeTruncated)?;
        if end > self.raw.len() {
            return Err(CryptError::EnvelopeTruncated);
        }
        let slice = &self.raw[self.offset..end];
        self.offset = end;
        Ok(slice)
    }

    fn take_u8(&mut self) -> CryptResult<u8> {
        Ok(self.take(1)?[0])
    }

    fn take_u16(&mut self) -> CryptResult<u16> {
        let bytes = self.take(WRAPPED_LEN_BYTES)?;
        Ok(u16::from_be_bytes([bytes[0], bytes[1]]))
    }

    fn remaining(&self) -> &'a [u8] {
        &self.raw[self.offset..]
    }
}