//! Spectra's identities and composite signatures (spec §8.0).

//!

//! An identity is a composite key pair, Ed25519 and ML-DSA-65, with a name,

//! the identity it succeeds and the keys it delegates to. Both halves sign

//! the same target and both must verify, so a signature stays sound while

//! either algorithm does. The design follows the lessons of Oxygen's

//! identities; the code, encoding and files are Spectra's own.

use crate::cbor::{self, Value};

use ed25519_dalek::ed25519::signature::Signer as _;

use ed25519_dalek::{Signature as EdSignature, SigningKey as EdSigningKey, VerifyingKey as EdVerifyingKey};

use ml_dsa::{

    EncodedSignature, EncodedVerifyingKey, MlDsa65, Seed, Signature as MlSignature, SigningKey as MlSigningKey,

    VerifyingKey as MlVerifyingKey,

};

use rand::RngCore;

const VERSION: u64 = 1;

/// An identity's name: `BLAKE3("spectra1 identity " || encoding)`.

#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]

pub struct KeyId(pub [u8; 32]);

impl KeyId {

    pub fn to_hex(&self) -> String {

        hex::encode(self.0)

    }

    pub fn from_hex(s: &str) -> Option<KeyId> {

        hex::decode(s.trim()).ok()?.try_into().ok().map(KeyId)

    }

    /// The first `n` hex digits, for display.

    pub fn short(&self, n: usize) -> String {

        let h = self.to_hex();

        h[..n.min(h.len())].to_string()

    }

}

impl std::fmt::Debug for KeyId {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        write!(f, "KeyId({})", self.short(12))

    }

}

impl std::fmt::Display for KeyId {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.write_str(&self.to_hex())

    }

}

/// A key an identity lets act for it, until `expires` if set (ms).

#[derive(Clone, Debug, PartialEq, Eq)]

pub struct Delegate {

    pub key: KeyId,

    pub expires: Option<u64>,

}

impl Delegate {

    pub fn allows(&self, at_ms: u64) -> bool {

        self.expires.is_none_or(|e| at_ms < e)

    }

}

#[derive(Clone, Debug, PartialEq, Eq)]

pub struct Identity {

    pub ed25519_pk: [u8; 32],

    pub mldsa65_pk: Vec<u8>,

    pub name: String,

    pub contact: String,

    pub created: u64,

    pub expires: Option<u64>,

    pub previous: Option<KeyId>,

    pub delegates: Vec<Delegate>,

}

/// Why an identity, a secret or a signature could not be read or checked.

#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]

pub enum Error {

    #[error(transparent)]

    Cbor(#[from] cbor::Error),

    #[error("{0}")]

    Shape(String),

    #[error("the signature is by another key")]

    KeyMismatch,

    #[error("a public key that does not decode")]

    BadKey,

    #[error("the Ed25519 half does not verify")]

    Ed25519,

    #[error("the ML-DSA-65 half does not verify")]

    MlDsa,

}

fn shape(s: &str) -> Error {

    Error::Shape(s.into())

}

fn bytes_of(v: Option<&Value>, what: &str) -> Result<Vec<u8>, Error> {

    v.and_then(Value::as_bytes).map(|b| b.to_vec()).ok_or_else(|| shape(what))

}

fn key_of(v: Option<&Value>, what: &str) -> Result<KeyId, Error> {

    bytes_of(v, what)?.try_into().map(KeyId).map_err(|_| shape(what))

}

fn opt_uint(v: Option<&Value>, what: &str) -> Result<Option<u64>, Error> {

    match v {

        None | Some(Value::Null) => Ok(None),

        Some(Value::Uint(u)) => Ok(Some(*u)),

        _ => Err(shape(what)),

    }

}

fn opt_val<T>(o: Option<T>, f: impl Fn(T) -> Value) -> Value {

    o.map(f).unwrap_or(Value::Null)

}

impl Identity {

