use rsa::{
    pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey, EncodeRsaPrivateKey, EncodeRsaPublicKey},
    pkcs8::{DecodePrivateKey, DecodePublicKey},
    signature::{Keypair, RandomizedSigner, SignatureEncoding, Verifier},
    Pkcs1v15Encrypt,
};
use sha2::Sha256;
use thiserror::Error;

pub enum RsaBitSize {
    Bit1024,
    Bit2048,
}

#[derive(Debug, Error)]
pub enum RsaCryptionError {
    #[error("Invalid public key")]
    InvalidPublicKey,
    #[error("Invalid private key")]
    InvalidPrivateKey,
    #[error("Unrecognized private key encoding or encrypted private key")]
    UnrecognizedPrivateKeyEncoding,
    #[error("Unrecognized public key encoding")]
    UnrecognizedPublicKeyEncoding,
    #[error("Private key required")]
    PrivateKeyRequired,
    #[error("Public key required")]
    PublicKeyRequired,
    #[error("Signing key required")]
    SigningKeyRequired,
    #[error("Verifying key required")]
    VerifyingKeyRequired,
    #[error("Failed to process data: {0}")]
    CryptionFailed(#[from] rsa::errors::Error),
    #[error("Unable to load signature")]
    UnableToLoadSignature,
    #[error("Signature verification failed")]
    VerificationFailed,
    #[error("Unable to export private key")]
    UnableToExportPrivateKey,
    #[error("Unable to export public key")]
    UnableToExportPublicKey,
}

#[derive(Debug)]
pub struct RsaCryptor {
    private_key: Option<rsa::RsaPrivateKey>,
    public_key: Option<rsa::RsaPublicKey>,
    signing_key: Option<rsa::pkcs1v15::SigningKey<Sha256>>,
    verifying_key: Option<rsa::pkcs1v15::VerifyingKey<Sha256>>,
}

fn load_binary_private_key(key: &[u8]) -> Result<rsa::RsaPrivateKey, RsaCryptionError> {
    let private_key = rsa::RsaPrivateKey::from_pkcs1_der(&key);
    if let Ok(private_key) = private_key {
        if private_key.validate().is_err() {
            return Err(RsaCryptionError::InvalidPrivateKey);
        }
        return Ok(private_key);
    }
    let private_key = rsa::RsaPrivateKey::from_pkcs8_der(&key);
    if let Ok(private_key) = private_key {
        if private_key.validate().is_err() {
            return Err(RsaCryptionError::InvalidPrivateKey);
        }
        return Ok(private_key);
    }
    Err(RsaCryptionError::UnrecognizedPrivateKeyEncoding)
}

fn load_text_private_key(pem: &str) -> Result<rsa::RsaPrivateKey, RsaCryptionError> {
    let private_key = rsa::RsaPrivateKey::from_pkcs1_pem(pem);
    if let Ok(private_key) = private_key {
        if private_key.validate().is_err() {
            return Err(RsaCryptionError::InvalidPrivateKey);
        }
        return Ok(private_key);
    }
    let private_key = rsa::RsaPrivateKey::from_pkcs8_pem(pem);
    if let Ok(private_key) = private_key {
        if private_key.validate().is_err() {
            return Err(RsaCryptionError::InvalidPrivateKey);
        }
        return Ok(private_key);
    }
    Err(RsaCryptionError::UnrecognizedPrivateKeyEncoding)
}

fn load_binary_public_key(key: &[u8]) -> Result<rsa::RsaPublicKey, RsaCryptionError> {
    let public_key = rsa::RsaPublicKey::from_pkcs1_der(&key);
    if let Ok(public_key) = public_key {
        return Ok(public_key);
    }
    let public_key = rsa::RsaPublicKey::from_public_key_der(&key);
    if let Ok(public_key) = public_key {
        return Ok(public_key);
    }
    Err(RsaCryptionError::UnrecognizedPublicKeyEncoding)
}

fn load_text_public_key(pem: &str) -> Result<rsa::RsaPublicKey, RsaCryptionError> {
    let public_key = rsa::RsaPublicKey::from_pkcs1_pem(pem);
    if let Ok(public_key) = public_key {
        return Ok(public_key);
    }
    let public_key = rsa::RsaPublicKey::from_public_key_pem(pem);
    if let Ok(public_key) = public_key {
        return Ok(public_key);
    }
    Err(RsaCryptionError::UnrecognizedPublicKeyEncoding)
}

impl RsaCryptor {
    /// 创建一个完全空白的 RSA 加密器。其中私钥、公钥、签名私钥和签名公钥都是空的，需要通过其他方法加载后方可使用。
    pub fn new_empty() -> Self {
        Self {
            private_key: None,
            public_key: None,
            signing_key: None,
            verifying_key: None,
        }
    }

