//! ML-DSA-44 cosignatures (type 0x06) of a checkpoint and of any subtree of
//! the log (log-seal §5.6), held to the bytes an independent ML-DSA
//! implementation computes for the same seed and inputs.

use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as B64;
use ed25519_dalek::SigningKey;
use jcp_logseal::note::{
    checkpoint, cosign_checkpoint, cosignatures, cosigned_subtree, read_checkpoint, sha256_root, verify_checkpoint, verify_subtree_cosignature,
    MlDsaCosigner, NoteError, NoteSigner, NoteVerifier, MAX_TIMESTAMP, MLDSA44_KEY_LEN, MLDSA44_SIG_LEN, TYPE_MLDSA44_COSIGNATURE,
};
use jcp_logseal::subtree::{
    evaluate_subtree_inclusion, leaf_hash, subtree_consistency_proof, subtree_hash, subtree_inclusion_proof, verify_subtree_consistency, Hash,
};
use jcp_schema::ContentId;
use ml_dsa::{B32, ExpandedSigningKey, MlDsa44};
use sha2::{Digest, Sha256};

const NAME: &str = "example.com/witness";
const ORIGIN: &str = "example.com/behind-the-sofa";
const CHECKPOINT: &str = "example.com/behind-the-sofa\n20852163\nCsUYapGGPo4dkMgIAUqom/Xajj7h2fB2MPA3j2jxq2I=\n";
const TIME: u64 = 1_679_315_147;

/// The seed the independent implementation was given: the bytes 1 to 32.
fn seed() -> [u8; 32] {
    core::array::from_fn(|i| u8::try_from(i + 1).unwrap())
}

fn sha(b: &[u8]) -> String {
    hex::encode(Sha256::digest(b))
}

/// Every value here was computed by an ML-DSA-44 implementation that shares
/// no code with this one, from the same seed and inputs: the key, its id and
/// verifier key, what a checkpoint cosignature and a subtree cosignature
/// sign, and the deterministic signatures themselves.
#[test]
fn an_independent_implementation_computes_the_same_bytes() {
    let c = MlDsaCosigner::from_seed(NAME, seed()).unwrap();
    assert_eq!(c.public_key().len(), MLDSA44_KEY_LEN);
    assert_eq!(sha(c.public_key()), "cec606d352aba0650d219bc77c12af36b36f8089bee9720cb0cbf65e49c20dd8");
    assert_eq!(format!("{:08x}", c.id()), "760e9e80");
    assert_eq!(sha(c.vkey().as_bytes()), "b6e3d201f816630b7daca3c4abfd884c8ac36ca7e4fedab4f8509ed1b9ba8d2f");

    let (origin, count, root) = read_checkpoint(CHECKPOINT).unwrap();
    assert_eq!(sha(&cosigned_subtree(NAME, TIME, &origin, 0, count, &root).unwrap()), "270c03b41f0308d9b8cefac32d5df3ead70315646c798bbac769fa32013a0633");
    let line = c.cosign_checkpoint(CHECKPOINT, TIME).unwrap();
    assert_eq!(sha(line.as_bytes()), "32f00d16f493812ccf7f60bfa9a94144e3f7ed64274b7860b478591291c07887");
    let payload = B64.decode(line.trim_end().rsplit_once(' ').unwrap().1).unwrap();
    assert_eq!(payload.len(), 4 + 8 + MLDSA44_SIG_LEN);
    assert_eq!(&payload[4..12], &TIME.to_be_bytes());
    assert_eq!(sha(&payload[12..]), "1cf7ee7f47d09d0064431c45c5f100693835b7368f7473f8b6e531b20730b81c");

    let sub: Hash = Sha256::digest(b"subtree").into();
    assert_eq!(sha(&cosigned_subtree(NAME, 0, ORIGIN, 8, 12, &sub).unwrap()), "97291f7a8f64ca6a6f9e01c3a0a7f7fbf3a3bc8befec7622953866910a780ae4");
    assert_eq!(sha(&c.cosign_subtree(ORIGIN, 8, 12, &sub).unwrap()), "c77c2d25f7bb5fd814d1a01bdf1f7aa8a7d5647599ca7a8bdb44fda59056dc9e");

    // and it reads back: the spec's example checkpoint, cosigned
    let vk = NoteVerifier::from_vkey(&c.vkey()).unwrap();
    assert_eq!((vk.kind, vk.key.len(), vk.id), (TYPE_MLDSA44_COSIGNATURE, MLDSA44_KEY_LEN, 0x760e_9e80));
    assert_eq!(cosignatures(&format!("{CHECKPOINT}\n{line}"), &[vk]).unwrap(), vec![(NAME.to_owned(), TIME)]);
}