    pub fn to_value(&self) -> Value {

        Value::map([

            (0, Value::Uint(VERSION)),

            (1, Value::Bytes(self.ed25519_pk.to_vec())),

            (2, Value::Bytes(self.mldsa65_pk.clone())),

            (3, Value::Text(self.name.clone())),

            (4, Value::Text(self.contact.clone())),

            (5, Value::Uint(self.created)),

            (6, opt_val(self.expires, Value::Uint)),

            (7, opt_val(self.previous, |k| Value::Bytes(k.0.to_vec()))),

            (

                8,

                Value::Array(

                    self.delegates

                        .iter()

                        .map(|d| Value::map([(0, Value::Bytes(d.key.0.to_vec())), (1, opt_val(d.expires, Value::Uint))]))

                        .collect(),

                ),

            ),

        ])

    }

    pub fn encode(&self) -> Vec<u8> {

        cbor::encode(&self.to_value())

    }

    pub fn hash(&self) -> KeyId {

        let mut h = blake3::Hasher::new();

        h.update(b"spectra1 identity ");

        h.update(&self.encode());

        KeyId(*h.finalize().as_bytes())

    }

    pub fn decode(b: &[u8]) -> Result<Identity, Error> {

        let v = cbor::decode(b)?;

        let m = v.as_map().ok_or_else(|| shape("an identity is a map"))?;

        if m.get(&0).and_then(Value::as_uint) != Some(VERSION) {

            return Err(shape("an identity of an unknown version"));

        }

        let text = |k: u64, what: &str| m.get(&k).and_then(Value::as_text).map(String::from).ok_or_else(|| shape(what));

        let delegates = m

            .get(&8)

            .and_then(Value::as_array)

            .ok_or_else(|| shape("delegates"))?

            .iter()

            .map(|d| {

                let dm = d.as_map().ok_or_else(|| shape("a delegate is a map"))?;

                Ok(Delegate { key: key_of(dm.get(&0), "a delegate's key")?, expires: opt_uint(dm.get(&1), "a delegate's expiry")? })

            })

            .collect::<Result<Vec<_>, Error>>()?;

        let id = Identity {

            ed25519_pk: bytes_of(m.get(&1), "the Ed25519 key")?.try_into().map_err(|_| shape("the Ed25519 key"))?,

            mldsa65_pk: bytes_of(m.get(&2), "the ML-DSA-65 key")?,

            name: text(3, "the name")?,

            contact: text(4, "the contact")?,

            created: m.get(&5).and_then(Value::as_uint).ok_or_else(|| shape("the creation time"))?,

            expires: opt_uint(m.get(&6), "the expiry")?,

            previous: match m.get(&7) {

                None | Some(Value::Null) => None,

                v => Some(key_of(v, "the previous identity")?),

            },

            delegates,

        };

        // One identity, one encoding: what was read must be what we write.

        if id.encode() != b {

            return Err(shape("an identity not in Spectra's own encoding"));

        }

        Ok(id)

    }

    pub fn delegate(&self, key: &KeyId) -> Option<&Delegate> {

        self.delegates.iter().find(|d| &d.key == key)

    }

    /// Whether this identity lets `key` act for it at `at_ms`.

    pub fn delegates_to(&self, key: &KeyId, at_ms: u64) -> bool {

        !self.is_expired_at(at_ms) && self.delegate(key).is_some_and(|d| d.allows(at_ms))

    }

    pub fn is_expired_at(&self, at_ms: u64) -> bool {

        self.expires.is_some_and(|e| at_ms >= e)

    }

    /// The next identity: the same key pair with these delegates.

    pub fn successor(&self, delegates: Vec<Delegate>, expires: Option<u64>, created: u64) -> Identity {

        Identity { delegates, expires, created, previous: Some(self.hash()), ..self.clone() }

    }

    fn ed_key(&self) -> Result<EdVerifyingKey, Error> {

        EdVerifyingKey::from_bytes(&self.ed25519_pk).map_err(|_| Error::BadKey)

    }