    /// 创建一个指定位数的 RSA 加密器，同时生成一套随机密钥。
    /// 如果不能正常生成私钥，则直接返回错误。
    ///
    /// - `bit_size`：密钥位数，目前支持 1024 位和 2048 位。
    pub fn new(bit_size: RsaBitSize) -> Result<Self, RsaCryptionError> {
        let mut rng = rand::thread_rng();
        let bit_size = match bit_size {
            RsaBitSize::Bit1024 => 1024,
            RsaBitSize::Bit2048 => 2048,
        };
        let private_key = rsa::RsaPrivateKey::new(&mut rng, bit_size)
            .map_err(|_| RsaCryptionError::InvalidPublicKey)?;
        let public_key = rsa::RsaPublicKey::from(&private_key);
        let signing_key = rsa::pkcs1v15::SigningKey::<Sha256>::new(private_key.clone());
        let verify_key = signing_key.verifying_key();
        Ok(Self {
            private_key: Some(private_key),
            public_key: Some(public_key),
            signing_key: Some(signing_key),
            verifying_key: Some(verify_key),
        })
    }

    /// 从字节数组中加载私钥，如果加载失败，则返回错误。
    /// 私钥加载成功后，公钥、签名私钥和签名公钥也会被同时自动生成。
    ///
    /// - `key`：私钥内容。
    pub fn load_binary_private_key<T: AsRef<[u8]>>(
        &mut self,
        key: T,
    ) -> Result<(), RsaCryptionError> {
        let private_key = load_binary_private_key(key.as_ref())?;
        let public_key = rsa::RsaPublicKey::from(&private_key);
        let signing_key = rsa::pkcs1v15::SigningKey::<Sha256>::new(private_key.clone());
        let verify_key = signing_key.verifying_key();
        self.private_key = Some(private_key);
        self.public_key = Some(public_key);
        self.signing_key = Some(signing_key);
        self.verifying_key = Some(verify_key);
        Ok(())
    }

    /// 从文本中加载私钥，如果加载失败，则返回错误。
    /// 私钥加载成功后，公钥、签名私钥和签名公钥也会被同时自动生成。
    ///
    /// - `pem`：私钥内容。
    pub fn load_text_private_key(&mut self, pem: &str) -> Result<(), RsaCryptionError> {
        let private_key = load_text_private_key(pem)?;
        let public_key = rsa::RsaPublicKey::from(&private_key);
        let signing_key = rsa::pkcs1v15::SigningKey::<Sha256>::new(private_key.clone());
        let verify_key = signing_key.verifying_key();
        self.private_key = Some(private_key);
        self.public_key = Some(public_key);
        self.signing_key = Some(signing_key);
        self.verifying_key = Some(verify_key);
        Ok(())
    }

    /// 从字节数组中加载公钥，如果加载失败，则返回错误。
    /// 公钥加载成功后，签名公钥也会被同时自动生成。私钥与签名私钥会被置为`None`。
    /// ! 没有私钥的情况下，无法对内容进行解密。
    ///
    /// - `key`：公钥内容。
    pub fn load_binary_public_key<T: AsRef<[u8]>>(
        &mut self,
        key: T,
    ) -> Result<(), RsaCryptionError> {
        let public_key = load_binary_public_key(key.as_ref())?;
        let verify_key = rsa::pkcs1v15::VerifyingKey::<Sha256>::from(public_key.clone());
        self.private_key = None;
        self.public_key = Some(public_key);
        self.signing_key = None;
        self.verifying_key = Some(verify_key);
        Ok(())
    }

    /// 从文本中加载公钥，如果加载失败，则返回错误。
    /// 公钥加载成功后，签名公钥也会被同时自动生成。私钥与签名私钥会被置为`None`。
    /// ! 没有私钥的情况下，无法对内容进行解密。
    ///
    /// - `pem`：公钥内容。
    pub fn load_text_public_key(&mut self, pem: &str) -> Result<(), RsaCryptionError> {
        let public_key = load_text_public_key(pem)?;
        let verify_key = rsa::pkcs1v15::VerifyingKey::<Sha256>::from(public_key.clone());
        self.private_key = None;
        self.public_key = Some(public_key);
        self.signing_key = None;
        self.verifying_key = Some(verify_key);
        Ok(())
    }