fn log(n: u32) -> Vec<Vec<u8>> {
    (0..n).map(|i| format!("receipt {i}").into_bytes()).collect()
}

fn leaves(log: &[Vec<u8>]) -> Vec<Hash> {
    log.iter().map(|e| leaf_hash(&ContentId::of_bytes(e).0)).collect()
}

fn vkey_of(c: &MlDsaCosigner) -> NoteVerifier {
    NoteVerifier::from_vkey(&c.vkey()).unwrap()
}

#[test]
fn a_checkpoint_cosigned_post_quantum_by_its_steward_and_witnesses_verifies() {
    let log = log(13);
    let text = checkpoint("jcp.example/steward", &log, 13).unwrap();
    let steward = MlDsaCosigner::from_seed("jcp.example/steward", [1; 32]).unwrap();
    let one = MlDsaCosigner::from_seed("witness.example/one", [2; 32]).unwrap();
    let two = NoteSigner::cosigner("witness.example/two", SigningKey::from_bytes(&[3; 32])).unwrap();
    let note = format!(
        "{text}\n{}{}{}",
        steward.cosign_checkpoint(&text, 100).unwrap(),
        one.cosign_checkpoint(&text, 101).unwrap(),
        cosign_checkpoint(&text, &two, 102).unwrap()
    );
    let sv = vkey_of(&steward);
    assert_eq!(verify_checkpoint(&note, &sv, "jcp.example/steward", &log), Ok(13));
    let known = [sv.clone(), vkey_of(&one), NoteVerifier::from_vkey(&two.vkey()).unwrap()];
    assert_eq!(
        cosignatures(&note, &known).unwrap(),
        vec![("jcp.example/steward".to_owned(), 100), ("witness.example/one".to_owned(), 101), ("witness.example/two".to_owned(), 102)]
    );

    // another history of the same length does not reach the stated root
    let forged: Vec<Vec<u8>> = (0..13).map(|i| format!("forged {i}").into_bytes()).collect();
    assert_eq!(verify_checkpoint(&note, &sv, "jcp.example/steward", &forged), Err(NoteError::BadCheckpoint("root")));
    // a steward's checkpoint without its own cosignature is not its
    let unsigned = format!("{text}\n{}", one.cosign_checkpoint(&text, 101).unwrap());
    assert_eq!(verify_checkpoint(&unsigned, &sv, "jcp.example/steward", &log), Err(NoteError::Unverified));
    // a cosignature moved to another checkpoint does not verify
    let shorter = checkpoint("jcp.example/steward", &log, 12).unwrap();
    let moved = format!("{shorter}\n{}", one.cosign_checkpoint(&text, 101).unwrap());
    assert_eq!(cosignatures(&moved, &known[1..2]), Err(NoteError::BadSignature("witness.example/one".into())));

    // an extension line: an ML-DSA cosignature does not cover it, an Ed25519 one does
    let extended = format!("{text}an extension\n\n{}{}", one.cosign_checkpoint(&text, 101).unwrap(), cosign_checkpoint(&text, &two, 102).unwrap());
    assert_eq!(cosignatures(&extended, &known[1..2]).unwrap(), vec![("witness.example/one".to_owned(), 101)]);
    assert_eq!(cosignatures(&extended, &known[2..]), Err(NoteError::BadSignature("witness.example/two".into())));
}

/// A line signed by the key itself, with whatever timestamp it states.
fn raw_line(name: &str, seed: [u8; 32], message: &[u8], time: u64) -> String {
    let key = ExpandedSigningKey::<MlDsa44>::from_seed(&B32::from(seed));
    let sig = key.sign_deterministic(message, b"").unwrap().encode().to_vec();
    let c = MlDsaCosigner::from_seed(name, seed).unwrap();
    let mut payload = c.id().to_be_bytes().to_vec();
    payload.extend_from_slice(&time.to_be_bytes());
    payload.extend_from_slice(&sig);
    format!("\u{2014} {name} {}\n", B64.encode(payload))
}