    fn ml_key(&self) -> Result<MlVerifyingKey<MlDsa65>, Error> {

        let enc = EncodedVerifyingKey::<MlDsa65>::try_from(self.mldsa65_pk.as_slice()).map_err(|_| Error::BadKey)?;

        Ok(MlVerifyingKey::decode(&enc))

    }

}

/// What both halves sign: the signer's KeyId and the message, under

/// Spectra's prefix, so a signature means one thing by one key.

fn target(key: &KeyId, message: &[u8]) -> [u8; 32] {

    let mut h = blake3::Hasher::new();

    h.update(b"spectra1 sign ");

    h.update(&key.0);

    h.update(message);

    *h.finalize().as_bytes()

}

#[derive(Clone, PartialEq, Eq)]

pub struct Signature {

    pub key: KeyId,

    pub ed: [u8; 64],

    pub ml: Vec<u8>,

}

impl std::fmt::Debug for Signature {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        write!(f, "Signature({:?})", self.key)

    }

}

impl Signature {

    pub fn to_value(&self) -> Value {

        Value::map([(0, Value::Bytes(self.key.0.to_vec())), (1, Value::Bytes(self.ed.to_vec())), (2, Value::Bytes(self.ml.clone()))])

    }

    pub fn from_value(v: &Value) -> Result<Signature, Error> {

        let m = v.as_map().ok_or_else(|| shape("a signature is a map"))?;

        Ok(Signature {

            key: key_of(m.get(&0), "the signature's key")?,

            ed: bytes_of(m.get(&1), "the Ed25519 signature")?.try_into().map_err(|_| shape("the Ed25519 signature"))?,

            ml: bytes_of(m.get(&2), "the ML-DSA-65 signature")?,

        })

    }

    pub fn encode(&self) -> Vec<u8> {

        cbor::encode(&self.to_value())

    }

    pub fn decode(b: &[u8]) -> Result<Signature, Error> {

        Signature::from_value(&cbor::decode(b)?)

    }

}

/// Check a signature by `identity` over `message`.

pub fn verify(identity: &Identity, sig: &Signature, message: &[u8]) -> Result<(), Error> {

    let key = identity.hash();

    if sig.key != key {

        return Err(Error::KeyMismatch);

    }

    let t = target(&key, message);

    identity.ed_key()?.verify_strict(&t, &EdSignature::from_bytes(&sig.ed)).map_err(|_| Error::Ed25519)?;

    let enc = EncodedSignature::<MlDsa65>::try_from(sig.ml.as_slice()).map_err(|_| Error::MlDsa)?;

    let ml = MlSignature::<MlDsa65>::decode(&enc).ok_or(Error::MlDsa)?;

    if !identity.ml_key()?.verify_with_context(&t, b"", &ml) {

        return Err(Error::MlDsa);

    }

    Ok(())

}

/// The seeds of a key pair; stored only in its identity home, 0600.

#[derive(Clone)]

pub struct SecretKey {

    ed_seed: [u8; 32],

    ml_seed: [u8; 32],

}

impl std::fmt::Debug for SecretKey {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.write_str("SecretKey(..)")

    }

}

impl SecretKey {

    pub fn generate() -> SecretKey {

        let (mut ed_seed, mut ml_seed) = ([0u8; 32], [0u8; 32]);

        rand::rngs::OsRng.fill_bytes(&mut ed_seed);

        rand::rngs::OsRng.fill_bytes(&mut ml_seed);

        SecretKey { ed_seed, ml_seed }

    }

    fn ed(&self) -> EdSigningKey {

        EdSigningKey::from_bytes(&self.ed_seed)

    }

    fn ml(&self) -> MlSigningKey<MlDsa65> {

        MlSigningKey::from_seed(&Seed::from(self.ml_seed))

    }

    /// The identity of this key pair, delegating to nobody.

