//! Issuer-key resolution, Ed25519 signing, and signature verification.

use crate::models::{DelegationToken, Violation};
use base64ct::{Base64UrlUnpadded, Encoding};
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use std::collections::HashMap;
use std::sync::RwLock;

/// Only signature algorithm accepted by the 0.2 wire format.
pub const SIGNATURE_ALG_ED25519: &str = "Ed25519";
const STAGE: &str = "verify_issuer";

/// Resolves issuer keys from trusted host configuration.
///
/// Implementations must never resolve keys from the request or token itself. Resolver
/// failures are authorization failures, not a reason to fall back to caller key material.
pub trait IssuerKeyResolver: Send + Sync {
    /// Resolves a verification key from trusted host state.
    ///
    /// `Ok(None)` means the issuer/key pair is not trusted. Backend failures should return an
    /// error; the evaluator treats both outcomes as denial.
    fn resolve(&self, issuer: &str, key_id: &str) -> Result<Option<VerifyingKey>, Violation>;
}

/// An in-process allowlist of pinned issuer keys.
#[derive(Default)]
pub struct PinnedIssuerKeys {
    keys: RwLock<HashMap<(String, String), VerifyingKey>>,
}

impl PinnedIssuerKeys {
    /// Creates an empty issuer-key allowlist.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds or replaces one trusted `(issuer, key_id)` mapping.
    pub fn insert(
        &self,
        issuer: impl Into<String>,
        key_id: impl Into<String>,
        key: VerifyingKey,
    ) -> Result<(), Violation> {
        self.keys
            .write()
            .map_err(|_| Violation::new(STAGE, "issuer key store lock poisoned"))?
            .insert((issuer.into(), key_id.into()), key);
        Ok(())
    }
}

impl IssuerKeyResolver for PinnedIssuerKeys {
    fn resolve(&self, issuer: &str, key_id: &str) -> Result<Option<VerifyingKey>, Violation> {
        Ok(self
            .keys
            .read()
            .map_err(|_| Violation::new(STAGE, "issuer key store lock poisoned"))?
            .get(&(issuer.to_owned(), key_id.to_owned()))
            .copied())
    }
}

/// Signs every authorization-relevant field of a capability token with an issuer key.
pub fn sign_token(
    token: &DelegationToken,
    issuer_signing_key: &SigningKey,
) -> Result<String, Violation> {
    let payload = token_signing_payload(token)?;
    Ok(Base64UrlUnpadded::encode_string(
        &issuer_signing_key.sign(&payload).to_bytes(),
    ))
}

/// Verifies a capability token against a key returned by trusted host configuration.
pub fn verify_token(
    token: &DelegationToken,
    resolver: &dyn IssuerKeyResolver,
) -> Result<(), Violation> {
    let key = resolver
        .resolve(&token.issuer, &token.key_id)?
        .ok_or_else(|| Violation::new(STAGE, "token issuer/key_id is not trusted"))?;
    let signature_bytes = Base64UrlUnpadded::decode_vec(&token.signature)
        .map_err(|_| Violation::new(STAGE, "token signature must be base64url-no-pad"))?;
    let signature_bytes: [u8; 64] = signature_bytes
        .try_into()
        .map_err(|_| Violation::new(STAGE, "token signature must decode to 64 bytes"))?;
    let signature = Signature::from_bytes(&signature_bytes);
    key.verify(&token_signing_payload(token)?, &signature)
        .map_err(|_| Violation::new(STAGE, "token signature verification failed"))
}

fn token_signing_payload(token: &DelegationToken) -> Result<Vec<u8>, Violation> {
    let mut unsigned = token.clone();
    unsigned.signature.clear();
    let value = serde_json::to_value(unsigned)
        .map_err(|error| Violation::new(STAGE, format!("token serialization failed: {error}")))?;
    Ok(canonical_json(&value))
}

fn canonical_json(value: &serde_json::Value) -> Vec<u8> {
    let mut output = Vec::new();
    write_canonical(&mut output, value);
    output
}

fn write_canonical(output: &mut Vec<u8>, value: &serde_json::Value) {
    match value {
        serde_json::Value::Object(map) => {
            output.push(b'{');
            let mut entries: Vec<_> = map.iter().collect();
            entries.sort_unstable_by_key(|(key, _)| *key);
            for (index, (key, value)) in entries.into_iter().enumerate() {
                if index > 0 {
                    output.push(b',');
                }
                output.extend_from_slice(
                    serde_json::to_string(key)
                        .expect("JSON object keys always serialize")
                        .as_bytes(),
                );
                output.push(b':');
                write_canonical(output, value);
            }
            output.push(b'}');
        }
        serde_json::Value::Array(values) => {
            output.push(b'[');
            for (index, value) in values.iter().enumerate() {
                if index > 0 {
                    output.push(b',');
                }
                write_canonical(output, value);
            }
            output.push(b']');
        }
        primitive => output.extend_from_slice(
            serde_json::to_string(primitive)
                .expect("JSON values always serialize")
                .as_bytes(),
        ),
    }
}
