//! Identities and composite Ed25519 + ML-DSA-65 signatures (spec §14.2, §14.3).
//!
//! Both halves sign the same message, which binds the composite `KeyId`, and
//! verification requires both. A signature with either half missing, wrong,
//! or made for a different KeyId is invalid.

use crate::cbor::{self, Value};
use crate::hash::{hash_object, KeyId, ObjectRef, ObjectType};
use crate::objects::{
    array, bytes, check_version, field, oref, text, uint, ObjectError, FORMAT_VERSION,
};
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;

pub const ED25519_PK_LEN: usize = 32;
pub const MLDSA65_PK_LEN: usize = 1952;
pub const ED25519_SIG_LEN: usize = 64;
pub const MLDSA65_SIG_LEN: usize = 3309;
/// `alg` value for Ed25519 + ML-DSA-65.
pub const ALG_COMPOSITE: u64 = 1;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Identity {
    pub ed25519_pk: [u8; ED25519_PK_LEN],
    pub mldsa65_pk: Vec<u8>,
    pub name: String,
    pub email: String,
    pub created: u64,
    pub expires: Option<u64>,
    pub previous: Option<KeyId>,
    pub delegates: Vec<KeyId>,
}

impl Identity {
    pub fn to_value(&self) -> Value {
        Value::map([
            (0, Value::Uint(FORMAT_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.email.clone())),
            (5, Value::Uint(self.created)),
            (6, self.expires.map(Value::Uint).unwrap_or(Value::Null)),
            (
                7,
                self.previous
                    .map(|p| Value::Bytes(p.0.to_vec()))
                    .unwrap_or(Value::Null),
            ),
            (
                8,
                Value::Array(
                    self.delegates
                        .iter()
                        .map(|d| Value::Bytes(d.0.to_vec()))
                        .collect(),
                ),
            ),
        ])
    }
    pub fn encode(&self) -> Vec<u8> {
        cbor::encode(&self.to_value())
    }
    pub fn hash(&self) -> KeyId {
        hash_object(ObjectType::Identity, &self.encode())
    }
    pub fn decode(b: &[u8]) -> Result<Identity, ObjectError> {
        let v = cbor::decode(b)?;
        let m = v
            .as_map()
            .ok_or_else(|| ObjectError::Schema("identity: expected map".into()))?;
        check_version(m)?;
        let ed = bytes(field(m, 1, "ed25519_pk")?, "ed25519_pk")?;
        let ml = bytes(field(m, 2, "mldsa65_pk")?, "mldsa65_pk")?;
        if ed.len() != ED25519_PK_LEN || ml.len() != MLDSA65_PK_LEN {
            return Err(ObjectError::Schema(
                "identity: wrong public key length".into(),
            ));
        }
        let opt_u = |v: &Value| -> Result<Option<u64>, ObjectError> {
            if v.is_null() {
                Ok(None)
            } else {
                uint(v, "expires").map(Some)
            }
        };
        let opt_r = |v: &Value| -> Result<Option<KeyId>, ObjectError> {
            if v.is_null() {
                Ok(None)
            } else {
                oref(v, "previous").map(Some)
            }
        };
        let mut delegates = array(field(m, 8, "delegates")?, "delegates")?
            .iter()
            .map(|d| oref(d, "delegate"))
            .collect::<Result<Vec<_>, _>>()?;
        delegates.sort();
        delegates.dedup();
        Ok(Identity {
            ed25519_pk: ed.try_into().unwrap(),
            mldsa65_pk: ml.to_vec(),
            name: text(field(m, 3, "name")?, "name")?.to_owned(),
            email: text(field(m, 4, "email")?, "email")?.to_owned(),
            created: uint(field(m, 5, "created")?, "created")?,
            expires: opt_u(field(m, 6, "expires")?)?,
            previous: opt_r(field(m, 7, "previous")?)?,
            delegates,
        })
    }

    /// Whether the identity is past its expiry at `now_ms`.
    pub fn is_expired_at(&self, now_ms: u64) -> bool {
        self.expires.is_some_and(|e| now_ms > e)
    }

    /// A successor identity with the same keys, `previous` set to this
    /// identity's KeyId, and the given delegates and expiry (spec §14.2).
    pub fn successor(&self, delegates: Vec<KeyId>, expires: Option<u64>, created: u64) -> Identity {
        let mut delegates = delegates;
        delegates.sort();
        delegates.dedup();
        Identity {
            previous: Some(self.hash()),
            delegates,
            expires,
            created,
            ..self.clone()
        }
    }

