package org.bouncycastle.crypto.generators;

import java.util.ArrayList;
import java.util.List;

import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.digests.Blake2bDigest;
import org.bouncycastle.crypto.params.Argon2Parameters;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Longs;
import org.bouncycastle.util.Pack;

/**
 * Argon2 PBKDF - Based on the results of https://password-hashing.net/ and https://www.ietf.org/archive/id/draft-irtf-cfrg-argon2-03.txt
 */
public class Argon2BytesGenerator
{
    private static final int ARGON2_BLOCK_SIZE = 1024;
    private static final int ARGON2_QWORDS_IN_BLOCK = ARGON2_BLOCK_SIZE / 8;

    private static final int ARGON2_ADDRESSES_IN_BLOCK = 128;

    private static final int ARGON2_PREHASH_DIGEST_LENGTH = 64;
    private static final int ARGON2_PREHASH_SEED_LENGTH = 72;

    private static final int ARGON2_SYNC_POINTS = 4;

    /* Minimum and maximum number of lanes (degree of parallelism) */
    private static final int MIN_PARALLELISM = 1;
    private static final int MAX_PARALLELISM = (1 << 24) - 1;

    /* Minimum and maximum digest size in bytes */
    private static final int MIN_OUTLEN = 4;

    /* Minimum and maximum number of passes */
    private static final int MIN_ITERATIONS = 1;

    private static final long M32L = 0xFFFFFFFFL;

    private static final byte[] ZERO_BYTES = new byte[4];

    private Argon2Parameters parameters;
    private BlockPool pool;
    private int memoryBlocks;
    private Block[] memory;
    private int segmentLength;
    private int laneLength;

    public Argon2BytesGenerator()
    {
    }

    /**
     * Initialise the Argon2BytesGenerator from the parameters.
     *
     * @param parameters Argon2 configuration.
     */
    public void init(Argon2Parameters parameters)
    {
        if (parameters.getVersion() != Argon2Parameters.ARGON2_VERSION_10 &&
            parameters.getVersion() != Argon2Parameters.ARGON2_VERSION_13)
        {
            throw new UnsupportedOperationException("unknown Argon2 version");
        }
        if (parameters.getType() != Argon2Parameters.ARGON2_d &&
            parameters.getType() != Argon2Parameters.ARGON2_i &&
            parameters.getType() != Argon2Parameters.ARGON2_id)
        {
            throw new UnsupportedOperationException("unknown Argon2 type");
        }

        if (parameters.getLanes() < MIN_PARALLELISM)
        {
            throw new IllegalStateException("lanes must be at least " + MIN_PARALLELISM);
        }
        else if (parameters.getLanes() > MAX_PARALLELISM)
        {
            throw new IllegalStateException("lanes must be at most " + MAX_PARALLELISM);
        }
        else if (parameters.getIterations() < MIN_ITERATIONS)
        {
            throw new IllegalStateException("iterations is less than: " + MIN_ITERATIONS);
        }

        this.parameters = parameters;

        // 2. Align memory size
        // Minimum memoryBlocks = 8L blocks, where L is the number of lanes
        int memoryBlocks = Math.max(parameters.getMemory(), 2 * ARGON2_SYNC_POINTS * parameters.getLanes());

        this.segmentLength = memoryBlocks / (ARGON2_SYNC_POINTS * parameters.getLanes());
        this.laneLength = segmentLength * ARGON2_SYNC_POINTS;

        // Ensure that all segments have equal length
        memoryBlocks = parameters.getLanes() * laneLength;
        this.memoryBlocks = memoryBlocks;

        BlockPool configured = parameters.getBlockPool();
        // if no pool is provided hold on to enough blocks for the primary memory
        this.pool = (configured != null) ? configured : new FixedBlockPool(memoryBlocks);
    }

    public int generateBytes(char[] password, byte[] out)
    {
        return generateBytes(parameters.getCharToByteConverter().convert(password), out);
    }

    public int generateBytes(char[] password, byte[] out, int outOff, int outLen)
    {
        return generateBytes(parameters.getCharToByteConverter().convert(password), out, outOff, outLen);
    }

    public int generateBytes(byte[] password, byte[] out)
    {
        return generateBytes(password, out, 0, out.length);
    }

