use std::{borrow::Cow, fmt::Display};

use rsa::signature::Verifier;
use rsa::{
    sha2::{Sha256, Sha512},
    signature::Signer,
};

use serde::{Deserialize, Serialize};
use strum::AsRefStr;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ApVerificationError {
    #[error("Signature algorithm and public key mismatch: {0} not compatible with {1}")]
    KeyAlgorithmMismatch(SigningAlgorithm, String),
    #[error("PKCS#1 v1.5 RSA signature verification failed: {0}")]
    RsaSignatureError(#[from] rsa::signature::Error),
}

#[derive(Debug, Error)]
pub enum ApSigningError {
    #[error("Signature algorithm and public key mismatch: {0} not compatible with {1}")]
    KeyAlgorithmMismatch(SigningAlgorithm, String),
    #[error("Signing error: {0}")]
    RsaSigningError(#[from] rsa::signature::Error),
}

#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub enum SigningAlgorithm {
    #[serde(rename = "hs2019")]
    Hs2019,
    #[serde(rename = "rsa-sha256")]
    RsaSha256,
}

impl Display for SigningAlgorithm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Hs2019 => write!(f, "hs2019"),
            Self::RsaSha256 => write!(f, "rsa-sha256"),
        }
    }
}

#[derive(Debug, Clone, AsRefStr)]
pub enum ApHttpVerificationKey<'a> {
    #[strum(serialize = "rsa-sha256")]
    RsaSha256(Cow<'a, rsa::pkcs1v15::VerifyingKey<Sha256>>),
    #[strum(serialize = "rsa-sha512")]
    RsaSha512(Cow<'a, rsa::pkcs1v15::VerifyingKey<Sha512>>),
    #[strum(serialize = "ed25519")]
    Ed25519(Cow<'a, ed25519_dalek::VerifyingKey>),
}

#[derive(Debug, Clone, AsRefStr)]
pub enum ApHttpPublicKey<'a> {
    #[strum(serialize = "rsa")]
    Rsa(Cow<'a, rsa::RsaPublicKey>),
    #[strum(serialize = "ed25519")]
    Ed25519(Cow<'a, ed25519_dalek::VerifyingKey>),
}

impl ApHttpVerificationKey<'_> {
    pub fn verify(&self, message: &[u8], signature: &[u8]) -> Result<(), ApVerificationError> {
        match self {
            ApHttpVerificationKey::RsaSha256(rsa_pubkey) => {
                let sig = rsa::pkcs1v15::Signature::try_from(signature)?;
                Ok(rsa_pubkey.verify(message, &sig)?)
            }
            ApHttpVerificationKey::RsaSha512(rsa_pubkey) => {
                let sig = rsa::pkcs1v15::Signature::try_from(signature)?;
                Ok(rsa_pubkey.verify(message, &sig)?)
            }
            ApHttpVerificationKey::Ed25519(ed25519_pubkey) => {
                let sig = ed25519_dalek::Signature::try_from(signature)?;
                Ok(ed25519_pubkey.verify_strict(message, &sig)?)
            }
        }
    }
}

impl ApHttpPublicKey<'_> {
    pub fn verify(
        &self,
        algorithm: SigningAlgorithm,
        message: &[u8],
        signature: &[u8],
    ) -> Result<(), ApVerificationError> {
        match (self, algorithm) {
            (Self::Rsa(key), SigningAlgorithm::Hs2019) => {
                let verification_key = ApHttpVerificationKey::RsaSha256(Cow::Owned(
                    rsa::pkcs1v15::VerifyingKey::new(key.clone().into_owned()),
                ));

                let Err(_) = verification_key.verify(message, signature) else {
                    return Ok(());
                };

                let verification_key = ApHttpVerificationKey::RsaSha512(Cow::Owned(
                    rsa::pkcs1v15::VerifyingKey::new(key.clone().into_owned()),
                ));

                Ok(verification_key.verify(message, signature)?)
            }
            (Self::Rsa(key), SigningAlgorithm::RsaSha256) => {
                let verification_key = ApHttpVerificationKey::RsaSha256(Cow::Owned(
                    rsa::pkcs1v15::VerifyingKey::new(key.clone().into_owned()),
                ));
                Ok(verification_key.verify(message, signature)?)
            }
            (_, SigningAlgorithm::RsaSha256) => {
                return Err(ApVerificationError::KeyAlgorithmMismatch(
                    algorithm,
                    self.as_ref().to_owned(),
                ));
            }
            (Self::Ed25519(key), SigningAlgorithm::Hs2019) => {
                let verification_key = ApHttpVerificationKey::Ed25519(Cow::Borrowed(key.as_ref()));
                Ok(verification_key.verify(message, signature)?)
            }
        }
    }
}

#[derive(Debug, Clone, AsRefStr)]
pub enum ApHttpPrivateKey<'a> {
    #[strum(serialize = "rsa")]
    Rsa(Box<Cow<'a, rsa::RsaPrivateKey>>),
    #[strum(serialize = "ed25519")]
    Ed25519(Cow<'a, ed25519_dalek::SecretKey>),
}

#[derive(Debug, Clone, AsRefStr)]
pub enum ApHttpSigningKey<'a> {
    #[strum(serialize = "rsa-sha256")]
    RsaSha256(Cow<'a, rsa::pkcs1v15::SigningKey<rsa::sha2::Sha256>>),
    #[strum(serialize = "rsa-sha512")]
    RsaSha512(Cow<'a, rsa::pkcs1v15::SigningKey<rsa::sha2::Sha512>>),
    #[strum(serialize = "ed25519")]
    Ed25519(Cow<'a, ed25519_dalek::SigningKey>),
}

#[derive(Debug, Clone)]
pub struct ApSigningKey<'a> {
    pub key: ApHttpSigningKey<'a>,
    pub key_id: Cow<'a, str>,
}

impl ApHttpSigningKey<'_> {
    pub fn sign(
        &self,
        algorithm: SigningAlgorithm,
        message: &[u8],
    ) -> Result<Vec<u8>, ApSigningError> {
        match (self, algorithm) {
            (Self::RsaSha256(key), SigningAlgorithm::RsaSha256 | SigningAlgorithm::Hs2019) => {
                Ok(Box::<[u8]>::from(key.sign(message)).into_vec())
            }
            (Self::RsaSha512(key), SigningAlgorithm::Hs2019) => {
                Ok(Box::<[u8]>::from(key.sign(message)).into_vec())
            }
            (Self::Ed25519(key), SigningAlgorithm::Hs2019) => {
                Ok(key.sign(message).to_bytes().to_vec())
            }
            (key, _) => Err(ApSigningError::KeyAlgorithmMismatch(
                algorithm,
                key.as_ref().to_string(),
            )),
        }
    }

    pub fn sign_base64(
        &self,
        algorithm: SigningAlgorithm,
        message: &[u8],
    ) -> Result<String, ApSigningError> {
        let signed = self.sign(algorithm, message)?;
        use base64::prelude::*;
        Ok(BASE64_STANDARD.encode(signed))
    }
}

impl ApHttpPrivateKey<'_> {
    pub fn create_signing_key<'a>(
        &'a self,
        key_id: &'a str,
        algorithm: SigningAlgorithm,
    ) -> Result<ApSigningKey<'a>, ApSigningError> {
        Ok(ApSigningKey {
            key_id: Cow::Borrowed(key_id),
            key: match (self, algorithm) {
                (Self::Rsa(key), SigningAlgorithm::RsaSha256 | SigningAlgorithm::Hs2019) => {
                    ApHttpSigningKey::RsaSha256(Cow::Owned(rsa::pkcs1v15::SigningKey::new(
                        key.clone().into_owned(),
                    )))
                }
                (Self::Ed25519(key), SigningAlgorithm::Hs2019) => ApHttpSigningKey::Ed25519(
                    Cow::Owned(ed25519_dalek::SigningKey::from_bytes(key.as_ref())),
                ),
                (key, _) => {
                    return Err(ApSigningError::KeyAlgorithmMismatch(
                        algorithm,
                        key.as_ref().to_owned(),
                    ))
                }
            },
        })
    }
}