fn by_hand(name: &str, time: u64, origin: &str, start: u64, end: u64, hash: &Hash) -> Vec<u8> {
    let mut m = b"subtree/v1\n\0".to_vec();
    m.push(u8::try_from(name.len()).unwrap());
    m.extend_from_slice(name.as_bytes());
    m.extend_from_slice(&time.to_be_bytes());
    m.push(u8::try_from(origin.len()).unwrap());
    m.extend_from_slice(origin.as_bytes());
    m.extend_from_slice(&start.to_be_bytes());
    m.extend_from_slice(&end.to_be_bytes());
    m.extend_from_slice(hash);
    m
}

#[test]
fn a_cosignature_or_key_out_of_shape_is_refused() {
    let c = MlDsaCosigner::from_seed(NAME, seed()).unwrap();
    let vk = vkey_of(&c);
    let (origin, count, root) = read_checkpoint(CHECKPOINT).unwrap();
    assert_eq!(cosigned_subtree(NAME, TIME, &origin, 0, count, &root).unwrap(), by_hand(NAME, TIME, &origin, 0, count, &root));

    // a timestamp past 2^63 - 1, signed by the key itself, is still refused; the bound itself is not
    assert!(c.cosign_checkpoint(CHECKPOINT, MAX_TIMESTAMP + 1).is_err());
    let past = raw_line(NAME, seed(), &by_hand(NAME, MAX_TIMESTAMP + 1, &origin, 0, count, &root), MAX_TIMESTAMP + 1);
    assert_eq!(cosignatures(&format!("{CHECKPOINT}\n{past}"), std::slice::from_ref(&vk)), Err(NoteError::BadSignature(NAME.into())));
    let edge = raw_line(NAME, seed(), &by_hand(NAME, MAX_TIMESTAMP, &origin, 0, count, &root), MAX_TIMESTAMP);
    assert_eq!(cosignatures(&format!("{CHECKPOINT}\n{edge}"), std::slice::from_ref(&vk)).unwrap(), vec![(NAME.to_owned(), MAX_TIMESTAMP)]);
    // the signature of the wrong length
    let mut payload = c.id().to_be_bytes().to_vec();
    payload.extend_from_slice(&[0; 8 + 64]);
    let short = format!("\u{2014} {NAME} {}\n", B64.encode(payload));
    assert_eq!(cosignatures(&format!("{CHECKPOINT}\n{short}"), std::slice::from_ref(&vk)), Err(NoteError::BadSignature(NAME.into())));
    // signed with a context string: not this format's signature
    let key = ExpandedSigningKey::<MlDsa44>::from_seed(&B32::from(seed()));
    let with_context = key.sign_deterministic(&by_hand(NAME, TIME, &origin, 0, count, &root), b"ctx").unwrap().encode().to_vec();
    let mut payload = c.id().to_be_bytes().to_vec();
    payload.extend_from_slice(&TIME.to_be_bytes());
    payload.extend_from_slice(&with_context);
    let contexted = format!("\u{2014} {NAME} {}\n", B64.encode(payload));
    assert_eq!(cosignatures(&format!("{CHECKPOINT}\n{contexted}"), std::slice::from_ref(&vk)), Err(NoteError::BadSignature(NAME.into())));

    // what a CosignedSubtree may say
    let h = [9; 32];
    assert!(cosigned_subtree(NAME, 0, ORIGIN, 0, 0, &h).is_err(), "an empty subtree");
    assert!(cosigned_subtree(NAME, 0, ORIGIN, 3, 5, &h).is_err(), "not a subtree");
    assert!(cosigned_subtree(NAME, 0, ORIGIN, 5, 3, &h).is_err(), "backwards");
    assert!(cosigned_subtree(NAME, 1, ORIGIN, 8, 12, &h).is_err(), "a timestamp on a subtree not starting at zero");
    assert!(cosigned_subtree(NAME, 1, ORIGIN, 0, 12, &h).is_ok());
    assert!(cosigned_subtree("has space", 0, ORIGIN, 0, 1, &h).is_err());
    assert!(cosigned_subtree(NAME, 0, &"o".repeat(256), 0, 1, &h).is_err());
    assert!(cosigned_subtree(NAME, 0, &"o".repeat(255), 0, 1, &h).is_ok());
    assert!(MlDsaCosigner::from_seed("a+b", seed()).is_err());

    // verifier keys: the wrong length, a key id not its own, a name not its own
    let mut body = vec![TYPE_MLDSA44_COSIGNATURE];
    body.extend_from_slice(&[0; 32]);
    assert_eq!(NoteVerifier::from_vkey(&format!("{NAME}+00000000+{}", B64.encode(&body))), Err(NoteError::BadKey("key length")));
    let good = c.vkey();
    assert_eq!(NoteVerifier::from_vkey(&good.replacen("760e9e80", "760e9e81", 1)), Err(NoteError::BadKey("key id")));
    assert_eq!(NoteVerifier::from_vkey(&good.replacen(NAME, "example.com/other", 1)), Err(NoteError::BadKey("key id")));

    // an Ed25519 cosigner verifies no subtree
    let ed = NoteVerifier::from_vkey(&NoteSigner::cosigner(NAME, SigningKey::from_bytes(&[3; 32])).unwrap().vkey()).unwrap();
    assert!(matches!(verify_subtree_cosignature(&ed, ORIGIN, 0, 1, &h, &[0; MLDSA44_SIG_LEN]), Err(NoteError::BadKey(_))));
}