    public int generateBytes(byte[] password, byte[] out, int outOff, int outLen)
    {
        if (outLen < MIN_OUTLEN)
        {
            throw new IllegalStateException("output length less than " + MIN_OUTLEN);
        }

        byte[] tmpBlockBytes = new byte[ARGON2_BLOCK_SIZE];

        allocateMemory();
        initialize(tmpBlockBytes, password, outLen);
        fillMemoryBlocks();
        digest(tmpBlockBytes, out, outOff, outLen);

        reset();

        return outLen;
    }

    // Allocate primary memory from the configured BlockPool.
    private void allocateMemory()
    {
        this.memory = new Block[memoryBlocks];

        for (int i = 0; i < memory.length; i++)
        {
            memory[i] = pool.allocate();
        }
    }

    // Return primary memory to the BlockPool.
    private void reset()
    {
        if (null != memory)
        {
            for (int i = 0; i < memory.length; i++)
            {
                Block b = memory[i];
                if (null != b)
                {
                    pool.deallocate(b);
                }
            }
        }
        memory = null;
    }

    private void fillMemoryBlocks()
    {
        FillBlock filler = new FillBlock(pool);
        Position position = new Position();
        for (int pass = 0; pass < parameters.getIterations(); ++pass)
        {
            position.pass = pass;

            for (int slice = 0; slice < ARGON2_SYNC_POINTS; ++slice)
            {
                position.slice = slice;

                for (int lane = 0; lane < parameters.getLanes(); ++lane)
                {
                    position.lane = lane;

                    fillSegment(filler, position);
                }
            }
        }
        filler.deallocate(pool);
    }

    private void fillSegment(FillBlock filler, Position position)
    {
        Block addressBlock = null, inputBlock = null;

        boolean dataIndependentAddressing = isDataIndependentAddressing(position);
        int startingIndex = getStartingIndex(position);
        int currentOffset = position.lane * laneLength + position.slice * segmentLength + startingIndex;
        int prevOffset = getPrevOffset(currentOffset);

        if (dataIndependentAddressing)
        {
            addressBlock = filler.addressBlock.clear();
            inputBlock = filler.inputBlock.clear();

            initAddressBlocks(filler, position, inputBlock, addressBlock);
        }

        final boolean withXor = isWithXor(position);

        for (int index = startingIndex; index < segmentLength; ++index)
        {
            long pseudoRandom = getPseudoRandom(filler, index, addressBlock, inputBlock, prevOffset, dataIndependentAddressing);
            int refLane = getRefLane(position, pseudoRandom);
            int refColumn = getRefColumn(position, index, pseudoRandom, refLane == position.lane);

            /* 2 Creating a new block */
            Block prevBlock = memory[prevOffset];
            Block refBlock = memory[((laneLength) * refLane + refColumn)];
            Block currentBlock = memory[currentOffset];

            if (withXor)
            {
                filler.fillBlockWithXor(prevBlock, refBlock, currentBlock);
            }
            else
            {
                filler.fillBlock(prevBlock, refBlock, currentBlock);
            }

            prevOffset = currentOffset;
            currentOffset++;
        }
    }

    private boolean isDataIndependentAddressing(Position position)
    {
        return (parameters.getType() == Argon2Parameters.ARGON2_i) ||
            (parameters.getType() == Argon2Parameters.ARGON2_id
                && (position.pass == 0)
                && (position.slice < ARGON2_SYNC_POINTS / 2)
            );
    }

    private void initAddressBlocks(FillBlock filler, Position position, Block inputBlock, Block addressBlock)
    {
        inputBlock.v[0] = intToLong(position.pass);
        inputBlock.v[1] = intToLong(position.lane);
        inputBlock.v[2] = intToLong(position.slice);
        inputBlock.v[3] = intToLong(memory.length);
        inputBlock.v[4] = intToLong(parameters.getIterations());
        inputBlock.v[5] = intToLong(parameters.getType());

        if ((position.pass == 0) && (position.slice == 0))
        {
            /* Don't forget to generate the first block of addresses: */
            nextAddresses(filler, inputBlock, addressBlock);
        }
    }

