use crate::error::KeyError;
use crate::{Algorithm, EllipticCurve, Jwk, Key, KeyMaterial, RsaPublicKey};
use aws_lc_rs::encoding::AsBigEndian;
use aws_lc_rs::signature::KeyPair;
use p256::elliptic_curve::array::typenum::Unsigned;
use p256::elliptic_curve::point::PointCompression;
use p256::elliptic_curve::sec1::{FromSec1Point, ModulusSize, ToSec1Point};
use p256::elliptic_curve::{Curve, CurveArithmetic, SecretKey};
use rsa::traits::{PrivateKeyParts, PublicKeyParts};

/// Generate a key pair for the given algorithm, returning the private and public keys.
pub fn generate(algorithm: Algorithm, id: Option<crate::KeyId>) -> crate::Result<Key> {
	let key = match algorithm {
		Algorithm::HS256 => generate_hmac_key::<32>(),
		Algorithm::HS384 => generate_hmac_key::<48>(),
		Algorithm::HS512 => generate_hmac_key::<64>(),
		Algorithm::RS256
		| Algorithm::RS384
		| Algorithm::RS512
		| Algorithm::PS256
		| Algorithm::PS384
		| Algorithm::PS512 => generate_rsa_key(2048),
		Algorithm::ES256 => generate_ec_key::<p256::NistP256>(EllipticCurve::P256),
		Algorithm::ES384 => generate_ec_key::<p384::NistP384>(EllipticCurve::P384),
		Algorithm::EdDSA => generate_ed25519_key(),
	};

	let mut jwk = Jwk::new(algorithm, key?);
	jwk.kid = id;
	jwk.try_into()
}

fn generate_hmac_key<const SIZE: usize>() -> crate::Result<KeyMaterial> {
	let mut key = [0u8; SIZE];
	aws_lc_rs::rand::fill(&mut key)?;
	Ok(KeyMaterial::OCT { secret: key.to_vec() })
}

struct AwsRng;

impl rsa::rand_core::RngCore for AwsRng {
	fn next_u32(&mut self) -> u32 {
		let mut buf = [0u8; 4];
		self.fill_bytes(&mut buf);
		u32::from_le_bytes(buf)
	}

	fn next_u64(&mut self) -> u64 {
		let mut buf = [0u8; 8];
		self.fill_bytes(&mut buf);
		u64::from_le_bytes(buf)
	}

	fn fill_bytes(&mut self, dest: &mut [u8]) {
		aws_lc_rs::rand::fill(dest).unwrap();
	}

	fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rsa::rand_core::Error> {
		aws_lc_rs::rand::fill(dest).map_err(|_| rsa::rand_core::Error::new("aws-lc-rs failed"))
	}
}

impl rsa::rand_core::CryptoRng for AwsRng {}

fn generate_rsa_key(size: usize) -> crate::Result<KeyMaterial> {
	let mut rng = AwsRng;
	let mut key = rsa::RsaPrivateKey::new(&mut rng, size)?;
	key.precompute()?;

	Ok(KeyMaterial::RSA {
		public: RsaPublicKey {
			e: key.e().to_bytes_be(),
			n: key.n().to_bytes_be(),
		},
		private: Some(crate::RsaPrivateKey {
			d: key.d().to_bytes_be(),
			p: key.primes()[0].to_bytes_be(),
			q: key.primes()[1].to_bytes_be(),
			dp: key.dp().ok_or(KeyError::MissingPrivateKey)?.to_bytes_be(),
			dq: key.dq().ok_or(KeyError::MissingPrivateKey)?.to_bytes_be(),
			qi: key.qinv().ok_or(KeyError::MissingPrivateKey)?.to_bytes_be().1,
			oth: None, // TODO https://datatracker.ietf.org/doc/html/rfc7518#section-6.3.2.7
		}),
	})
}

fn generate_ec_key<C>(curve: EllipticCurve) -> crate::Result<KeyMaterial>
where
	C: Curve + CurveArithmetic + PointCompression,
	C::AffinePoint: ToSec1Point<C> + FromSec1Point<C>,
	C::FieldBytesSize: ModulusSize,
{
	let mut bytes = vec![0u8; C::FieldBytesSize::to_usize()];
	let secret = loop {
		aws_lc_rs::rand::fill(&mut bytes)?;

		if let Ok(secret) = SecretKey::<C>::from_slice(&bytes) {
			break secret;
		}
	};

	let point = secret.public_key().to_sec1_point(false);

	let x = point.x().ok_or(KeyError::MissingEcX)?.to_vec();
	let y = point.y().ok_or(KeyError::MissingEcY)?.to_vec();
	let d = secret.to_bytes().to_vec();

	Ok(KeyMaterial::EC {
		curve,
		x,
		y,
		d: Some(d),
	})
}

fn generate_ed25519_key() -> crate::Result<KeyMaterial> {
	let key_pair = aws_lc_rs::signature::Ed25519KeyPair::generate()?;

	let public_key = key_pair.public_key().as_ref().to_vec();
	let seed = key_pair.seed()?.as_be_bytes()?.as_ref().to_vec();

	Ok(KeyMaterial::OKP {
		curve: EllipticCurve::Ed25519,
		x: public_key,
		d: Some(seed),
	})
}
