//the base type for all WOTS+ signer/verifier types.
//

use crate::{WotsPrf,WotsRng};
use rand_core::{RngCore,CryptoRng};
//fill in b with the result of base changing message followed by the checksum
//used by both SignerBase and VerifierBase, hence the nasty generics
fn get_exponents<const M_BYTES: usize, const L1: usize, const L2: usize, const L: usize, const W: usize, const W_EXP: usize>(message: &[u8; M_BYTES], b: &mut [u16; L]) {
    let mut b_ctr = 0usize;
	let mut this_mask = (W - 1) as u16;
    let base_shift = this_mask.leading_zeros() as usize - 8; //number of bits to shift left to put the mask at the start of an 8-bit block
    let mut this_shift = (base_shift + 8) as isize;
    this_mask <<= this_shift;
    for i in 0..(M_BYTES - 1) { //extract all the bits from message[i] that we can
        let this_halfword = ((message[i] as u16) << 8) | message[i+1] as u16;
        while this_shift > (8 - W_EXP) as isize {
            b[b_ctr] = (this_halfword & this_mask) >> this_shift;
            this_mask >>= W_EXP;
            this_shift -= W_EXP as isize;
            b_ctr += 1;
        }
        //adjust mask & shift for message[i+1]
        this_mask = ((W - 1) as u16) << (base_shift as isize + this_shift + W_EXP as isize) as usize;
        this_shift = base_shift as isize + this_shift + W_EXP as isize; //guaranteed to be positive
    }
    //final byte of message: same principle, but with zeros at the end
    let last_halfword = (message[M_BYTES - 1] as u16) << 8;
    while this_shift > (8 - W_EXP as isize) {
        b[b_ctr] = (last_halfword & this_mask) >> this_shift;
        this_mask >>= W_EXP;
        this_shift -= W_EXP as isize;
        b_ctr += 1;
    }
    //sanity check!
    //assert!(b_ctr == L1);
    //create Ci
    //this is done a little differently, since C has a bit width that is precisely
    //proscribed by m and w and may not really match up with 32 in any way
    //plus C is a numeric type instead of a string of bytes, so bit shifts are easier
    let mut c: u32 = (L1*(W - 1)) as u32;
    for i in 0..L1 {
        c -= b[i] as u32;
    }
    for i in 1usize..=(L2) {
        b[L - i] = ((c >> (W_EXP*i)) as u16) & ((W - 1) as u16);
    }
}

//because we can't use associated constants in the definition of the structs
pub const fn getW(W_EXP: usize) -> usize {
    1usize << W_EXP
}

pub const fn getL1(W_EXP: usize, M_BYTES: usize) -> usize {
    let W: usize = getW(W_EXP);
    (8*M_BYTES + W_EXP - 1)/W_EXP
}

pub const fn getL2(W_EXP: usize, M_BYTES: usize) -> usize {
    let W: usize = getW(W_EXP);
    let L1: usize = getL1(W_EXP, M_BYTES);
    (usize::BITS - (L1*(W - 1)).leading_zeros()) as usize/W_EXP + 1 //floor(log2(l1*(w-1))/log2(w)) + 1, BITS - leading_zeros gives the integeral part of log2
}

pub const fn getL(W_EXP: usize, M_BYTES: usize) -> usize {
    let L1: usize = getL1(W_EXP,M_BYTES);
    let L2: usize = getL2(W_EXP,M_BYTES);
    L1 + L2
}