    private boolean isWithXor(Position position)
    {
        return !(position.pass == 0 || parameters.getVersion() == Argon2Parameters.ARGON2_VERSION_10);
    }

    private int getPrevOffset(int currentOffset)
    {
        if (currentOffset % laneLength == 0)
        {
            /* Last block in this lane */
            return currentOffset + laneLength - 1;
        }
        else
        {
            /* Previous block */
            return currentOffset - 1;
        }
    }

    private static int getStartingIndex(Position position)
    {
        if ((position.pass == 0) && (position.slice == 0))
        {
            return 2; /* we have already generated the first two blocks */
        }
        else
        {
            return 0;
        }
    }

    private void nextAddresses(FillBlock filler, Block inputBlock, Block addressBlock)
    {
        inputBlock.v[6]++;
        filler.fillBlock(inputBlock, addressBlock);
        filler.fillBlock(addressBlock, addressBlock);
    }

    /* 1.2 Computing the index of the reference block */
    /* 1.2.1 Taking pseudo-random value from the previous block */
    private long getPseudoRandom(FillBlock filler, int index, Block addressBlock, Block inputBlock, int prevOffset,
        boolean dataIndependentAddressing)
    {
        if (dataIndependentAddressing)
        {
            int addressIndex = index % ARGON2_ADDRESSES_IN_BLOCK;
            if (addressIndex == 0)
            {
                nextAddresses(filler, inputBlock, addressBlock);
            }
            return addressBlock.v[addressIndex];
        }
        else
        {
            return memory[prevOffset].v[0];
        }
    }

    private int getRefLane(Position position, long pseudoRandom)
    {
        int refLane = (int)(((pseudoRandom >>> 32)) % parameters.getLanes());

        if ((position.pass == 0) && (position.slice == 0))
        {
            /* Can not reference other lanes yet */
            refLane = position.lane;
        }
        return refLane;
    }

    private int getRefColumn(Position position, int index, long pseudoRandom, boolean sameLane)
    {
        int referenceAreaSize;
        int startPosition;

        if (position.pass == 0)
        {
            startPosition = 0;

            if (sameLane)
            {
                /* The same lane => add current segment */
                referenceAreaSize = position.slice * segmentLength + index - 1;
            }
            else
            {
                /* pass == 0 && !sameLane => position.slice > 0*/
                referenceAreaSize = position.slice * segmentLength + ((index == 0) ? (-1) : 0);
            }
        }
        else
        {
            startPosition = ((position.slice + 1) * segmentLength) % laneLength;

            if (sameLane)
            {
                referenceAreaSize = laneLength - segmentLength + index - 1;
            }
            else
            {
                referenceAreaSize = laneLength - segmentLength + ((index == 0) ? (-1) : 0);
            }
        }

        long relativePosition = pseudoRandom & 0xFFFFFFFFL;
        relativePosition = (relativePosition * relativePosition) >>> 32;
        relativePosition = referenceAreaSize - 1 - ((referenceAreaSize * relativePosition) >>> 32);

        return (int)(startPosition + relativePosition) % laneLength;
    }

    private void digest(byte[] tmpBlockBytes, byte[] out, int outOff, int outLen)
    {
        Block finalBlock = memory[laneLength - 1];

        /* XOR the last blocks */
        for (int i = 1; i < parameters.getLanes(); i++)
        {
            int lastBlockInLane = i * laneLength + (laneLength - 1);
            finalBlock.xorWith(memory[lastBlockInLane]);
        }

        finalBlock.toBytes(tmpBlockBytes);

        hash(tmpBlockBytes, out, outOff, outLen);
    }

