package org.bouncycastle.pqc.crypto.gmss;

import java.util.Enumeration;
import java.util.Vector;

import org.bouncycastle.crypto.Digest;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Integers;
import org.bouncycastle.util.encoders.Hex;


This class computes a whole Merkle tree and saves the needed values for
AuthPath computation. It is used for precomputation of the root of a
following tree. After initialization, 2^H updates are required to complete
the root. Every update requires one leaf value as parameter. While computing
the root all initial values for the authentication path algorithm (treehash,
auth, retain) are stored for later use.
/**
 * This class computes a whole Merkle tree and saves the needed values for
 * AuthPath computation. It is used for precomputation of the root of a
 * following tree. After initialization, 2^H updates are required to complete
 * the root. Every update requires one leaf value as parameter. While computing
 * the root all initial values for the authentication path algorithm (treehash,
 * auth, retain) are stored for later use.
 */
public class GMSSRootCalc
{

    
max height of the tree
/**
     * max height of the tree
     */
    private int heightOfTree;

    
length of the messageDigest
/**
     * length of the messageDigest
     */
    private int mdLength;

    
the treehash instances of the tree
/**
     * the treehash instances of the tree
     */
    private Treehash[] treehash;

    
stores the retain nodes for authPath computation
/**
     * stores the retain nodes for authPath computation
     */
    private Vector[] retain;

    
finally stores the root of the tree when finished
/**
     * finally stores the root of the tree when finished
     */
    private byte[] root;

    
stores the authentication path y_1(i), i = 0..H-1
/**
     * stores the authentication path y_1(i), i = 0..H-1
     */
    private byte[][] AuthPath;

    
the value K for the authentication path computation
/**
     * the value K for the authentication path computation
     */
    private int K;

    
Vector element that stores the nodes on the stack
/**
     * Vector element that stores the nodes on the stack
     */
    private Vector tailStack;

    
stores the height of all nodes laying on the tailStack
/**
     * stores the height of all nodes laying on the tailStack
     */
    private Vector heightOfNodes;
    
The hash function used for the construction of the authentication trees
/**
     * The hash function used for the construction of the authentication trees
     */
    private Digest messDigestTree;

    
An array of strings containing the name of the hash function used to
construct the authentication trees and used by the OTS.
/**
     * An array of strings containing the name of the hash function used to
     * construct the authentication trees and used by the OTS.
     */
    private GMSSDigestProvider digestProvider;

    
stores the index of the current node on each height of the tree
/**
     * stores the index of the current node on each height of the tree
     */
    private int[] index;

    
true if instance was already initialized, false otherwise
/**
     * true if instance was already initialized, false otherwise
     */
    private boolean isInitialized;

    
true it instance was finished
/**
     * true it instance was finished
     */
    private boolean isFinished;

    
Integer that stores the index of the next seed that has to be omitted to
the treehashs
/**
     * Integer that stores the index of the next seed that has to be omitted to
     * the treehashs
     */
    private int indexForNextSeed;

    
temporary integer that stores the height of the next treehash instance
that gets initialized with a seed
/**
     * temporary integer that stores the height of the next treehash instance
     * that gets initialized with a seed
     */
    private int heightOfNextSeed;

    
Constructor
Params:
heightOfTree – maximal height of the tree
digestProvider –       an array of strings, containing the name of the used hash
                    function and PRNG and the name of the corresponding
                    provider/**
     * Constructor
     *
     * @param heightOfTree maximal height of the tree
     * @param digestProvider       an array of strings, containing the name of the used hash
     *                     function and PRNG and the name of the corresponding
     *                     provider
     */
    public GMSSRootCalc(int heightOfTree, int K, GMSSDigestProvider digestProvider)
    {
        this.heightOfTree = heightOfTree;
        this.digestProvider = digestProvider;
        this.messDigestTree = digestProvider.get();
        this.mdLength = messDigestTree.getDigestSize();
        this.K = K;
        this.index = new int[heightOfTree];
        this.AuthPath = new byte[heightOfTree][mdLength];
        this.root = new byte[mdLength];
        // this.treehash = new Treehash[this.heightOfTree - this.K];
        this.retain = new Vector[this.K - 1];
        for (int i = 0; i < K - 1; i++)
        {
            this.retain[i] = new Vector();
        }

    }

    
Initializes the calculation of a new root
Params:
sharedStack – the stack shared by all treehash instances of this tree/**
     * Initializes the calculation of a new root
     *
     * @param sharedStack the stack shared by all treehash instances of this tree
     */
    public void initialize(Vector sharedStack)
    {
        this.treehash = new Treehash[this.heightOfTree - this.K];
        for (int i = 0; i < this.heightOfTree - this.K; i++)
        {
            this.treehash[i] = new Treehash(sharedStack, i, this.digestProvider.get());
        }

        this.index = new int[heightOfTree];
        this.AuthPath = new byte[heightOfTree][mdLength];
        this.root = new byte[mdLength];

        this.tailStack = new Vector();
        this.heightOfNodes = new Vector();
        this.isInitialized = true;
        this.isFinished = false;

        for (int i = 0; i < heightOfTree; i++)
        {
            this.index[i] = -1;
        }

        this.retain = new Vector[this.K - 1];
        for (int i = 0; i < K - 1; i++)
        {
            this.retain[i] = new Vector();
        }

        this.indexForNextSeed = 3;
        this.heightOfNextSeed = 0;
    }

    
updates the root with one leaf and stores needed values in retain,
treehash or authpath. Additionally counts the seeds used. This method is
used when performing the updates for TREE++.
Params:
seed – the initial seed for treehash: seedNext
leaf – the height of the treehash/**
     * updates the root with one leaf and stores needed values in retain,
     * treehash or authpath. Additionally counts the seeds used. This method is
     * used when performing the updates for TREE++.
     *
     * @param seed the initial seed for treehash: seedNext
     * @param leaf the height of the treehash
     */
    public void update(byte[] seed, byte[] leaf)
    {
        if (this.heightOfNextSeed < (this.heightOfTree - this.K)
            && this.indexForNextSeed - 2 == index[0])
        {
            this.initializeTreehashSeed(seed, this.heightOfNextSeed);
            this.heightOfNextSeed++;
            this.indexForNextSeed *= 2;
        }
        // now call the simple update
        this.update(leaf);
    }

    
Updates the root with one leaf and stores the needed values in retain,
treehash or authpath
/**
     * Updates the root with one leaf and stores the needed values in retain,
     * treehash or authpath
     */
    public void update(byte[] leaf)
    {

        if (isFinished)
        {
            System.out.print("Too much updates for Tree!!");
            return;
        }
        if (!isInitialized)
        {
            System.err.println("GMSSRootCalc not initialized!");
            return;
        }

        // a new leaf was omitted, so raise index on lowest layer
        index[0]++;

        // store the nodes on the lowest layer in treehash or authpath
        if (index[0] == 1)
        {
            System.arraycopy(leaf, 0, AuthPath[0], 0, mdLength);
        }
        else if (index[0] == 3)
        {
            // store in treehash only if K < H
            if (heightOfTree > K)
            {
                treehash[0].setFirstNode(leaf);
            }
        }

        if ((index[0] - 3) % 2 == 0 && index[0] >= 3)
        {
            // store in retain if K = H
            if (heightOfTree == K)
            // TODO: check it
            {
                retain[0].insertElementAt(leaf, 0);
            }
        }

        // if first update to this tree is made
        if (index[0] == 0)
        {
            tailStack.addElement(leaf);
            heightOfNodes.addElement(Integers.valueOf(0));
        }
        else
        {

            byte[] help = new byte[mdLength];
            byte[] toBeHashed = new byte[mdLength << 1];

            // store the new leaf in help
            System.arraycopy(leaf, 0, help, 0, mdLength);
            int helpHeight = 0;
            // while top to nodes have same height
            while (tailStack.size() > 0
                && helpHeight == ((Integer)heightOfNodes.lastElement())
                .intValue())
            {

                // help <-- hash(stack top element || help)
                System.arraycopy(tailStack.lastElement(), 0, toBeHashed, 0,
                    mdLength);
                tailStack.removeElementAt(tailStack.size() - 1);
                heightOfNodes.removeElementAt(heightOfNodes.size() - 1);
                System.arraycopy(help, 0, toBeHashed, mdLength, mdLength);

                messDigestTree.update(toBeHashed, 0, toBeHashed.length);
                help = new byte[messDigestTree.getDigestSize()];
                messDigestTree.doFinal(help, 0);

                // the new help node is one step higher
                helpHeight++;
                if (helpHeight < heightOfTree)
                {
                    index[helpHeight]++;

                    // add index 1 element to initial authpath
                    if (index[helpHeight] == 1)
                    {
                        System.arraycopy(help, 0, AuthPath[helpHeight], 0,
                            mdLength);
                    }

                    if (helpHeight >= heightOfTree - K)
                    {
                        if (helpHeight == 0)
                        {
                            System.out.println("M���P");
                        }
                        // add help element to retain stack if it is a right
                        // node
                        // and not stored in treehash
                        if ((index[helpHeight] - 3) % 2 == 0
                            && index[helpHeight] >= 3)
                        // TODO: check it
                        {
                            retain[helpHeight - (heightOfTree - K)]
                                .insertElementAt(help, 0);
                        }
                    }
                    else
                    {
                        // if element is third in his line add it to treehash
                        if (index[helpHeight] == 3)
                        {
                            treehash[helpHeight].setFirstNode(help);
                        }
                    }
                }
            }
            // push help element to the stack
            tailStack.addElement(help);
            heightOfNodes.addElement(Integers.valueOf(helpHeight));

            // is the root calculation finished?
            if (helpHeight == heightOfTree)
            {
                isFinished = true;
                isInitialized = false;
                root = (byte[])tailStack.lastElement();
            }
        }

    }

    
initializes the seeds for the treehashs of the tree precomputed by this
class
Params:
seed –  the initial seed for treehash: seedNext
index – the height of the treehash/**
     * initializes the seeds for the treehashs of the tree precomputed by this
     * class
     *
     * @param seed  the initial seed for treehash: seedNext
     * @param index the height of the treehash
     */
    public void initializeTreehashSeed(byte[] seed, int index)
    {
        treehash[index].initializeSeed(seed);
    }

    
Method to check whether the instance has been initialized or not
Returns:
true if treehash was already initialized/**
     * Method to check whether the instance has been initialized or not
     *
     * @return true if treehash was already initialized
     */
    public boolean wasInitialized()
    {
        return isInitialized;
    }

    
Method to check whether the instance has been finished or not
Returns:
true if tree has reached its maximum height/**
     * Method to check whether the instance has been finished or not
     *
     * @return true if tree has reached its maximum height
     */
    public boolean wasFinished()
    {
        return isFinished;
    }

    
returns the authentication path of the first leaf of the tree
Returns:
the authentication path of the first leaf of the tree/**
     * returns the authentication path of the first leaf of the tree
     *
     * @return the authentication path of the first leaf of the tree
     */
    public byte[][] getAuthPath()
    {
        return GMSSUtils.clone(AuthPath);
    }

    
returns the initial treehash instances, storing value y_3(i)
Returns:
the initial treehash instances, storing value y_3(i)/**
     * returns the initial treehash instances, storing value y_3(i)
     *
     * @return the initial treehash instances, storing value y_3(i)
     */
    public Treehash[] getTreehash()
    {
        return GMSSUtils.clone(treehash);
    }

    
returns the retain stacks storing all right nodes near to the root
Returns:
the retain stacks storing all right nodes near to the root/**
     * returns the retain stacks storing all right nodes near to the root
     *
     * @return the retain stacks storing all right nodes near to the root
     */
    public Vector[] getRetain()
    {
        return GMSSUtils.clone(retain);
    }

    
returns the finished root value
Returns:
the finished root value/**
     * returns the finished root value
     *
     * @return the finished root value
     */
    public byte[] getRoot()
    {
        return Arrays.clone(root);
    }

    
returns the shared stack
Returns:
the shared stack/**
     * returns the shared stack
     *
     * @return the shared stack
     */
    public Vector getStack()
    {
        Vector copy = new Vector();
        for (Enumeration en = tailStack.elements(); en.hasMoreElements();)
        {
            copy.addElement(en.nextElement());
        }
        return copy;
    }

    
Returns the status byte array used by the GMSSPrivateKeyASN.1 class
Returns:
The status bytes/**
     * Returns the status byte array used by the GMSSPrivateKeyASN.1 class
     *
     * @return The status bytes
     */
    public byte[][] getStatByte()
    {

        int tailLength;
        if (tailStack == null)
        {
            tailLength = 0;
        }
        else
        {
            tailLength = tailStack.size();
        }
        byte[][] statByte = new byte[1 + heightOfTree + tailLength][64]; //FIXME: messDigestTree.getByteLength()
        statByte[0] = root;

        for (int i = 0; i < heightOfTree; i++)
        {
            statByte[1 + i] = AuthPath[i];
        }
        for (int i = 0; i < tailLength; i++)
        {
            statByte[1 + heightOfTree + i] = (byte[])tailStack.elementAt(i);
        }

        return statByte;
    }

    
Returns the status int array used by the GMSSPrivateKeyASN.1 class
Returns:
The status ints/**
     * Returns the status int array used by the GMSSPrivateKeyASN.1 class
     *
     * @return The status ints
     */
    public int[] getStatInt()
    {

        int tailLength;
        if (tailStack == null)
        {
            tailLength = 0;
        }
        else
        {
            tailLength = tailStack.size();
        }
        int[] statInt = new int[8 + heightOfTree + tailLength];
        statInt[0] = heightOfTree;
        statInt[1] = mdLength;
        statInt[2] = K;
        statInt[3] = indexForNextSeed;
        statInt[4] = heightOfNextSeed;
        if (isFinished)
        {
            statInt[5] = 1;
        }
        else
        {
            statInt[5] = 0;
        }
        if (isInitialized)
        {
            statInt[6] = 1;
        }
        else
        {
            statInt[6] = 0;
        }
        statInt[7] = tailLength;

        for (int i = 0; i < heightOfTree; i++)
        {
            statInt[8 + i] = index[i];
        }
        for (int i = 0; i < tailLength; i++)
        {
            statInt[8 + heightOfTree + i] = ((Integer)heightOfNodes
                .elementAt(i)).intValue();
        }

        return statInt;
    }

    
Returns:
a human readable version of the structure/**
     * @return a human readable version of the structure
     */
    public String toString()
    {
        String out = "";
        int tailLength;
        if (tailStack == null)
        {
            tailLength = 0;
        }
        else
        {
            tailLength = tailStack.size();
        }

        for (int i = 0; i < 8 + heightOfTree + tailLength; i++)
        {
            out = out + getStatInt()[i] + " ";
        }
        for (int i = 0; i < 1 + heightOfTree + tailLength; i++)
        {
            out = out + new String(Hex.encode(getStatByte()[i])) + " ";
        }
        out = out + "  " + digestProvider.get().getDigestSize();
        return out;
    }
}