//type parameters:
//W_EXP = exponent of w. w = 2^W_EXP.
//M = length of a message, in bits
//N_BYTES = number of bytes that the PRF takes as input and outputs
//K_BYTES = number of bytes that the PRF takes as a key
//PRNG_BYTES = # of bytes to use for the PRNG seed. 
//K_BYTES and PRNG_BYTES can be zero if you're using some other kind of initialization method
//such as XMSS
//W,L1,L2,L MUST be the result of getW,getL1,getL2,getL respectively (this is because
//associated & generic consts apparently don't work together too well)
//PRF is the PRF you want to use (hash function)
//RNG is the PRNG generator you want to use (generates random bytes)
#[repr(C)]
#[derive(Clone)]
pub struct SignerBase<const W_EXP: usize, const M_BYTES: usize, const N_BYTES: usize, const K_BYTES: usize, const PRNG_BYTES: usize, const W: usize, const L1: usize, const L2: usize, const L: usize, PRF: WotsPrf<K_BYTES,N_BYTES> + Clone, RNG: WotsRng<PRNG_BYTES> + Clone> {
    rng_seed: [u8; PRNG_BYTES],
    K: [u8; K_BYTES],
    sk: [[u8; N_BYTES]; L],
    pub pk: [[u8; N_BYTES]; L],
    phantom_prf: core::marker::PhantomData<PRF>,
    phantom_rng: core::marker::PhantomData<RNG>
}

#[repr(C)]
#[derive(Clone)]
pub struct VerifierBase<const W_EXP: usize, const M_BYTES: usize, const N_BYTES: usize, const K_BYTES: usize, const PRNG_BYTES: usize, const W:usize, const L1: usize, const L2:usize, const L: usize, PRF: WotsPrf<K_BYTES,N_BYTES> + Clone, RNG: WotsRng<PRNG_BYTES> + Clone> {
    rng_seed: [u8; PRNG_BYTES],
    K: [u8; K_BYTES],
    pk: [[u8; N_BYTES]; L],
    phantom_prf: core::marker::PhantomData<PRF>,
    phantom_rng: core::marker::PhantomData<RNG>
}

impl<const W_EXP: usize, const M_BYTES: usize, const N_BYTES: usize, const K_BYTES: usize, const PRNG_BYTES: usize, const W: usize, const L1: usize, const L2: usize, const L: usize, PRF: WotsPrf<K_BYTES,N_BYTES> + Clone, RNG: WotsRng<PRNG_BYTES> + Clone> SignerBase<W_EXP,M_BYTES,N_BYTES,K_BYTES,PRNG_BYTES,W,L1,L2,L,PRF,RNG> {
    pub const PKLEN: usize = L*N_BYTES + PRNG_BYTES + K_BYTES;
    pub const SIGLEN: usize = L*N_BYTES;
    pub const KEYPAIRLEN: usize = Self::PKLEN + Self::SIGLEN;
    //generate a new keypair. Uses seed_rng to choose random values, signer_rng to actually make
    //the signature
    pub fn generate_inplace<R: RngCore + CryptoRng>(&mut self, seed_rng: &mut R, rng_obj: Option<&RNG>) {
        seed_rng.fill_bytes(&mut self.K);
        seed_rng.fill_bytes(&mut self.rng_seed);
        for slot in self.sk.iter_mut() {
            seed_rng.fill_bytes(slot);
        }
        //copy secself key into private key
        for (sk_slot,pk_slot) in self.sk.iter().zip(self.pk.iter_mut()) {
            pk_slot.copy_from_slice(sk_slot);
        }
        let mut my_prng = if let Some(obj) = rng_obj {
                RNG::mk_rng_with_obj(obj,&self.rng_seed)
            } else {
                RNG::mk_rng(&self.rng_seed)
            };
        for i in 1..W { //for each chain level 1 to w-1
            let mut ri = [0u8; N_BYTES];
            my_prng.fill_bytes(&mut ri);
            for pk_slot in self.pk.iter_mut() { //for each pk slot
                for (pkj,rij) in pk_slot.iter_mut().zip(ri.iter()) { //xor pk_slot with ri
                    *pkj ^= *rij;
                }
                //now apply PRF
                PRF::do_prf(&self.K,pk_slot);
            }
        }
    }