    /**
     * H' - hash - variable length hash function
     */
    private void hash(byte[] input, byte[] out, int outOff, int outLen)
    {
        byte[] outLenBytes = new byte[4];
        Pack.intToLittleEndian(outLen, outLenBytes, 0);

        int blake2bLength = 64;

        if (outLen <= blake2bLength)
        {
            Blake2bDigest blake = new Blake2bDigest(outLen * 8);

            blake.update(outLenBytes, 0, outLenBytes.length);
            blake.update(input, 0, input.length);
            blake.doFinal(out, outOff);
        }
        else
        {
            Blake2bDigest digest = new Blake2bDigest(blake2bLength * 8);
            byte[] outBuffer = new byte[blake2bLength];

            /* V1 */
            digest.update(outLenBytes, 0, outLenBytes.length);
            digest.update(input, 0, input.length);
            digest.doFinal(outBuffer, 0);

            int halfLen = blake2bLength / 2, outPos = outOff;
            System.arraycopy(outBuffer, 0, out, outPos, halfLen);
            outPos += halfLen;

            int r = ((outLen + 31) / 32) - 2;

            for (int i = 2; i <= r; i++, outPos += halfLen)
            {
                /* V2 to Vr */
                digest.update(outBuffer, 0, outBuffer.length);
                digest.doFinal(outBuffer, 0);

                System.arraycopy(outBuffer, 0, out, outPos, halfLen);
            }

            int lastLength = outLen - 32 * r;

            /* Vr+1 */
            digest = new Blake2bDigest(lastLength * 8);

            digest.update(outBuffer, 0, outBuffer.length);
            digest.doFinal(out, outPos);
        }
    }

    /*
     * One BLAKE2 round over the 16 working words at the given (distinct) block
     * positions. The words are loaded into locals once, mixed entirely in
     * registers through the 8 G applications, and stored back once - the
     * arithmetic and ordering are identical to the previous array-indexed
     * roundFunction/F/quarterRound triple, only the per-step v[] loads/stores
     * are removed. This is the data-independent Argon2 compression, so holding
     * the words in registers introduces no data-dependent index or branch.
     *
     * Each G(a,b,c,d) is the BLAKE2b mixing with the Lyra PHC fBlaMka step
     * (a <- a + b + 2*aL*bL, aL = least 32 bits, + modulo 2^64); the (a & M32L)
     * masks are load-bearing and must stay.
     */
    private static void roundFunction(Block block,
                                      int v0, int v1, int v2, int v3,
                                      int v4, int v5, int v6, int v7,
                                      int v8, int v9, int v10, int v11,
                                      int v12, int v13, int v14, int v15)
    {
        final long[] v = block.v;

        /* column step */
        G(v, v0, v4, v8, v12);
        G(v, v1, v5, v9, v13);
        G(v, v2, v6, v10, v14);
        G(v, v3, v7, v11, v15);

        /* diagonal step */
        G(v, v0, v5, v10, v15);
        G(v, v1, v6, v11, v12);
        G(v, v2, v7, v8, v13);
        G(v, v3, v4, v9, v14);
    }

    private static void G(long[] v, int pa, int pb, int pc, int pd)
    {
        long a = v[pa], b = v[pb], c = v[pc], d = v[pd];

        a += b + 2 * (a & M32L) * (b & M32L);
        d = Longs.rotateRight(d ^ a, 32);
        c += d + 2 * (c & M32L) * (d & M32L);
        b = Longs.rotateRight(b ^ c, 24);
        a += b + 2 * (a & M32L) * (b & M32L);
        d = Longs.rotateRight(d ^ a, 16);
        c += d + 2 * (c & M32L) * (d & M32L);
        b = Longs.rotateRight(b ^ c, 63);

        v[pa] = a; v[pb] = b; v[pc] = c; v[pd] = d;
    }

    private void initialize(byte[] tmpBlockBytes, byte[] password, int outputLength)
    {
        /*
         * H0 = H64(p, τ, m, t, v, y, |P|, P, |S|, S, |L|, K, |X|, X)
         * -> 64 byte (ARGON2_PREHASH_DIGEST_LENGTH)
         */

        Blake2bDigest blake = new Blake2bDigest(ARGON2_PREHASH_DIGEST_LENGTH * 8);

        int[] values = { parameters.getLanes(), outputLength, parameters.getMemory(), parameters.getIterations(),
            parameters.getVersion(), parameters.getType() };

        Pack.intToLittleEndian(values, tmpBlockBytes, 0);
        blake.update(tmpBlockBytes, 0, values.length * 4);

        addByteString(tmpBlockBytes, blake, password);
        addByteString(tmpBlockBytes, blake, parameters.getSalt());
        addByteString(tmpBlockBytes, blake, parameters.getSecret());
        addByteString(tmpBlockBytes, blake, parameters.getAdditional());

        byte[] initialHashWithZeros = new byte[ARGON2_PREHASH_SEED_LENGTH];
        blake.doFinal(initialHashWithZeros, 0);

        fillFirstBlocks(tmpBlockBytes, initialHashWithZeros);
    }

