//! RFC 8391 XMSS-SHA2_10_256 signature VERIFICATION core.
//!
//! This is the exact inner computation the Post-Quantum Finality aggregation zkVM
//! guest runs PER VALIDATOR: recover a validator's long-term XMSS root from one
//! hash-based signature over the domain-bound block message. A validator's
//! attestation is VALID iff the recovered root equals the root the registry
//! committed for that validator.
//!
//! Parameter set XMSS-SHA2_10_256 (RFC 8391 OID 0x00000001):
//!   n = 32 (SHA-256), Winternitz w = 16, len = 67 chains, single tree height h = 10.
//!
//! It is a byte-for-byte model of the already-validated on-chain
//! `AereXmssVerifier.sol` and its independent JS oracle
//! (`contracts/test/xmssVerifier.test.js`), which are checked against the official
//! github.com/XMSS/xmss-reference known-answer vector. Deterministic, `no_std`,
//! allocation-free, float-free, hash-only: exactly the shape a zkVM circuit needs.
//!
//! RFC 8391 keyed hash toolbox (padding_len = n = 32):
//!   F(KEY, M)      = SHA256( toByte(0,32) || KEY || M )
//!   H(KEY, L||R)   = SHA256( toByte(1,32) || KEY || L || R )
//!   H_msg(KEY, M)  = SHA256( toByte(2,32) || KEY || M )
//!   PRF(KEY, ADRS) = SHA256( toByte(3,32) || KEY || ADRS )
//! with per-address keys/bitmasks derived by PRF(SEED, ADRS) over the 32-byte hash
//! address (type + OTS/L-tree/hash-tree fields + key_and_mask in {0,1,2}).

use crate::sha256::sha256;

/// Hash output size / node size in bytes (SHA-256).
pub const N: usize = 32;
/// Winternitz parameter.
pub const W: usize = 16;
/// Total WOTS+ chains (64 message digits + 3 checksum digits).
pub const LEN: usize = 67;
/// Single-tree Merkle height.
pub const H: usize = 10;

/// A 32-byte hash / node.
pub type Hash = [u8; N];

// XMSS hash-address types (RFC 8391 section 2.5).
const ADDR_OTS: u32 = 0;
const ADDR_LTREE: u32 = 1;
const ADDR_HASHTREE: u32 = 2;

/// An XMSS-SHA2_10_256 public key: the long-term root plus the public SEED.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct XmssPubKey {
    /// Long-term XMSS public root (first 32 bytes of the XMSS public key).
    pub root: Hash,
    /// XMSS public SEED (last 32 bytes of the XMSS public key).
    pub seed: Hash,
}

/// A single-tree XMSS-SHA2_10_256 signature.
#[derive(Clone, Copy, Debug)]
pub struct XmssSig {
    /// Leaf index used by the signer (0 .. 2^h - 1).
    pub idx: u32,
    /// Per-signature randomizer R.
    pub r: Hash,
    /// The 67 WOTS+ one-time-signature chain values.
    pub wots: [Hash; LEN],
    /// The h = 10 Merkle authentication-path nodes.
    pub auth: [Hash; H],
}

// --------------------------------------------------------------------------
//                    RFC 8391 keyed hash toolbox (SHA-256)
// --------------------------------------------------------------------------

/// A 32-byte big-endian encoding of a small domain-separation prefix
/// (`toByte(v, 32)`): 31 zero bytes then `v`.
#[inline]
fn to_byte32(v: u8) -> Hash {
    let mut b = [0u8; 32];
    b[31] = v;
    b
}

/// Pack an XMSS 32-byte hash address. layer (bytes 0..4) and tree (bytes 4..12)
/// are always zero for single-tree XMSS-SHA2_10_256; then type (bytes 12..16),
/// the type-specific words a4/a5/a6 (bytes 16..20 / 20..24 / 24..28) and
/// key_and_mask (bytes 28..32). Matches `AereXmssVerifier._addr` exactly:
///   value = (typ<<128) | (a4<<96) | (a5<<64) | (a6<<32) | km  (big-endian bytes32).
#[inline]
fn addr(typ: u32, a4: u32, a5: u32, a6: u32, km: u32) -> Hash {
    let mut a = [0u8; 32];
    a[12..16].copy_from_slice(&typ.to_be_bytes());
    a[16..20].copy_from_slice(&a4.to_be_bytes());
    a[20..24].copy_from_slice(&a5.to_be_bytes());
    a[24..28].copy_from_slice(&a6.to_be_bytes());
    a[28..32].copy_from_slice(&km.to_be_bytes());
    a
}

