use crate::{schema, UIDCError};
use microrm::prelude::*;
use ring::{
    rand::SecureRandom,
    signature::{Ed25519KeyPair, KeyPair},
};
use sha2::Digest;

use itertools::Itertools;

#[derive(Debug)]
pub enum KeyError {
    Plain(&'static str),
}

#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum HMacType {
    Sha256,
    Sha512,
}

impl HMacType {
    pub fn digest_width(&self) -> usize {
        match self {
            Self::Sha256 => 32,
            Self::Sha512 => 64,
        }
    }
}

#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum KeyType {
    HMac(HMacType),
    RSA2048,
    RSA4096,
    Ed25519,
}

impl std::fmt::Display for KeyType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        <Self as std::fmt::Debug>::fmt(self, f)
    }
}

pub const KEY_TYPE_NAMES: &[(&str, KeyType)] = &[
    ("hmac_sha256", KeyType::HMac(HMacType::Sha256)),
    ("hmac_sha512", KeyType::HMac(HMacType::Sha512)),
    ("rsa2048", KeyType::RSA2048),
    ("rsa4096", KeyType::RSA4096),
    ("ed25519", KeyType::Ed25519),
];

impl std::str::FromStr for KeyType {
    type Err = UIDCError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        for (spec, kty) in KEY_TYPE_NAMES {
            if s == *spec {
                return Ok(*kty);
            }
        }
        Err(UIDCError::AbortString(format!(
            "invalid keytype: must be one of {}",
            KEY_TYPE_NAMES.iter().map(|v| v.0).join(",")
        )))
    }
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
struct RsaPublicData<'l> {
    modulus: &'l [u8],
    exponent: &'l [u8],
}

impl schema::Key {
    pub fn into_jwk(self) -> jsonwebtoken::jwk::Jwk {
        use jsonwebtoken::jwk;

        match self.key_type.as_ref() {
            KeyType::RSA2048 | KeyType::RSA4096 => {
                let rpubdata: RsaPublicData =
                    bincode::deserialize(self.public_data.as_slice()).unwrap();

                jwk::Jwk {
                    common: jwk::CommonParameters {
                        public_key_use: Some(jwk::PublicKeyUse::Signature),
                        key_algorithm: Some(jwk::KeyAlgorithm::RS256),
                        key_id: Some(self.key_id),
                        ..Default::default()
                    },
                    algorithm: jwk::AlgorithmParameters::RSA(jwk::RSAKeyParameters {
                        key_type: jwk::RSAKeyType::RSA,
                        n: base64::encode(rpubdata.modulus),
                        e: base64::encode(rpubdata.exponent),
                    }),
                }
            }
            KeyType::Ed25519 => jwk::Jwk {
                common: jwk::CommonParameters {
                    public_key_use: Some(jwk::PublicKeyUse::Signature),
                    key_algorithm: Some(jwk::KeyAlgorithm::EdDSA),
                    key_id: Some(self.key_id),
                    ..Default::default()
                },
                algorithm: jwk::AlgorithmParameters::OctetKeyPair(jwk::OctetKeyPairParameters {
                    key_type: jwk::OctetKeyPairType::OctetKeyPair,
                    curve: jwk::EllipticCurve::Ed25519,
                    x: base64::encode(self.public_data),
                }),
            },
            _ => todo!(),
        }
    }
}

fn pubkey_id(data: &[u8]) -> String {
    let mut key_hasher = sha2::Sha256::new();
    key_hasher.update(data);
    let mut key_id = base64::encode(key_hasher.finalize());
    key_id.truncate(16);
    key_id
}

