use std::fmt;

use rand_core::{OsRng, RngCore};
use rsa::pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey};
use rsa::pkcs8::spki::{DecodePublicKey, EncodePublicKey};
use rsa::pkcs8::{DecodePrivateKey, EncodePrivateKey, LineEnding};
use rsa::traits::PublicKeyParts;
use rsa::{Oaep, RsaPrivateKey, RsaPublicKey};
use sha2::{Digest, Sha256};

use crate::error::{CryptError, CryptResult};

pub const DEFAULT_KEY_BITS: usize = 4096;
pub const MIN_KEY_BITS: usize = 2048;
pub const FINGERPRINT_BYTES: usize = 8;

#[derive(Clone)]
pub struct RsaKeyPair {
    key_id: String,
    private: RsaPrivateKey,
}

impl RsaKeyPair {
    pub fn generate(key_id: impl Into<String>, bits: usize) -> CryptResult<Self> {
        if bits < MIN_KEY_BITS {
            return Err(CryptError::KeySize(bits, MIN_KEY_BITS));
        }
        let private = RsaPrivateKey::new(&mut OsRng, bits)
            .map_err(|error| CryptError::Rsa(error.to_string()))?;
        Ok(Self {
            key_id: key_id.into(),
            private,
        })
    }

    pub fn from_private_pem(key_id: impl Into<String>, pem: &str) -> CryptResult<Self> {
        let private = parse_private_pem(pem)?;
        Ok(Self {
            key_id: key_id.into(),
            private,
        })
    }

    pub fn from_pem_pair(
        key_id: impl Into<String>,
        private_pem: &str,
        public_pem: &str,
    ) -> CryptResult<Self> {
        let pair = Self::from_private_pem(key_id, private_pem)?;
        let expected = parse_public_pem(public_pem)?;
        if expected != pair.private.to_public_key() {
            return Err(CryptError::KeyMismatch);
        }
        Ok(pair)
    }

    pub fn key_id(&self) -> &str {
        &self.key_id
    }

    pub fn public(&self) -> RsaPublicKey {
        self.private.to_public_key()
    }

    pub fn private_pem(&self) -> CryptResult<String> {
        let pem = self
            .private
            .to_pkcs8_pem(LineEnding::LF)
            .map_err(|error| CryptError::Pem(error.to_string()))?;
        Ok(pem.to_string())
    }

    pub fn public_pem(&self) -> CryptResult<String> {
        encode_public_pem(&self.private.to_public_key())
    }

    pub fn fingerprint(&self) -> String {
        fingerprint_of(&self.private.to_public_key())
    }

    pub fn decrypt(&self, ciphertext: &[u8]) -> CryptResult<Vec<u8>> {
        self.private
            .decrypt(Oaep::new::<Sha256>(), ciphertext)
            .map_err(|error| CryptError::Rsa(error.to_string()))
    }

    pub fn encrypt(&self, plaintext: &[u8]) -> CryptResult<Vec<u8>> {
        self.private
            .to_public_key()
            .encrypt(&mut OsRng, Oaep::new::<Sha256>(), plaintext)
            .map_err(|error| CryptError::Rsa(error.to_string()))
    }
}

impl fmt::Debug for RsaKeyPair {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RsaKeyPair")
            .field("key_id", &self.key_id)
            .field("bits", &self.private.n().bits())
            .field("fingerprint", &self.fingerprint())
            .finish_non_exhaustive()
    }
}

#[derive(Clone)]
pub struct RsaEncryptor {
    key_id: String,
    public: RsaPublicKey,
}

impl RsaEncryptor {
    pub fn from_public_pem(key_id: impl Into<String>, pem: &str) -> CryptResult<Self> {
        Ok(Self {
            key_id: key_id.into(),
            public: parse_public_pem(pem)?,
        })
    }

    pub fn key_id(&self) -> &str {
        &self.key_id
    }

    pub fn fingerprint(&self) -> String {
        fingerprint_of(&self.public)
    }

    pub fn public_pem(&self) -> CryptResult<String> {
        encode_public_pem(&self.public)
    }

    pub fn encrypt(&self, plaintext: &[u8]) -> CryptResult<Vec<u8>> {
        self.public
            .encrypt(&mut OsRng, Oaep::new::<Sha256>(), plaintext)
            .map_err(|error| CryptError::Rsa(error.to_string()))
    }
}

impl fmt::Debug for RsaEncryptor {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RsaEncryptor")
            .field("key_id", &self.key_id)
            .field("fingerprint", &self.fingerprint())
            .finish()
    }
}

pub struct RsaDecryptor {
    keys: Vec<RsaKeyPair>,
}

impl RsaDecryptor {
    pub fn new(keys: Vec<RsaKeyPair>) -> CryptResult<Self> {
        if keys.is_empty() {
            return Err(CryptError::NoKeys);
        }
        Ok(Self { keys })
    }

    pub fn key_ids(&self) -> Vec<&str> {
        self.keys.iter().map(RsaKeyPair::key_id).collect()
    }

    pub fn pair(&self, key_id: &str) -> Option<&RsaKeyPair> {
        self.keys.iter().find(|key| key.key_id == key_id)
    }

    pub fn single(&self) -> &RsaKeyPair {
        &self.keys[0]
    }

    pub fn decrypt(&self, key_id: &str, ciphertext: &[u8]) -> CryptResult<Vec<u8>> {
        if let Some(key) = self.pair(key_id) {
            return key.decrypt(ciphertext);
        }
        let mut last = CryptError::KeyNotFound(key_id.to_owned());
        for key in &self.keys {
            match key.decrypt(ciphertext) {
                Ok(plaintext) => return Ok(plaintext),
                Err(error) => last = error,
            }
        }
        Err(last)
    }
}

impl fmt::Debug for RsaDecryptor {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RsaDecryptor")
            .field("key_ids", &self.key_ids())
            .finish()
    }
}

pub fn fingerprint_of(public: &RsaPublicKey) -> String {
    let der = public.to_public_key_der();
    let Ok(document) = der else {
        return String::new();
    };
    let digest = Sha256::digest(document.as_bytes());
    digest[..FINGERPRINT_BYTES]
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect()
}

pub fn random_bytes(len: usize) -> Vec<u8> {
    let mut bytes = vec![0u8; len];
    OsRng.fill_bytes(&mut bytes);
    bytes
}

fn parse_private_pem(pem: &str) -> CryptResult<RsaPrivateKey> {
    let trimmed = pem.trim();
    RsaPrivateKey::from_pkcs8_pem(trimmed)
        .or_else(|_| RsaPrivateKey::from_pkcs1_pem(trimmed))
        .map_err(|error| CryptError::Pem(error.to_string()))
}

fn parse_public_pem(pem: &str) -> CryptResult<RsaPublicKey> {
    let trimmed = pem.trim();
    RsaPublicKey::from_public_key_pem(trimmed)
        .or_else(|_| RsaPublicKey::from_pkcs1_pem(trimmed))
        .map_err(|error| CryptError::Pem(error.to_string()))
}

fn encode_public_pem(public: &RsaPublicKey) -> CryptResult<String> {
    public
        .to_public_key_pem(LineEnding::LF)
        .map_err(|error| CryptError::Pem(error.to_string()))
}