    fn ed_key(&self) -> Result<EdVerifyingKey, VerifyError> {
        EdVerifyingKey::from_bytes(&self.ed25519_pk).map_err(|_| VerifyError::BadKey)
    }
    fn ml_key(&self) -> Result<MlVerifyingKey<MlDsa65>, VerifyError> {
        let enc = EncodedVerifyingKey::<MlDsa65>::try_from(self.mldsa65_pk.as_slice())
            .map_err(|_| VerifyError::BadKey)?;
        Ok(MlVerifyingKey::decode(&enc))
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Domain {
    Change,
    Attest,
    Policy,
    /// A server-issued challenge, proving possession of a key to a host
    /// (hub spec §5.3, §5.4). Never a stored object.
    Challenge,
}

impl Domain {
    pub fn name(self) -> &'static str {
        match self {
            Domain::Change => "change",
            Domain::Attest => "attest",
            Domain::Policy => "policy",
            Domain::Challenge => "challenge",
        }
    }
}

/// `"oxy1 sig " || domain || 0x00 || KeyId || target` (spec §14.3).
pub fn signing_message(domain: Domain, key: &KeyId, target: &ObjectRef) -> Vec<u8> {
    let mut m = Vec::with_capacity(80);
    m.extend_from_slice(b"oxy1 sig ");
    m.extend_from_slice(domain.name().as_bytes());
    m.push(0);
    m.extend_from_slice(&key.0);
    m.extend_from_slice(&target.0);
    m
}

#[derive(Clone, PartialEq, Eq)]
pub struct Signature {
    pub key: KeyId,
    pub ed: [u8; ED25519_SIG_LEN],
    pub ml: Vec<u8>,
}

impl std::fmt::Debug for Signature {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Signature(key {:?})", self.key)
    }
}

