use crate::RingspaceError;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD as BASE64URL};
use hex;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

/// JSON Web Key structure for Ed25519 and X25519 keys
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Jwk {
    pub kty: String,
    pub crv: String,
    pub x: String,
    pub kid: String,
    pub alg: String,
    pub r#use: String,
}

/// JSON Web Key Set containing multiple keys
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Jwks {
    pub keys: Vec<Jwk>,
}

impl Jwks {
    /// Extract signing key from array, returning the first
    /// key with a `sig` use.
    pub fn get_signing_key(&self) -> Option<&Jwk> {
        self.keys.iter().filter(|&k| k.r#use == "sig").nth(0)
    }
    /// Extract encryption key from array, returning the first key with
    /// "enc" as a `use` value
    pub fn get_encryption_key(&self) -> Option<&Jwk> {
        self.keys.iter().filter(|&k| k.r#use == "enc").nth(0)
    }
}

impl Jwk {
    /// Creates a JWK from ED25519 signing keys.
    /// The `kid` is a [sha2::Sha256] hash, hex-encoded.
    pub fn new_ed25519(public_signing_key: &[u8]) -> Self {
        let kid = hex::encode(Sha256::digest(public_signing_key));
        Jwk {
            kty: "OKP".to_string(), // Octet Key Pair
            crv: "Ed25519".to_string(),
            x: BASE64URL.encode(public_signing_key),
            kid,
            alg: "EdDSA".to_string(),
            r#use: "sig".to_string(),
        }
    }

    /// Creates a JWK from X25519 encryption keys.
    /// The `kid` is a [sha2::Sha256] hash, hex-encoded.
    pub fn new_x25519(public_encryption_key: &[u8]) -> Self {
        let kid = hex::encode(Sha256::digest(public_encryption_key));
        Jwk {
            kty: "OKP".to_string(),
            crv: "X25519".to_string(),
            x: BASE64URL.encode(public_encryption_key),
            kid,
            alg: "ECDH-ES".to_string(),
            r#use: "enc".to_string(),
        }
    }

    /// Extract public key bytes from the JWK
    pub fn get_public_key(&self) -> Result<Vec<u8>, String> {
        BASE64URL
            .decode(&self.x)
            .map_err(|e| format!("Failed to decode public key: {}", e))
    }
}

/// Convert a base64-encoded Ed25519 public key to JWK format
pub fn ed25519_public_key_to_jwk(public_signing_key_b64: &str) -> Result<Jwk, RingspaceError> {
    let public_key = BASE64URL
        .decode(public_signing_key_b64)
        .map_err(|e| RingspaceError::new(format!("Failed to decode public signing key: {}", e)))?;

    if public_key.len() != 32 {
        return Err(RingspaceError::from("Invalid Ed25519 public key length"));
    }

    Ok(Jwk::new_ed25519(&public_key))
}

/// Convert a base64-encoded X25519 public key to JWK format
pub fn x25519_public_key_to_jwk(public_encryption_key_b64: &str) -> Result<Jwk, RingspaceError> {
    let public_key = BASE64URL.decode(public_encryption_key_b64).map_err(|e| {
        RingspaceError::new(format!("Failed to decode public encryption key: {}", e))
    })?;

    if public_key.len() != 32 {
        return Err(RingspaceError::from("Invalid X25519 public key length"));
    }

    Ok(Jwk::new_x25519(&public_key))
}