/// A subtree cosignature states one subtree of the log; with a consistency
/// proof it holds for the log's checkpoint, and with an inclusion proof for
/// an entry in it.
#[test]
fn a_subtree_cosignature_covers_one_subtree_of_the_log() {
    let log = log(13);
    let d = leaves(&log);
    let ids: Vec<ContentId> = log.iter().map(|e| ContentId::of_bytes(e)).collect();
    assert_eq!(subtree_hash(&d, 0, 13), Some(sha256_root(&ids)), "the checkpoint's root is the subtree [0, 13)");

    let w = MlDsaCosigner::from_seed("witness.example/one", [2; 32]).unwrap();
    let wv = vkey_of(&w);
    let h = subtree_hash(&d, 8, 12).unwrap();
    let sig = w.cosign_subtree("jcp.example/steward", 8, 12, &h).unwrap();
    assert_eq!(sig.len(), MLDSA44_SIG_LEN);
    assert_eq!(verify_subtree_cosignature(&wv, "jcp.example/steward", 8, 12, &h, &sig), Ok(()));
    let bad = Err(NoteError::BadSignature("witness.example/one".into()));
    assert_eq!(verify_subtree_cosignature(&wv, "jcp.example/steward", 8, 13, &h, &sig), bad);
    assert_eq!(verify_subtree_cosignature(&wv, "jcp.example/steward", 8, 12, &subtree_hash(&d, 0, 8).unwrap(), &sig), bad);
    assert_eq!(verify_subtree_cosignature(&wv, "jcp.example/other", 8, 12, &h, &sig), bad);
    let other = vkey_of(&MlDsaCosigner::from_seed("witness.example/one", [4; 32]).unwrap());
    assert_eq!(verify_subtree_cosignature(&other, "jcp.example/steward", 8, 12, &h, &sig), bad);
    // a subtree from zero is cosigned without a time too, unlike the checkpoint it may equal
    let h0 = subtree_hash(&d, 0, 8).unwrap();
    let from_zero = w.cosign_subtree("jcp.example/steward", 0, 8, &h0).unwrap();
    assert_eq!(verify_subtree_cosignature(&wv, "jcp.example/steward", 0, 8, &h0, &from_zero), Ok(()));

    // consistent with the checkpoint
    let text = checkpoint("jcp.example/steward", &log, 13).unwrap();
    let (_, n, root) = read_checkpoint(&text).unwrap();
    let proof = subtree_consistency_proof(&d, 8, 12).unwrap();
    assert!(verify_subtree_consistency(n, 8, 12, &proof, &h, &root));
    // and holding receipt 10
    let incl = subtree_inclusion_proof(&d, 8, 12, 10).unwrap();
    assert_eq!(evaluate_subtree_inclusion(&d[10], 10, 8, 12, &incl), Some(h));
    assert_eq!(evaluate_subtree_inclusion(&leaf_hash(&ContentId::of_bytes(b"forged").0), 10, 8, 12, &incl).filter(|r| *r == h), None);
}