impl Signature {
    pub fn to_value(&self) -> Value {
        Value::map([
            (0, Value::Bytes(self.key.0.to_vec())),
            (1, Value::Uint(ALG_COMPOSITE)),
            (2, Value::Bytes(self.ed.to_vec())),
            (3, Value::Bytes(self.ml.clone())),
        ])
    }
    pub fn from_value(v: &Value) -> Result<Signature, ObjectError> {
        let m = v
            .as_map()
            .ok_or_else(|| ObjectError::Schema("signature: expected map".into()))?;
        let alg = uint(field(m, 1, "alg")?, "alg")?;
        if alg != ALG_COMPOSITE {
            return Err(ObjectError::Schema(format!(
                "unknown signature algorithm {alg}"
            )));
        }
        let ed = bytes(field(m, 2, "ed")?, "ed")?;
        let ml = bytes(field(m, 3, "ml")?, "ml")?;
        if ed.len() != ED25519_SIG_LEN || ml.len() != MLDSA65_SIG_LEN {
            return Err(ObjectError::Schema("signature: wrong length".into()));
        }
        Ok(Signature {
            key: oref(field(m, 0, "key")?, "key")?,
            ed: ed.try_into().unwrap(),
            ml: ml.to_vec(),
        })
    }
    pub fn encode_list(sigs: &[Signature]) -> Vec<u8> {
        cbor::encode(&Value::Array(sigs.iter().map(|s| s.to_value()).collect()))
    }
    pub fn decode_list(b: &[u8]) -> Result<Vec<Signature>, ObjectError> {
        let v = cbor::decode(b)?;
        array(&v, "signatures")?
            .iter()
            .map(Signature::from_value)
            .collect()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum VerifyError {
    #[error("signature names a different key than the identity")]
    KeyMismatch,
    #[error("identity holds a malformed public key")]
    BadKey,
    #[error("Ed25519 signature invalid")]
    Ed25519,
    #[error("ML-DSA-65 signature invalid")]
    MlDsa,
    #[error("a challenge is at least 16 bytes")]
    ShortChallenge,
}

/// Verify a composite signature by `identity` over `target` (spec §14.3).
pub fn verify(
    identity: &Identity,
    sig: &Signature,
    domain: Domain,
    target: &ObjectRef,
) -> Result<(), VerifyError> {
    let key = identity.hash();
    if sig.key != key {
        return Err(VerifyError::KeyMismatch);
    }
    let msg = signing_message(domain, &key, target);
    let ed_sig = EdSignature::from_bytes(&sig.ed);
    identity
        .ed_key()?
        .verify_strict(&msg, &ed_sig)
        .map_err(|_| VerifyError::Ed25519)?;
    let enc =
        EncodedSignature::<MlDsa65>::try_from(sig.ml.as_slice()).map_err(|_| VerifyError::MlDsa)?;
    let ml_sig = MlSignature::<MlDsa65>::decode(&enc).ok_or(VerifyError::MlDsa)?;
    if !identity.ml_key()?.verify_with_context(&msg, b"", &ml_sig) {
        return Err(VerifyError::MlDsa);
    }
    Ok(())
}

/// The seeds of a composite key pair. Stored only in the user's key directory.
#[derive(Clone)]
pub struct SecretKey {
    pub ed_seed: [u8; 32],
    pub 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 = [0u8; 32];
        let mut ml_seed = [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 object for these keys.
    pub fn identity(&self, name: &str, email: &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(),
            email: email.into(),
            created,
            expires: None,
            previous: None,
            delegates: vec![],
        }
    }

    pub fn sign(&self, domain: Domain, key: &KeyId, target: &ObjectRef) -> Signature {
        let msg = signing_message(domain, key, target);
        let ed = self.ed().sign(&msg).to_bytes();
        let ml = self
            .ml()
            .expanded_key()
            .sign_deterministic(&msg, b"")
            .expect("empty context is valid")
            .encode()
            .as_slice()
            .to_vec();
        Signature { key: *key, ed, ml }
    }

    pub fn encode(&self, key: &KeyId) -> Vec<u8> {
        cbor::encode(&Value::map([
            (0, Value::Uint(FORMAT_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), ObjectError> {
        let v = cbor::decode(b)?;
        let m = v
            .as_map()
            .ok_or_else(|| ObjectError::Schema("secret key: expected map".into()))?;
        check_version(m)?;
        let seed = |k: u64| -> Result<[u8; 32], ObjectError> {
            bytes(field(m, k, "seed")?, "seed")?
                .try_into()
                .map_err(|_| ObjectError::Schema("seed: expected 32 bytes".into()))
        };
        Ok((
            oref(field(m, 1, "key")?, "key")?,
            SecretKey {
                ed_seed: seed(2)?,
                ml_seed: seed(3)?,
            },
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn fixed_key() -> SecretKey {
        SecretKey {
            ed_seed: [1; 32],
            ml_seed: [2; 32],
        }
    }

    #[test]
    fn identity_round_trip_and_sizes() {
        let sk = fixed_key();
        let id = sk.identity("Ada", "ada@example.org", 1_700_000_000_000);
        assert_eq!(id.mldsa65_pk.len(), MLDSA65_PK_LEN);
        let d = Identity::decode(&id.encode()).unwrap();
        assert_eq!(d, id);
        assert_eq!(d.hash(), id.hash());
        // Deterministic from seeds.
        assert_eq!(
            fixed_key()
                .identity("Ada", "ada@example.org", 1_700_000_000_000)
                .hash(),
            id.hash()
        );
        let mut other = id.clone();
        other.delegates.push(ObjectRef([9; 32]));
        assert_ne!(other.hash(), id.hash());
    }

    #[test]
    fn sign_and_verify() {
        let sk = fixed_key();
        let id = sk.identity("Ada", "ada@example.org", 1);
        let key = id.hash();
        let target = ObjectRef([7; 32]);
        let sig = sk.sign(Domain::Change, &key, &target);
        assert_eq!(sig.ml.len(), MLDSA65_SIG_LEN);
        assert_eq!(verify(&id, &sig, Domain::Change, &target), Ok(()));
        // Different target, different domain: invalid.
        assert_eq!(
            verify(&id, &sig, Domain::Change, &ObjectRef([8; 32])),
            Err(VerifyError::Ed25519)
        );
        assert_eq!(
            verify(&id, &sig, Domain::Attest, &target),
            Err(VerifyError::Ed25519)
        );
        // Deterministic signatures.
        assert_eq!(sk.sign(Domain::Change, &key, &target), sig);
        // Signature list round trip.
        let l =
            Signature::decode_list(&Signature::encode_list(std::slice::from_ref(&sig))).unwrap();
        assert_eq!(l, vec![sig]);
    }

    #[test]
    fn stripping_or_forging_either_half_fails() {
        let sk = fixed_key();
        let id = sk.identity("Ada", "a@b", 1);
        let key = id.hash();
        let target = ObjectRef([7; 32]);
        let good = sk.sign(Domain::Change, &key, &target);

        let mut no_ml = good.clone();
        no_ml.ml = vec![0; MLDSA65_SIG_LEN];
        assert_eq!(
            verify(&id, &no_ml, Domain::Change, &target),
            Err(VerifyError::MlDsa)
        );

        let mut bad_ed = good.clone();
        bad_ed.ed[10] ^= 1;
        assert_eq!(
            verify(&id, &bad_ed, Domain::Change, &target),
            Err(VerifyError::Ed25519)
        );

        // A signature made for another KeyId over the same target does not transfer.
        let other = SecretKey {
            ed_seed: [3; 32],
            ml_seed: [4; 32],
        };
        let other_id = other.identity("Bob", "b@b", 1);
        let theirs = other.sign(Domain::Change, &other_id.hash(), &target);
        assert_eq!(
            verify(&id, &theirs, Domain::Change, &target),
            Err(VerifyError::KeyMismatch)
        );
        let mut relabelled = theirs.clone();
        relabelled.key = key;
        assert_eq!(
            verify(&id, &relabelled, Domain::Change, &target),
            Err(VerifyError::Ed25519)
        );

        // Wrong length is rejected at decode time.
        let mut short = good.to_value();
        if let Value::Map(m) = &mut short {
            m.insert(3, Value::Bytes(vec![0; 10]));
        }
        assert!(Signature::from_value(&short).is_err());
    }

    #[test]
    fn secret_key_round_trip() {
        let sk = fixed_key();
        let key = sk.identity("x", "y", 0).hash();
        let (k, d) = SecretKey::decode(&sk.encode(&key)).unwrap();
        assert_eq!(k, key);
        assert_eq!(d.ed_seed, sk.ed_seed);
        assert_eq!(d.ml_seed, sk.ml_seed);
        assert_ne!(SecretKey::generate().ed_seed, SecretKey::generate().ed_seed);
    }
}

/// The signing target of a challenge: `BLAKE3-256("oxy1 challenge " || bytes)`,
/// so a challenge can never collide with a stored object's reference.
pub fn challenge_target(challenge: &[u8]) -> ObjectRef {
    let mut h = blake3::Hasher::new();
    h.update(b"oxy1 challenge ");
    h.update(challenge);
    ObjectRef(*h.finalize().as_bytes())
}

impl SecretKey {
    /// Sign a host's challenge (at least 16 bytes) with this key.
    pub fn sign_challenge(&self, key: &KeyId, challenge: &[u8]) -> Result<Signature, VerifyError> {
        if challenge.len() < 16 {
            return Err(VerifyError::ShortChallenge);
        }
        Ok(self.sign(Domain::Challenge, key, &challenge_target(challenge)))
    }
}

/// Check a challenge signature against the identity that claims it.
pub fn verify_challenge(
    identity: &Identity,
    sig: &Signature,
    challenge: &[u8],
) -> Result<(), VerifyError> {
    if challenge.len() < 16 {
        return Err(VerifyError::ShortChallenge);
    }
    if sig.key != identity.hash() {
        return Err(VerifyError::KeyMismatch);
    }
    verify(
        identity,
        sig,
        Domain::Challenge,
        &challenge_target(challenge),
    )
}

#[cfg(test)]
mod challenge_tests {
    use super::*;

    #[test]
    fn a_challenge_signature_binds_key_and_bytes() {
        let sk = SecretKey::generate();
        let id = sk.identity("A", "a@example", 1);
        let key = id.hash();
        let challenge = [7u8; 32];
        let sig = sk.sign_challenge(&key, &challenge).unwrap();
        assert!(verify_challenge(&id, &sig, &challenge).is_ok());
        assert!(
            verify_challenge(&id, &sig, &[8u8; 32]).is_err(),
            "other bytes"
        );
        let other = SecretKey::generate().identity("B", "b@example", 1);
        assert!(
            verify_challenge(&other, &sig, &challenge).is_err(),
            "other identity"
        );
        assert!(sk.sign_challenge(&key, b"short").is_err());
        // A challenge target never equals a change signature's target for the same bytes.
        assert_ne!(challenge_target(&challenge), ObjectRef(challenge));
    }
}