    private static void addByteString(byte[] tmpBlockBytes, Digest digest, byte[] octets)
    {
        if (null == octets)
        {
            digest.update(ZERO_BYTES, 0, 4);
            return;
        }

        Pack.intToLittleEndian(octets.length, tmpBlockBytes, 0);
        digest.update(tmpBlockBytes, 0, 4);
        digest.update(octets, 0, octets.length);
    }

    /**
     * (H0 || 0 || i) 72 byte -> 1024 byte
     * (H0 || 1 || i) 72 byte -> 1024 byte
     */
    private void fillFirstBlocks(byte[] tmpBlockBytes, byte[] initialHashWithZeros)
    {
        byte[] initialHashWithOnes = new byte[ARGON2_PREHASH_SEED_LENGTH];
        System.arraycopy(initialHashWithZeros, 0, initialHashWithOnes, 0, ARGON2_PREHASH_DIGEST_LENGTH);
//        Pack.intToLittleEndian(1, initialHashWithOnes, ARGON2_PREHASH_DIGEST_LENGTH);
        initialHashWithOnes[ARGON2_PREHASH_DIGEST_LENGTH] = 1;

        for (int i = 0; i < parameters.getLanes(); i++)
        {
            Pack.intToLittleEndian(i, initialHashWithZeros, ARGON2_PREHASH_DIGEST_LENGTH + 4);
            Pack.intToLittleEndian(i, initialHashWithOnes, ARGON2_PREHASH_DIGEST_LENGTH + 4);

            hash(initialHashWithZeros, tmpBlockBytes, 0, ARGON2_BLOCK_SIZE);
            memory[i * laneLength + 0].fromBytes(tmpBlockBytes);

            hash(initialHashWithOnes, tmpBlockBytes, 0, ARGON2_BLOCK_SIZE);
            memory[i * laneLength + 1].fromBytes(tmpBlockBytes);
        }
    }

    private long intToLong(int x)
    {
        return (long)(x & M32L);
    }

    private static class FillBlock
    {
        final Block R;
        final Block Z;

        final Block addressBlock;
        final Block inputBlock;

        FillBlock(BlockPool pool)
        {
            R = pool.allocate();
            Z = pool.allocate();
            addressBlock = pool.allocate();
            inputBlock = pool.allocate();
        }

        void deallocate(BlockPool pool)
        {
            pool.deallocate(addressBlock);
            pool.deallocate(inputBlock);
            pool.deallocate(R);
            pool.deallocate(Z);
        }

        private void applyBlake()
        {
            /* Apply Blake2 on columns of 64-bit words: (0,1,...,15) , then
            (16,17,..31)... finally (112,113,...127) */
            for (int i = 0; i < 8; i++)
            {

                int i16 = 16 * i;
                roundFunction(Z,
                    i16, i16 + 1, i16 + 2,
                    i16 + 3, i16 + 4, i16 + 5,
                    i16 + 6, i16 + 7, i16 + 8,
                    i16 + 9, i16 + 10, i16 + 11,
                    i16 + 12, i16 + 13, i16 + 14,
                    i16 + 15
                );
            }

            /* Apply Blake2 on rows of 64-bit words: (0,1,16,17,...112,113), then
            (2,3,18,19,...,114,115).. finally (14,15,30,31,...,126,127) */
            for (int i = 0; i < 8; i++)
            {

                int i2 = 2 * i;
                roundFunction(Z,
                    i2, i2 + 1, i2 + 16,
                    i2 + 17, i2 + 32, i2 + 33,
                    i2 + 48, i2 + 49, i2 + 64,
                    i2 + 65, i2 + 80, i2 + 81,
                    i2 + 96, i2 + 97, i2 + 112,
                    i2 + 113
                );
            }
        }

        private void fillBlock(Block Y, Block currentBlock)
        {
            Z.copyBlock(Y);
            applyBlake();
            currentBlock.xor(Y, Z);
        }

        private void fillBlock(Block X, Block Y, Block currentBlock)
        {
            R.xor(X, Y);
            Z.copyBlock(R);
            applyBlake();
            currentBlock.xor(R, Z);
        }

