//! 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));
    }
}