    //same as generate_inplace but with a PRF object
    pub fn generate_inplace_with_prf_obj<R: RngCore + CryptoRng>(&mut self, seed_rng: &mut R, prf_obj: &PRF, rng_obj: Option<&RNG>) {
        seed_rng.fill_bytes(&mut self.K);
        seed_rng.fill_bytes(&mut self.rng_seed);
        for slot in self.sk.iter_mut() {
            seed_rng.fill_bytes(slot);
        }
        //copy secself key into private key
        for (sk_slot,pk_slot) in self.sk.iter().zip(self.pk.iter_mut()) {
            pk_slot.copy_from_slice(sk_slot);
        }
        let mut my_prng = if let Some(obj) = rng_obj {
                RNG::mk_rng_with_obj(obj,&self.rng_seed)
            } else {
                RNG::mk_rng(&self.rng_seed)
            };
        for i in 1..W { //for each chain level 1 to w-1
            let mut ri = [0u8; N_BYTES];
            my_prng.fill_bytes(&mut ri);
            for pk_slot in self.pk.iter_mut() { //for each pk slot
                for (pkj,rij) in pk_slot.iter_mut().zip(ri.iter()) { //xor pk_slot with ri
                    *pkj ^= *rij;
                }
                //now apply PRF
                PRF::do_prf_with_obj(prf_obj,&self.K,pk_slot);
            }
        }
    }


    //same as generate(), but allocates and returns the SignerBase on the stack.
    pub fn generate<R: RngCore + CryptoRng>(seed_rng: &mut R, rng_obj: Option<&RNG>) -> Self {
        let mut ret = Self {
            rng_seed: [0u8; PRNG_BYTES],
            K: [0u8; K_BYTES],
            sk: [[0u8; N_BYTES]; L],
            pk: [[0u8; N_BYTES]; L],
            phantom_prf: core::marker::PhantomData::<PRF>::default(),
            phantom_rng: core::marker::PhantomData::<RNG>::default()
        };
        ret.generate_inplace(seed_rng, rng_obj);
        ret
    }

    //same as generate(), but with a PRF object
    pub fn generate_with_prf_obj<R: RngCore + CryptoRng>(seed_rng: &mut R, prf_obj: &PRF, rng_obj: Option<&RNG>) -> Self {
        let mut ret = Self {
            rng_seed: [0u8; PRNG_BYTES],
            K: [0u8; K_BYTES],
            sk: [[0u8; N_BYTES]; L],
            pk: [[0u8; N_BYTES]; L],
            phantom_prf: core::marker::PhantomData::<PRF>::default(),
            phantom_rng: core::marker::PhantomData::<RNG>::default()
        };
        ret.generate_inplace_with_prf_obj(seed_rng, prf_obj, rng_obj);
        ret
    }

    //create a SignerBase instance from a keypair that has been formatted into buf
    pub fn from_formatted_keypair(buf: &[u8]) -> Option<Self> {
        if buf.len() < Self::KEYPAIRLEN {
            None
        } else {
            let mut ret = Self {
                rng_seed: [0u8; PRNG_BYTES],
                K: [0u8; K_BYTES],
                sk: [[0u8; N_BYTES]; L],
                pk: [[0u8; N_BYTES]; L],
                phantom_prf: core::marker::PhantomData::<PRF>::default(),
                phantom_rng: core::marker::PhantomData::<RNG>::default()
            };
            Self::mogrify_from(buf.get(0..Self::KEYPAIRLEN).unwrap()).map(|x| { ret.clone_from(x) });
            Some(ret)
        }
    }