/// PRF(SEED, ADRS) = SHA256( toByte(3,32) || SEED || ADRS ).
#[inline]
fn prf(seed: &Hash, address: &Hash) -> Hash {
    let mut buf = [0u8; 96];
    buf[0..32].copy_from_slice(&to_byte32(3));
    buf[32..64].copy_from_slice(seed);
    buf[64..96].copy_from_slice(address);
    sha256(&buf)
}

#[inline]
fn xor32(a: &Hash, b: &Hash) -> Hash {
    let mut o = [0u8; 32];
    for i in 0..32 {
        o[i] = a[i] ^ b[i];
    }
    o
}

/// RFC 8391 F: one hash-chain step. key = PRF(SEED, addr(km=0)),
/// mask = PRF(SEED, addr(km=1)), then F = SHA256( toByte(0,32) || key || (x^mask) ).
#[inline]
fn f_hash(seed: &Hash, typ: u32, a4: u32, a5: u32, a6: u32, x: &Hash) -> Hash {
    let key = prf(seed, &addr(typ, a4, a5, a6, 0));
    let mask = prf(seed, &addr(typ, a4, a5, a6, 1));
    let masked = xor32(x, &mask);
    let mut buf = [0u8; 96];
    // toByte(0,32) is all zero.
    buf[32..64].copy_from_slice(&key);
    buf[64..96].copy_from_slice(&masked);
    sha256(&buf)
}

/// RFC 8391 RAND_HASH / H over two n-byte children. key = PRF(addr(km=0)),
/// m0 = PRF(addr(km=1)), m1 = PRF(addr(km=2)), then
/// H = SHA256( toByte(1,32) || key || (L^m0) || (R^m1) ).
#[inline]
fn h_hash(seed: &Hash, typ: u32, a4: u32, a5: u32, a6: u32, left: &Hash, right: &Hash) -> Hash {
    let key = prf(seed, &addr(typ, a4, a5, a6, 0));
    let m0 = prf(seed, &addr(typ, a4, a5, a6, 1));
    let m1 = prf(seed, &addr(typ, a4, a5, a6, 2));
    let l = xor32(left, &m0);
    let r = xor32(right, &m1);
    let mut buf = [0u8; 128];
    buf[0..32].copy_from_slice(&to_byte32(1));
    buf[32..64].copy_from_slice(&key);
    buf[64..96].copy_from_slice(&l);
    buf[96..128].copy_from_slice(&r);
    sha256(&buf)
}

// --------------------------------------------------------------------------
//                              WOTS+ (leaf)
// --------------------------------------------------------------------------

/// Derive the 67 base-16 chain lengths (64 message digits + 3 checksum digits)
/// from a 32-byte message hash, exactly per RFC 8391 chain_lengths.
pub fn chain_lengths(m: &Hash) -> [u32; LEN] {
    let mut d = [0u32; LEN];
    let mut csum: u32 = 0;
    for i in 0..32 {
        let b = m[i];
        let hi = (b >> 4) as u32;
        let lo = (b & 0x0f) as u32;
        d[2 * i] = hi;
        d[2 * i + 1] = lo;
        csum += (15 - hi) + (15 - lo);
    }
    // csum << (8 - (len2*log_w % 8)) == csum << 4, then base_w over 2 big-endian bytes.
    let c = csum << 4;
    d[64] = (c >> 12) & 0x0f;
    d[65] = (c >> 8) & 0x0f;
    d[66] = (c >> 4) & 0x0f;
    d
}

/// WOTS_PKFromSig: complete each of the 67 chains from the signature value to the
/// chain end (position w-1), returning the 67 WOTS+ public-key chain values.
fn wots_pk_from_sig(seed: &Hash, idx_leaf: u32, mhash: &Hash, sig: &[Hash; LEN]) -> [Hash; LEN] {
    let lens = chain_lengths(mhash);
    let mut pk = [[0u8; 32]; LEN];
    for i in 0..LEN {
        let mut x = sig[i];
        // gen_chain: for s in [lens[i], w-1): F with hash-address s, chain-address i.
        let mut s = lens[i];
        while s < (W as u32) - 1 {
            x = f_hash(seed, ADDR_OTS, idx_leaf, i as u32, s, &x);
            s += 1;
        }
        pk[i] = x;
    }
    pk
}