fn generate_rsa(realm: &schema::Realm, kty: KeyType, bits: usize) -> Result<String, UIDCError> {
    // ring does not generate RSA keypairs, so we need to shell out to openssl for this.
    let openssl_output = std::process::Command::new("sh")
        .arg("-c")
        .arg(format!(
            "openssl genpkey \
             -algorithm RSA \
             -pkeyopt rsa_keygen_pubexp:65537 \
             -pkeyopt rsa_keygen_bits:{bits} \
             -outform der"
        ))
        .output()
        .map_err(|_| UIDCError::Abort("couldn't invoke openssl"))?;

    let secret = openssl_output.stdout;

    let keypair = ring::signature::RsaKeyPair::from_der(&secret)
        .map_err(|_| UIDCError::Abort("couldn't parse PKCS#8 output from openssl"))?;
    let public = keypair.public_key();
    let key_id = pubkey_id(public.as_ref());
    let expiry = time::OffsetDateTime::now_utc() + time::Duration::days(730);

    let public_data = bincode::serialize(&RsaPublicData {
        modulus: public.modulus().big_endian_without_leading_zero(),
        exponent: public.exponent().big_endian_without_leading_zero(),
    })
    .unwrap();

    realm.keys.insert(schema::Key {
        key_id: key_id.clone(),
        key_type: kty.into(),
        key_state: schema::KeyState::Active.into(),
        public_data,
        secret_data: secret,
        expiry,
    })?;

    Ok(key_id)
}

pub fn generate_in(realm: &schema::Realm, kty: KeyType) -> Result<String, UIDCError> {
    let rng = ring::rand::SystemRandom::new();
    match kty {
        KeyType::HMac(hmty) => {
            let rng = ring::rand::SystemRandom::new();
            let mut key_id = [0u8; 16];
            rng.fill(&mut key_id)
                .map_err(|_| UIDCError::Abort("couldn't generate random values"))?;
            let key_id = pubkey_id(&key_id);

            let mut keydata = vec![0; hmty.digest_width()];
            rng.fill(keydata.as_mut_slice())
                .map_err(|_| UIDCError::Abort("couldn't generate random values"))?;

            let expiry = time::OffsetDateTime::now_utc() + time::Duration::days(730);
            realm.keys.insert(schema::Key {
                key_id: key_id.clone(),
                key_type: kty.into(),
                key_state: schema::KeyState::Active.into(),
                // no separate public data for HMAC keys
                public_data: vec![],
                secret_data: keydata.clone(),
                expiry,
            })?;
            Ok(key_id)
        }
        KeyType::RSA2048 => generate_rsa(realm, KeyType::RSA2048, 2048),
        KeyType::RSA4096 => generate_rsa(realm, KeyType::RSA4096, 4096),
        KeyType::Ed25519 => {
            let generated_keypair = Ed25519KeyPair::generate_pkcs8(&rng)
                .map_err(|_| KeyError::Plain("failed to generate key"))?;

            let keydata = generated_keypair.as_ref().to_owned();

            let loaded_key = Ed25519KeyPair::from_pkcs8(keydata.as_slice())
                .expect("couldn't load just-generated key");
            let pubkey = loaded_key.public_key();
            let key_id = pubkey_id(pubkey.as_ref());
            let expiry = time::OffsetDateTime::now_utc() + time::Duration::days(730);

            realm.keys.insert(schema::Key {
                key_id: key_id.clone(),
                key_type: kty.into(),
                key_state: schema::KeyState::Active.into(),
                public_data: pubkey.as_ref().into(),
                secret_data: keydata,
                expiry,
            })?;

            Ok(key_id)
        }
    }
}

pub fn list(realm: &schema::Realm) -> Result<(), UIDCError> {
    let keys = realm.keys.get()?;

    for key in keys {
        println!(
            "- [{}] {:?}, expires {}",
            key.key_id,
            key.key_type,
            key.expiry
                .format(&time::format_description::well_known::Rfc3339)
                .unwrap()
        );
    }
    Ok(())
}

pub fn remove(realm: &schema::Realm, key_id: &String) -> Result<(), UIDCError> {
    realm.keys.with(schema::Key::KeyId, key_id).delete()?;
    Ok(())
}