    //transmogrify a &'a mut [u8] into &'a mut Self, returning Some(ref) iff ptr is exactly wide
    //enough to hold Self
    //SignerBase has C representation and there is no chance of alignment issues,
    //and all possible values of a slice the length of size_of<Self>() are valid,
    //and Self::Drop has no side effects (just zeroizes self)
    //therefore this function is NOT unsafe.
    pub fn mogrify_from_mut<'a>(ptr: &'a mut [u8]) -> Option<&'a mut Self> {
        if ptr.len() != core::mem::size_of::<Self>() {
            None
        } else {
            unsafe {
                let p = ptr.as_mut_ptr() as *mut Self;
                p.as_mut()
            }
        }
    }

    //same as mogrify_from_mut but the refs are all immutable
    pub fn mogrify_from<'a>(ptr: &'a [u8]) -> Option<&'a Self> {
        if ptr.len() != core::mem::size_of::<Self>() {
            None
        } else {
            unsafe {
                let p = ptr.as_ptr() as *const Self;
                p.as_ref()
            }
        }
    }

    //formats the public key into buf
    //returns true iff buf was big enough and the keypair was copied in
    pub fn format_public_key(&self, buf: &mut [u8]) -> bool {
        if buf.len() < Self::PKLEN {
            false
        } else {
            buf[0..PRNG_BYTES].copy_from_slice(&self.rng_seed);
            buf[PRNG_BYTES..(PRNG_BYTES + K_BYTES)].copy_from_slice(&self.K);
            for i in 0..L {
                buf[(PRNG_BYTES + K_BYTES + i*N_BYTES)..(PRNG_BYTES + K_BYTES + (i+1)*N_BYTES)].copy_from_slice(&self.pk[i]);
            }
            true
        }
    }

    //formats the both public key and secret key into buf
    //returns true iff buf was big enough and the keypair was copied in
    //buf should should be at least Self::KEYPAIRLEN long.
    pub fn format_secret_keypair(&self, buf: &mut [u8]) -> bool {
        if buf.len() < Self::KEYPAIRLEN {
            false
        } else if let Some(buf_keypair) = Self::mogrify_from_mut(buf.get_mut(0..Self::KEYPAIRLEN).unwrap()) {
            buf_keypair.clone_from(self);
            true
        } else {
            false
        }
    }

    //signs a message and formats the signature into sig
    //returns Ok(()) if the message was signed, Err(self) if there isn't enough space in sig.
    //
    //AFTER THIS FUNCTION IS CALLED THE SignerBase WILL NOT BE USABLE
    pub fn sign(&mut self, message: &[u8], sig: &mut [u8], rng_obj: Option<&RNG>) -> Result<(),()> {
        if sig.len() < Self::SIGLEN || message.len() < M_BYTES {
            Err(())
        } else {
           //turn message into an array
           let message2 = <&[u8; M_BYTES]>::try_from(message.get(0..M_BYTES).unwrap()).unwrap();
           //get exponents
           let mut b: [u16; L] = [0u16; L];
           get_exponents::<M_BYTES,L1,L2,L,W,W_EXP>(message2,&mut b);
           //get rng
           let mut my_prng = if let Some(obj) = rng_obj {
                   RNG::mk_rng_with_obj(obj,&self.rng_seed)
               } else {
                   RNG::mk_rng(&self.rng_seed)
               };
           //use sk to compute the signature
           for i in 1..W { //for each chain level 1 to w-1
	           let mut ri = [0u8; N_BYTES];
	           my_prng.fill_bytes(&mut ri);
	           for (j,sk_slot) in self.sk.iter_mut().enumerate() {
                   if i as u16 <= b[j] { //prf should only be applied b[j] times, in other words once per element of 1..(b[j] + 1)
		               for (skj,rij) in sk_slot.iter_mut().zip(ri.iter()) { //xor sk_slot with ri
		                   *skj ^= *rij;
		               }
                       PRF::do_prf(&self.K,sk_slot);
                   }
	           }
           }
           //copy sk (which is now the signature) into sig
           for (j,sk_slot) in self.sk.iter().enumerate() {
               sig[(j*N_BYTES)..((j+1)*N_BYTES)].copy_from_slice(sk_slot);
           }
           Ok(())
        }
    }

    //same as sign() but PRF gets an object
    pub fn sign_with_prf_obj(&mut self, message: &[u8], sig: &mut [u8], prf_obj: &PRF, rng_obj: Option<&RNG>) -> Result<(),()> {
        if sig.len() < Self::SIGLEN || message.len() < M_BYTES {
            Err(())
        } else {
           //turn message into an array
           let message2 = <&[u8; M_BYTES]>::try_from(message.get(0..M_BYTES).unwrap()).unwrap();
           //get exponents
           let mut b: [u16; L] = [0u16; L];
           get_exponents::<M_BYTES,L1,L2,L,W,W_EXP>(message2,&mut b);
           //get rng
           let mut my_prng = if let Some(obj) = rng_obj {
                   RNG::mk_rng_with_obj(obj,&self.rng_seed)
               } else {
                   RNG::mk_rng(&self.rng_seed)
               };
           //use sk to compute the signature
           for i in 1..W { //for each chain level 1 to w-1
	           let mut ri = [0u8; N_BYTES];
	           my_prng.fill_bytes(&mut ri);
	           for (j,sk_slot) in self.sk.iter_mut().enumerate() {
                   if i as u16 <= b[j] { //prf should only be applied b[j] times, in other words once per element of 1..(b[j] + 1)
		               for (skj,rij) in sk_slot.iter_mut().zip(ri.iter()) { //xor sk_slot with ri
		                   *skj ^= *rij;
		               }
                       PRF::do_prf_with_obj(prf_obj,&self.K,sk_slot);
                   }
	           }
           }
           //copy sk (which is now the signature) into sig
           for (j,sk_slot) in self.sk.iter().enumerate() {
               sig[(j*N_BYTES)..((j+1)*N_BYTES)].copy_from_slice(sk_slot);
           }
           Ok(())
        }
    }
}