        private void fillBlockWithXor(Block X, Block Y, Block currentBlock)
        {
            R.xor(X, Y);
            Z.copyBlock(R);
            applyBlake();
            currentBlock.xorWith(R, Z);
        }
    }

    /**
     * A single 1024-byte memory block used by the Argon2 mixing function.
     * Made public so callers can supply their own {@link BlockPool}
     * implementations - see {@link Argon2Parameters.Builder#withBlockPool}.
     */
    public static class Block
    {
        private static final int SIZE = ARGON2_QWORDS_IN_BLOCK;

        /* 128 * 8 Byte QWords */
        private final long[] v;

        public Block()
        {
            v = new long[SIZE];
        }

        void fromBytes(byte[] input)
        {
            if (input.length < ARGON2_BLOCK_SIZE)
            {
                throw new IllegalArgumentException("input shorter than blocksize");
            }

            Pack.littleEndianToLong(input, 0, v);
        }

        void toBytes(byte[] output)
        {
            if (output.length < ARGON2_BLOCK_SIZE)
            {
                throw new IllegalArgumentException("output shorter than blocksize");
            }

            Pack.longToLittleEndian(v, output, 0);
        }

        private void copyBlock(Block other)
        {
            System.arraycopy(other.v, 0, v, 0, SIZE);
        }

        private void xor(Block b1, Block b2)
        {
            long[] v0 = v, v1 = b1.v, v2 = b2.v;
            for (int i = 0; i < SIZE; i++)
            {
                v0[i] = v1[i] ^ v2[i];
            }
        }

        private void xorWith(Block b1)
        {
            long[] v0 = v, v1 = b1.v;
            for (int i = 0; i < SIZE; i++)
            {
                v0[i] ^= v1[i];
            }
        }

        private void xorWith(Block b1, Block b2)
        {
            long[] v0 = v, v1 = b1.v, v2 = b2.v;
            for (int i = 0; i < SIZE; i++)
            {
                v0[i] ^= v1[i] ^ v2[i];
            }
        }

        private Block clear()
        {
            Arrays.fill(v, 0);
            return this;
        }
    }

    private static class Position
    {
        int pass;
        int lane;
        int slice;

        Position()
        {
        }
    }

    /**
     * Strategy for allocating and recycling Argon2 {@link Block} objects.
     * <p>
     * An implementation may simply return fresh blocks from
     * {@link #allocate()} (the default behaviour) or pool them so the
     * underlying {@code long[]} buffers can be reused across successive
     * {@code generateBytes} calls. Implementations must accept matching
     * allocate/deallocate pairs - the generator does not guard against
     * double-deallocation of the same block.
     */
    public static interface BlockPool
    {
        Block allocate();

        void deallocate(Block block);
    }

    /**
     * Bounded pool that recycles up to {@code maxBlocks} {@link Block} objects.
     * Excess blocks returned via {@link #deallocate(Block)} are dropped and
     * left for the garbage collector. Returned blocks are zeroised both on
     * deallocation and again on allocation, so a recycled block is never
     * observed with stale data.
     */
    public static class FixedBlockPool
        implements BlockPool
    {
        // a synchronized list rather than java.util.concurrent (Java 5+), so the
        // class still compiles/runs on the legacy pre-1.5 distributions.
        private final int maxBlocks;
        private final List<Block> blocks;

        public FixedBlockPool(int maxBlocks)
        {
            this.maxBlocks = maxBlocks;
            this.blocks = new ArrayList<Block>(maxBlocks);
        }

        public Block allocate()
        {
            Block block = null;
            synchronized (blocks)
            {
                if (!blocks.isEmpty())
                {
                    block = (Block)blocks.remove(blocks.size() - 1);
                }
            }
            if (block == null)
            {
                return new Block();
            }
            // a deallocate() in another thread may not have published its clear()
            // - re-clear here so callers see a zeroised block.
            block.clear();
            return block;
        }

        public void deallocate(Block block)
        {
            block.clear();
            synchronized (blocks)
            {
                if (blocks.size() < maxBlocks)
                {
                    blocks.add(block);
                }
            }
        }
    }
}
