//! Provides methods to encrypt/decrypt Bitmessage messages.

use std::{
    fmt,
    io::{self, Read, Write},
    mem::size_of,
};

use block_modes::BlockMode;
use hmac::{Mac, NewMac};

use crate::{
    error::TooLongError,
    hash::sha256,
    io::{LenBm, ReadFrom, SizedReadFrom, WriteTo},
    message,
    var_type::VarInt,
};

pub use crate::error::TooShortError;

const KEY_LEN: usize = 32;

/// A private key.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct PrivateKey {
    bytes: [u8; KEY_LEN],
}

impl PrivateKey {
    /// Constructs a private key from a byte array.
    pub fn new(bytes: [u8; KEY_LEN]) -> Self {
        Self { bytes }
    }
}

impl AsRef<[u8]> for PrivateKey {
    fn as_ref(&self) -> &[u8] {
        &self.bytes
    }
}

/// A public key.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct PublicKey {
    x: [u8; KEY_LEN],
    y: [u8; KEY_LEN],
}

impl PublicKey {
    const PREFIX: u8 = 0x04;

    /// Encodes the public key.
    pub fn encode(&self) -> Vec<u8> {
        let mut encoded = Vec::with_capacity(1 + KEY_LEN + KEY_LEN);
        encoded.push(Self::PREFIX);
        encoded.extend_from_slice(&self.x);
        encoded.extend_from_slice(&self.y);
        encoded
    }
}

impl WriteTo for PublicKey {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        self.x.write_to(w)?;
        self.y.write_to(w)?;
        Ok(())
    }
}

impl ReadFrom for PublicKey {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let x = <[u8; KEY_LEN]>::read_from(r)?;
        let y = <[u8; KEY_LEN]>::read_from(r)?;
        Ok(Self { x, y })
    }
}

type Cipher = block_modes::Cbc<aes::Aes256, block_modes::block_padding::Pkcs7>;

/// This error indicates
/// that the decryption failed.
#[derive(Clone, Debug)]
pub enum DecryptError {
    /// Indicates that at least one of the lengths of the components of the public key was invalid.
    /// The actual lengths of the components are returned as payloads of this variant.
    InvalidKeyLength {
        /// The actual length of the component x.
        x_len: usize,
        /// The actual length of the component y.
        y_len: usize,
    },
    /// An error was caught during initialization of the MAC verifier.
    /// The actual error caught is returned as a payload of this variant.
    InvalidMacKeyLength(crypto_mac::InvalidKeyLength),
    /// Indicates that the secret key is invalid.
    InvalidSecretKey,
    /// An error was caught during verification of the MAC.
    /// The actual error caught is returned as a payload of this variant.
    MacError(crypto_mac::MacError),
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
    /// An error was caught during a process using a block mode.
    /// The actual error caught is returned as a payload of this variant.
    BlockModeError(block_modes::BlockModeError),
}

impl fmt::Display for DecryptError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidKeyLength { x_len, y_len } => {
                write!(f, "invalid key length x: {}, y: {}", x_len, y_len)
            }
            Self::InvalidMacKeyLength(err) => err.fmt(f),
            Self::InvalidSecretKey => "invalid secret key".fmt(f),
            Self::MacError(err) => err.fmt(f),
            Self::Secp256k1Error(err) => err.fmt(f),
            Self::BlockModeError(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for DecryptError {}

impl From<crypto_mac::InvalidKeyLength> for DecryptError {
    fn from(err: crypto_mac::InvalidKeyLength) -> Self {
        Self::InvalidMacKeyLength(err)
    }
}

impl From<crypto_mac::MacError> for DecryptError {
    fn from(err: crypto_mac::MacError) -> Self {
        Self::MacError(err)
    }
}

impl From<secp256k1::Error> for DecryptError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

impl From<block_modes::BlockModeError> for DecryptError {
    fn from(err: block_modes::BlockModeError) -> Self {
        Self::BlockModeError(err)
    }
}

const IV_LEN: usize = 16;
const MAC_LEN: usize = 32;

/// A encrypted message.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Encrypted {
    iv: [u8; IV_LEN],
    curve_type: u16,
    //x_length: u16,
    x: Vec<u8>,
    //y_length: u16,
    y: Vec<u8>,
    encrypted: Vec<u8>,
    mac: [u8; MAC_LEN],
}

impl Encrypted {
    /// Decrypts this encrypted object and returns the decrypted bytes.
    pub fn decrypt(&self, private_key: &PrivateKey) -> Result<Vec<u8>, DecryptError> {
        if self.x.len() != KEY_LEN || self.y.len() != KEY_LEN {
            return Err(DecryptError::InvalidKeyLength {
                x_len: self.x.len(),
                y_len: self.y.len(),
            });
        }
        let mut public_key_bytes = [0; 65];
        public_key_bytes[0] = 0x04;
        public_key_bytes[1..33].copy_from_slice(&self.x);
        public_key_bytes[33..].copy_from_slice(&self.y);
        let public_key = secp256k1::PublicKey::from_slice(&public_key_bytes)?;
        let secret_key = secp256k1::SecretKey::from_slice(private_key.as_ref())?;
        let shared_secret =
            secp256k1::ecdh::SharedSecret::new_with_hash(&public_key, &secret_key, |x, _y| {
                x.into()
            })?;

        let h = crate::hash::sha512(shared_secret);
        let mac_key = &h[32..];

        let mut mac = hmac::Hmac::<sha2::Sha256>::new_varkey(mac_key)?;
        mac.update(&self.iv);
        mac.update(&self.curve_type.to_be_bytes());
        mac.update(&(self.x.len() as u16).to_be_bytes());
        mac.update(&self.x);
        mac.update(&(self.y.len() as u16).to_be_bytes());
        mac.update(&self.y);
        mac.update(&self.encrypted);
        mac.verify(&self.mac)?;

        let encryption_key = &h[..32];
        let cipher = Cipher::new_fix(encryption_key.into(), self.iv.as_ref().into());
        Ok(cipher.decrypt_vec(&self.encrypted)?)
    }
}