    /// 对指定内容进行加密，返回加密后的内容。
    ///
    /// - `data`：待加密的内容。
    pub fn encrypt<T: AsRef<[u8]>>(&self, data: T) -> Result<Vec<u8>, RsaCryptionError> {
        if self.public_key.is_none() {
            return Err(RsaCryptionError::PublicKeyRequired);
        }
        let mut rng = rand::thread_rng();
        let encrypted_data = self
            .public_key
            .as_ref()
            .unwrap()
            .encrypt(&mut rng, Pkcs1v15Encrypt, data.as_ref())
            .map_err(|e| RsaCryptionError::CryptionFailed(e))?;
        Ok(encrypted_data)
    }

    /// 对指定内容进行解密，返回解密后的内容。如果没有私钥，则返回错误。
    ///
    /// - `data`：待解密的内容。
    pub fn decrypt<T: AsRef<[u8]>>(&self, data: T) -> Result<Vec<u8>, RsaCryptionError> {
        if self.private_key.is_none() {
            return Err(RsaCryptionError::PrivateKeyRequired);
        }
        let decrypted_data = self
            .private_key
            .as_ref()
            .unwrap()
            .decrypt(Pkcs1v15Encrypt, data.as_ref())
            .map_err(|e| RsaCryptionError::CryptionFailed(e))?;
        Ok(decrypted_data)
    }

    /// 对指定内容进行签名，返回签名后的内容。如果没有签名私钥，则返回错误。
    ///
    /// - `data`：待签名的内容。
    pub fn sign<T: AsRef<[u8]>>(&self, data: T) -> Result<Vec<u8>, RsaCryptionError> {
        if self.signing_key.is_none() {
            return Err(RsaCryptionError::SigningKeyRequired);
        }
        let mut rng = rand::thread_rng();
        let sign_data = self
            .signing_key
            .as_ref()
            .unwrap()
            .sign_with_rng(&mut rng, data.as_ref())
            .to_vec();
        Ok(sign_data)
    }

    /// 对指定内容和签名进行验证，如果验证失败，则返回错误。
    ///
    /// - `signature`：签名内容。
    /// - `data`：待验证的内容。
    pub fn verify<T: AsRef<[u8]>>(&self, signature: T, data: T) -> Result<(), RsaCryptionError> {
        if self.verifying_key.is_none() {
            return Err(RsaCryptionError::VerifyingKeyRequired);
        }
        let signature = rsa::pkcs1v15::Signature::try_from(signature.as_ref())
            .map_err(|_| RsaCryptionError::UnableToLoadSignature)?;
        self.verifying_key
            .as_ref()
            .unwrap()
            .verify(data.as_ref(), &signature)
            .map_err(|_| RsaCryptionError::VerificationFailed)?;
        Ok(())
    }

    /// 以字节数组方式导出私钥，如果没有私钥，则返回错误。
    pub fn export_private_key_binary(&self) -> Result<Vec<u8>, RsaCryptionError> {
        if self.private_key.is_none() {
            return Err(RsaCryptionError::PrivateKeyRequired);
        }
        let doc = self
            .private_key
            .as_ref()
            .unwrap()
            .to_pkcs1_der()
            .map_err(|_| RsaCryptionError::UnableToExportPrivateKey)?;
        Ok(doc.as_bytes().to_vec())
    }

    /// 以PEM文本方式导出私钥，如果没有私钥，则返回错误。
    pub fn export_private_key_pem(&self) -> Result<String, RsaCryptionError> {
        if self.private_key.is_none() {
            return Err(RsaCryptionError::PrivateKeyRequired);
        }
        let doc = self
            .private_key
            .as_ref()
            .unwrap()
            .to_pkcs1_pem(rsa::pkcs1::LineEnding::CRLF)
            .map_err(|_| RsaCryptionError::UnableToExportPrivateKey)?;
        Ok(doc.to_string())
    }

    /// 以字节数组的方式导出公钥，如果没有公钥，则返回错误。
    pub fn export_public_key_binary(&self) -> Result<Vec<u8>, RsaCryptionError> {
        if self.public_key.is_none() {
            return Err(RsaCryptionError::PublicKeyRequired);
        }
        let doc = self
            .public_key
            .as_ref()
            .unwrap()
            .to_pkcs1_der()
            .map_err(|_| RsaCryptionError::UnableToExportPublicKey)?;
        Ok(doc.into_vec())
    }

    /// 以PEM文本方式导出公钥，如果没有公钥，则返回错误。
    pub fn export_public_key_pem(&self) -> Result<String, RsaCryptionError> {
        if self.public_key.is_none() {
            return Err(RsaCryptionError::PublicKeyRequired);
        }
        let doc = self
            .public_key
            .as_ref()
            .unwrap()
            .to_pkcs1_pem(rsa::pkcs1::LineEnding::CRLF)
            .map_err(|_| RsaCryptionError::UnableToExportPublicKey)?;
        Ok(doc)
    }
}