use crate::AppService;
use crate::error::AppError;
use base64::Engine;
use rsa::pkcs1::{DecodeRsaPrivateKey, EncodeRsaPrivateKey, EncodeRsaPublicKey};
use rsa::{Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey};
use serde::{Deserialize, Serialize};
use session::Session;

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct RsaResponse {
    pub public_key: String,
}

impl AppService {
    pub const RSA_PUBLIC_KEY: &'static str = "rsa:public";
    pub const RSA_PRIVATE_KEY: &'static str = "rsa:private";
    const RSA_BIT_SIZE: usize = 2048;
    pub async fn auth_rsa(&self, context: &Session) -> Result<RsaResponse, AppError> {
        #[allow(deprecated)]
        let mut rng = rsa::rand_core::OsRng::default();
        let Ok(priv_key) = RsaPrivateKey::new(&mut rng, Self::RSA_BIT_SIZE) else {
            tracing::error!("RSA key generation failed");
            return Err(AppError::RsaGenerationError);
        };
        let pub_key = RsaPublicKey::from(&priv_key);
        let priv_pem = priv_key
            .to_pkcs1_pem(Default::default())
            .map_err(|_| AppError::RsaGenerationError)?
            .to_string();
        let pub_pem = pub_key
            .to_pkcs1_pem(Default::default())
            .map_err(|_| AppError::RsaGenerationError)?
            .to_string();
        context
            .insert(Self::RSA_PUBLIC_KEY, pub_pem.clone())
            .map_err(|_| AppError::RsaGenerationError)?;
        context
            .insert(Self::RSA_PRIVATE_KEY, priv_pem)
            .map_err(|_| AppError::RsaGenerationError)?;
        Ok(RsaResponse {
            public_key: pub_pem,
        })
    }
    pub async fn auth_rsa_decode(
        &self,
        context: &Session,
        data: String,
    ) -> Result<String, AppError> {
        let priv_key = context
            .get::<String>(Self::RSA_PRIVATE_KEY)
            .map_err(|_| AppError::RsaDecodeError)?
            .ok_or(AppError::RsaDecodeError)?;
        let Ok(priv_key) = RsaPrivateKey::from_pkcs1_pem(&priv_key) else {
            tracing::warn!(ip = ?context.ip_address(), "RSA decode failed: invalid private key");
            return Err(AppError::RsaDecodeError);
        };
        let cipher = base64::engine::general_purpose::STANDARD
            .decode(&data)
            .map_err(|_| AppError::RsaDecodeError)?;
        let Ok(decrypted) = priv_key.decrypt(Pkcs1v15Encrypt, &cipher) else {
            tracing::warn!(ip = ?context.ip_address(), "RSA decrypt failed");
            return Err(AppError::RsaDecodeError);
        };
        Ok(String::from_utf8_lossy(&decrypted).to_string())
    }
}