impl SizedReadFrom for Encrypted {
    fn sized_read_from(r: &mut dyn Read, len: usize) -> io::Result<Self>
    where
        Self: Sized,
    {
        let mut iv = [0; IV_LEN];
        r.read_exact(&mut iv)?;

        let curve_type = u16::read_from(r)?;

        let x_length = u16::read_from(r)?;

        let mut rt = r.take(x_length as u64);
        let mut x = Vec::with_capacity(x_length as usize);
        rt.read_to_end(&mut x)?;

        let y_length = u16::read_from(r)?;

        let mut rt = r.take(y_length as u64);
        let mut y = Vec::with_capacity(y_length as usize);
        rt.read_to_end(&mut y)?;

        let sub_len = IV_LEN
            + size_of::<u16>()
            + size_of::<u16>()
            + x_length as usize
            + size_of::<u16>()
            + y_length as usize
            + MAC_LEN;
        if sub_len > len {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooShortError::new(sub_len, len),
            ));
        }
        let encrypted_len = len - sub_len;
        let mut rt = r.take(encrypted_len as u64);
        let mut encrypted = Vec::with_capacity(encrypted_len);
        rt.read_to_end(&mut encrypted)?;

        let mut mac = [0; MAC_LEN];
        r.read_exact(&mut mac)?;

        Ok(Self {
            iv,
            curve_type,
            x,
            y,
            encrypted,
            mac,
        })
    }
}

/// This error indicates
/// that the signing failed.
#[derive(Clone, Debug)]
pub enum SignError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for SignError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for SignError {}

impl From<secp256k1::Error> for SignError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// This error indicates
/// that the verification failed.
#[derive(Clone, Debug)]
pub enum VerifyError {
    /// An error was caught during a process using secp256k1.
    /// The actual error caught is returned as a payload of this variant.
    Secp256k1Error(secp256k1::Error),
}

impl fmt::Display for VerifyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1Error(err) => err.fmt(f),
        }
    }
}

impl std::error::Error for VerifyError {}

impl From<secp256k1::Error> for VerifyError {
    fn from(err: secp256k1::Error) -> Self {
        Self::Secp256k1Error(err)
    }
}

/// A signature.
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Signature {
    bytes: Vec<u8>,
}

fn sign_sha256(data: impl AsRef<[u8]>, private_key: &PrivateKey) -> Result<Signature, SignError> {
    let hash = sha256(data);
    let message = secp256k1::Message::from_slice(hash.as_ref())?;
    let secret_key = secp256k1::SecretKey::from_slice(private_key.as_ref())?;
    let secp = secp256k1::Secp256k1::signing_only();
    let signature = secp.sign(&message, &secret_key);
    let bytes = Vec::from(signature.serialize_der().as_ref());
    Ok(Signature { bytes })
}

/// Signs on data using a private key.
pub fn sign(data: impl AsRef<[u8]>, private_key: &PrivateKey) -> Result<Signature, SignError> {
    sign_sha256(data, private_key)
}

/// Verifies a signature of data using a public key.
pub fn verify(
    data: impl AsRef<[u8]>,
    signature: &Signature,
    public_key: &PublicKey,
) -> Result<(), VerifyError> {
    let hash = sha256(data);
    let message = secp256k1::Message::from_slice(hash.as_ref())?;
    let mut sig = secp256k1::Signature::from_der(&signature.bytes)?;
    sig.normalize_s();
    let mut public_key_bytes = [0; 65];
    public_key_bytes[0] = 0x04;
    public_key_bytes[1..33].copy_from_slice(&public_key.x);
    public_key_bytes[33..].copy_from_slice(&public_key.y);
    let pk = secp256k1::PublicKey::from_slice(&public_key_bytes)?;
    let secp = secp256k1::Secp256k1::verification_only();
    secp.verify(&message, &sig, &pk)?;
    Ok(())
}

impl WriteTo for Signature {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let len: VarInt = self.bytes.len().into();
        len.write_to(w)?;
        self.bytes.write_to(w)?;
        Ok(())
    }
}

impl ReadFrom for Signature {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let len = VarInt::read_from(r)?;
        if len.as_u64() > message::Object::MAX_OBJECT_PAYLOAD_LENGTH as u64 {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                TooLongError::new(
                    message::Object::MAX_OBJECT_PAYLOAD_LENGTH,
                    len.as_u64() as usize,
                ),
            ));
        }
        let bytes = Vec::<u8>::sized_read_from(r, len.as_u64() as usize)?;
        Ok(Self { bytes })
    }
}

impl LenBm for Signature {
    fn len_bm(&self) -> usize {
        let len: VarInt = self.bytes.len().into();
        len.len_bm() + self.bytes.len()
    }
}