/// L-tree: compress the 67 WOTS+ public-key values into a single n-byte leaf.
fn l_tree(seed: &Hash, idx_leaf: u32, pk: &[Hash; LEN]) -> Hash {
    let mut nodes = *pk;
    let mut l = LEN;
    let mut height: u32 = 0;
    while l > 1 {
        let parent = l >> 1;
        for i in 0..parent {
            nodes[i] = h_hash(
                seed,
                ADDR_LTREE,
                idx_leaf,
                height,
                i as u32,
                &nodes[2 * i],
                &nodes[2 * i + 1],
            );
        }
        if l & 1 == 1 {
            nodes[parent] = nodes[l - 1];
            l = parent + 1;
        } else {
            l = parent;
        }
        height += 1;
    }
    nodes[0]
}

// --------------------------------------------------------------------------
//                        Merkle authentication path
// --------------------------------------------------------------------------

/// compute_root: fold the leaf with the h=10 authentication path up to the root,
/// using the hash-tree address (type 2) with the correct per-level height/index.
fn compute_root(seed: &Hash, leaf: &Hash, leaf_idx: u32, auth: &[Hash; H]) -> Hash {
    let mut left: Hash;
    let mut right: Hash;
    if leaf_idx & 1 == 1 {
        left = auth[0];
        right = *leaf;
    } else {
        left = *leaf;
        right = auth[0];
    }
    let mut li = leaf_idx;
    for i in 0..(H - 1) {
        li >>= 1;
        let node = h_hash(seed, ADDR_HASHTREE, 0, i as u32, li, &left, &right);
        if li & 1 == 1 {
            left = auth[i + 1];
            right = node;
        } else {
            left = node;
            right = auth[i + 1];
        }
    }
    li >>= 1;
    h_hash(seed, ADDR_HASHTREE, 0, (H - 1) as u32, li, &left, &right)
}

// --------------------------------------------------------------------------
//                                Verify
// --------------------------------------------------------------------------

/// Recover the candidate XMSS root from a signature over `message`, binding to the
/// claimed public `root` in the message hash exactly as RFC 8391 H_msg requires:
///   M' = SHA256( toByte(2,32) || R || root || toByte(idx, 32) || M )
/// then WOTS_PKFromSig -> L-tree leaf -> compute_root along the auth path.
///
/// The signature is VALID iff the returned value equals `pk.root`. This function
/// is the exact per-validator obligation the aggregation circuit discharges.
pub fn xmss_recover_root(pk: &XmssPubKey, sig: &XmssSig, message: &[u8]) -> Hash {
    // M' = H_msg( R || root || toByte(idx,32) || M ).
    let mut msg_buf = [0u8; 32 + 32 + 32 + 32]; // prefix || R || root || toByte(idx,32)
    msg_buf[0..32].copy_from_slice(&to_byte32(2));
    msg_buf[32..64].copy_from_slice(&sig.r);
    msg_buf[64..96].copy_from_slice(&pk.root);
    // toByte(idx, 32): 28 zero bytes then the 4-byte big-endian index.
    msg_buf[124..128].copy_from_slice(&sig.idx.to_be_bytes());

    let mut hasher = crate::sha256::Sha256::new();
    hasher.update(&msg_buf);
    hasher.update(message);
    let mhash = hasher.finalize();

    let idx_leaf = sig.idx & ((1u32 << H) - 1);
    let wpk = wots_pk_from_sig(&pk.seed, idx_leaf, &mhash, &sig.wots);
    let leaf = l_tree(&pk.seed, idx_leaf, &wpk);
    compute_root(&pk.seed, &leaf, idx_leaf, &sig.auth)
}

/// Verify an XMSS-SHA2_10_256 signature: true iff the recovered root equals the
/// public key's committed root. Constant-shape, hash-only, deterministic.
pub fn xmss_verify(pk: &XmssPubKey, sig: &XmssSig, message: &[u8]) -> bool {
    xmss_recover_root(pk, sig, message) == pk.root
}