    pub fn identity(&self, name: &str, contact: &str, created: u64) -> Identity {

        Identity {

            ed25519_pk: self.ed().verifying_key().to_bytes(),

            mldsa65_pk: self.ml().expanded_key().verifying_key().encode().as_slice().to_vec(),

            name: name.into(),

            contact: contact.into(),

            created,

            expires: None,

            previous: None,

            delegates: vec![],

        }

    }

    /// Sign `message` as the identity named `key` (this key pair's).

    pub fn sign(&self, key: &KeyId, message: &[u8]) -> Signature {

        let t = target(key, message);

        let ed = self.ed().sign(&t).to_bytes();

        let ml = self.ml().expanded_key().sign_deterministic(&t, b"").expect("an empty context is valid").encode().as_slice().to_vec();

        Signature { key: *key, ed, ml }

    }

    /// Whether this is the key pair of `identity`.

    pub fn matches(&self, identity: &Identity) -> bool {

        let mine = self.identity("", "", 0);

        mine.ed25519_pk == identity.ed25519_pk && mine.mldsa65_pk == identity.mldsa65_pk

    }

    pub fn encode(&self, key: &KeyId) -> Vec<u8> {

        cbor::encode(&Value::map([

            (0, Value::Uint(VERSION)),

            (1, Value::Bytes(key.0.to_vec())),

            (2, Value::Bytes(self.ed_seed.to_vec())),

            (3, Value::Bytes(self.ml_seed.to_vec())),

        ]))

    }

    pub fn decode(b: &[u8]) -> Result<(KeyId, SecretKey), Error> {

        let v = cbor::decode(b)?;

        let m = v.as_map().ok_or_else(|| shape("a secret key is a map"))?;

        if m.get(&0).and_then(Value::as_uint) != Some(VERSION) {

            return Err(shape("a secret key of an unknown version"));

        }

        let seed = |k: u64| -> Result<[u8; 32], Error> { bytes_of(m.get(&k), "a seed")?.try_into().map_err(|_| shape("a seed")) };

        Ok((key_of(m.get(&1), "the secret's key")?, SecretKey { ed_seed: seed(2)?, ml_seed: seed(3)? }))

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn a_signature_binds_the_key_and_the_message_and_needs_both_halves() {

        let sk = SecretKey::generate();

        let id = sk.identity("Andy", "a@example", 1);

        let key = id.hash();

        let sig = sk.sign(&key, b"spectra1 test hello");

        verify(&id, &sig, b"spectra1 test hello").unwrap();

        assert_eq!(verify(&id, &sig, b"spectra1 test other"), Err(Error::Ed25519));

        let other = SecretKey::generate().identity("B", "", 1);

        assert_eq!(verify(&other, &sig, b"spectra1 test hello"), Err(Error::KeyMismatch));

        let mut half = sig.clone();

        half.ml[0] ^= 1;

        assert!(verify(&id, &half, b"spectra1 test hello").is_err(), "a broken ML-DSA half fails");

        assert_eq!(Signature::decode(&sig.encode()).unwrap(), sig);

    }

    #[test]

    fn identities_round_trip_and_successors_chain() {

        let sk = SecretKey::generate();

        let id = sk.identity("Andy", "", 5);

        assert_eq!(Identity::decode(&id.encode()).unwrap(), id);

        let agent = SecretKey::generate().identity("claude", "", 5).hash();

        let next = id.successor(vec![Delegate { key: agent, expires: Some(100) }], None, 6);

        assert_eq!(next.previous, Some(id.hash()));

        assert_ne!(next.hash(), id.hash());

        assert!(next.delegates_to(&agent, 99));

        assert!(!next.delegates_to(&agent, 100), "expired");

        assert!(!id.delegates_to(&agent, 1));

        assert!(sk.matches(&next));

        let (k, back) = SecretKey::decode(&sk.encode(&next.hash())).unwrap();

        assert_eq!(k, next.hash());

        assert!(back.matches(&id));

    }

}