impl<const W_EXP: usize, const M_BYTES: usize, const N_BYTES: usize, const K_BYTES: usize, const PRNG_BYTES: usize, const W: usize, const L1: usize, const L2: usize, const L: usize, PRF: WotsPrf<K_BYTES,N_BYTES> + Clone, RNG: WotsRng<PRNG_BYTES> + Clone> VerifierBase<W_EXP,M_BYTES,N_BYTES,K_BYTES,PRNG_BYTES,W,L1,L2,L, PRF, RNG> {
    pub const PKLEN: usize = L*N_BYTES + PRNG_BYTES + K_BYTES;
    pub const SIGLEN: usize = L*N_BYTES;
    //transmogrify a &'a [u8] into &'a Self
    //safe for the same reasons SignerBase::mogrify_from() is safe
    pub fn mogrify_from<'a>(ptr: &'a [u8]) -> Option<&'a Self> {
        if ptr.len() != core::mem::size_of::<Self>() {
            None
        } else {
            unsafe {
                let p = ptr.as_ptr() as *const Self;
                p.as_ref()
            }
        }
    }

    //transmogrify a &'a mut [u8] into &'a mut Self
    //safe for the same reasons SignerBase::mogrify_from() is safe
    pub fn mogrify_from_mut<'a>(ptr: &'a mut [u8]) -> Option<&'a mut Self> {
        if ptr.len() != core::mem::size_of::<Self>() {
            None
        } else {
            unsafe {
                let p = ptr.as_mut_ptr() as *mut Self;
                p.as_mut()
            }
        }
    }

    //create a new VerifierBase out of a slice
    //returns Some(VerifierBase) iff buf was big enough to contain a public key
    pub fn from_formatted_pubkey(buf: &[u8]) -> Option<Self> {
        if buf.len() < Self::PKLEN {
            None
        } else {
            let mut ret = Self {
                rng_seed: [0u8; PRNG_BYTES],
                K: [0u8; K_BYTES],
                pk: [[0u8; N_BYTES]; L],
                phantom_prf: core::marker::PhantomData::<PRF>::default(),
                phantom_rng: core::marker::PhantomData::<RNG>::default()
            };
            Self::mogrify_from(buf.get(0..Self::PKLEN).unwrap()).map(|x| { ret.clone_from(x) });
            Some(ret)
        }
    }

    //copy the contents of this VerifierBase into a slice
    //returns true iff buf was big enough and the public key was written
    pub fn format_pubkey(&self, buf: &mut [u8]) -> bool {
        if buf.len() < Self::PKLEN {
            false
        } else if let Some(buf_pubkey) = Self::mogrify_from_mut(buf.get_mut(0..Self::PKLEN).unwrap()) {
            buf_pubkey.clone_from(self);
            true
        } else {
            false
        }
    }

    //verify a slice
    //Returns Err(()) if message or sig are too small to contain a message or sig
    //(respectively)
    //Ok(false) if the buffers were the right size, but the signature does not check out
    //Ok(true) if the buffers were the right size and the signature checks out
    //
    //this is called "verify_destructive" because it overwrites sig with the public key derivable
    //from sig.
    //Unfortunately I can't provide a version which doesn't overwrite sig on the BASE type.
    pub fn verify_destructive(&self, message: &[u8], sig: &mut [u8], rng_obj: Option<&RNG>) -> Result<bool,()> {
        if sig.len() < Self::SIGLEN || message.len() < M_BYTES {
            Err(())
        } else {
           //turn message into an array
           let message2 = <&[u8; M_BYTES]>::try_from(message.get(0..M_BYTES).unwrap()).unwrap();
           //get exponents
           let mut b = [0u16; L];
           get_exponents::<M_BYTES,L1,L2,L,W,W_EXP>(message2,&mut b);
           //get rng
           let mut my_prng = if let Some(obj) = rng_obj {
                   RNG::mk_rng_with_obj(obj,&self.rng_seed)
                } else {
                   RNG::mk_rng(&self.rng_seed)
                };
           //use sig to compute the signature
           for i in 1..W { //for each chain level 1 to w-1
	           let mut ri = [0u8; N_BYTES];
	           my_prng.fill_bytes(&mut ri);
               for j in 0..L {
                   if b[j] < i as u16 { //prf should be applied W - 1 - b[j] times, in other words once per element of (b[j] + 1)..W
                       //sig_slot is the array version of sig_j
                       let sig_slot = <&mut [u8;N_BYTES]>::try_from(sig.get_mut((j*N_BYTES)..((j+1)*N_BYTES)).unwrap()).unwrap();
		               for (skj,rij) in sig_slot.iter_mut().zip(ri.iter()) { //xor sig_slot with ri
		                   *skj ^= *rij;
		               }
                       //now apply prf
                       PRF::do_prf(&self.K,sig_slot);
                   }
               }
           }
           //check sig against the public key
           for (i,pk_slot) in self.pk.iter().enumerate() {
               let sig_slot = <&[u8;N_BYTES]>::try_from(sig.get((i*N_BYTES)..((i+1)*N_BYTES)).unwrap()).unwrap();
               if sig_slot != pk_slot {
                   return Ok(false);
               }
           }
           Ok(true)
        }
    }
    
    //same as verify_destructive, but with a PRF object
    pub fn verify_destructive_with_prf_obj(&self, message: &[u8], sig: &mut [u8], prf_obj: &PRF, rng_obj: Option<&RNG>) -> Result<bool,()> {
        if sig.len() < Self::SIGLEN || message.len() < M_BYTES {
            Err(())
        } else {
           //turn message into an array
           let message2 = <&[u8; M_BYTES]>::try_from(message.get(0..M_BYTES).unwrap()).unwrap();
           //get exponents
           let mut b = [0u16; L];
           get_exponents::<M_BYTES,L1,L2,L,W,W_EXP>(message2,&mut b);
           //get rng
           let mut my_prng = if let Some(obj) = rng_obj {
                   RNG::mk_rng_with_obj(obj,&self.rng_seed)
                } else {
                   RNG::mk_rng(&self.rng_seed)
                };
           //use sig to compute the signature
           for i in 1..W { //for each chain level 1 to w-1
	           let mut ri = [0u8; N_BYTES];
	           my_prng.fill_bytes(&mut ri);
               for j in 0..L {
                   if b[j] < i as u16 { //prf should be applied W - 1 - b[j] times, in other words once per element of (b[j] + 1)..W
                       //sig_slot is the array version of sig_j
                       let sig_slot = <&mut [u8;N_BYTES]>::try_from(sig.get_mut((j*N_BYTES)..((j+1)*N_BYTES)).unwrap()).unwrap();
		               for (skj,rij) in sig_slot.iter_mut().zip(ri.iter()) { //xor sig_slot with ri
		                   *skj ^= *rij;
		               }
                       //now apply prf
                       PRF::do_prf_with_obj(prf_obj,&self.K,sig_slot);
                   }
               }
           }
           //check sig against the public key
           for (i,pk_slot) in self.pk.iter().enumerate() {
               let sig_slot = <&[u8;N_BYTES]>::try_from(sig.get((i*N_BYTES)..((i+1)*N_BYTES)).unwrap()).unwrap();
               if sig_slot != pk_slot {
                   return Ok(false);
               }
           }
           Ok(true)
        }
    }
}

