package jdk.jfr.internal.util;

import java.util.Arrays;

import java.util.HashMap;

import java.util.Map;

import java.util.Objects;

import java.util.function.BiConsumer;

import java.util.function.Consumer;

public class PerfectHashMap<V> {

    private static final long COLLISION_SHIFT = 63;

    private static final long COLLISION_BIT = 1L << COLLISION_SHIFT;

    private static final long COLLISION_MASK = COLLISION_BIT - 1;

    private static final int MAX_REMAP_ATTEMPTS = 100000;

    private static final int  MAX_ATTEMPS_BEFORE_RESIZE = 100;

    static final long W = 64L;

    static class LinkedValue<V> {

        final V value;

        long next;

        LinkedValue(V value) {

            this.value = value;

            this.next = 0;

        }

    }

    private UniversalHashFamily hashFamily = new UniversalHashFamily();

    private PrimitiveHashMap<LinkedValue<V>> loadMap;

    private Object[] valueTable;

    private long[] routeTable;

    private long shift;

    private long shiftMask;

    private int tableLengthMask;

    private long primaryHashFunction = 0;

    private int collisions = 0;

    private int retries = 0;

    private int sizeFactor = 1;

    private boolean minimal;

    public V get(long key) {

        LinkedValue<V> v = loadMap.get(key);

        return v != null ? v.value : null;

    }

    public V put(long key, V value) {

        LinkedValue<V> existing = loadMap.put(key, new LinkedValue<V>(value));

        return existing != null ? existing.value : null;

    }

    public void forEach(BiConsumer<? super Long, ? super V> action) {

        //loadMap.forEach(PerfectHashMap<V>::callback);

    }

    public final void forEach(Consumer<? super V> action) {

        //loadMap.forEach(action);

    }

    public final long[] keys() {

        return loadMap.keys();

    }

    static class Log2 {

        private static final int MAX_SIZE_EXPONENT = 32;

        static long log2base10(long exponent) {

            return 1L << exponent;

        }

        static int log2(int value) {

            int i = 0;

            int lastMultiple = 0;

            while (i < MAX_SIZE_EXPONENT) {

                int multiple = (int)log2base10(i);

                if ((value & multiple) != 0) {

                    lastMultiple = i;

                }

                ++i;

            }

            return ((int)log2base10(lastMultiple) ^ value) != 0 ? lastMultiple + 1 : lastMultiple;

        }

            }

    static final int tableExponent(int cap) {

        return Log2.log2(cap);

    }

    PerfectHashMap() {

        this(false, 101);

    }

    PerfectHashMap(int size) {

        this(false, size);

    }

    PerfectHashMap(boolean minimal, int size) {

        this.minimal = minimal;

        this.loadMap = new PrimitiveHashMap<>(size);

        this.primaryHashFunction = hashFamily.getRandomHashFunction();

    }

    @SuppressWarnings("unchecked")

    public V getPerfect(long key) {

        int routeIdx = getIndex(key, primaryHashFunction);

        assert(routeIdx >= 0);

        assert(routeIdx < routeTable.length);

        long element = routeTable[routeIdx];

        int valueIdx = element < 0 ? getIndex(key, -element - 1) : (int)element;

        assert(valueIdx >= 0);

        assert(valueIdx < valueTable.length);

        return (V)valueTable[valueIdx];

    }

    private long getRandomHashFunction() {

        return hashFamily.getRandomHashFunction();

    }

    private int getIndex(long key, long hashFunction) {

       final int idx = UniversalHashFamily.getIndex(key, hashFunction, shift, shiftMask);

       assert(idx >= 0);

       assert(idx < routeTable.length);

       return idx;

    }

    private static boolean isColliding(long entry) {

        return entry < 0;

    }

    private boolean isNonColliding(long entry) {

        return entry > 0;

    }

    private static long setColliding(long entry) {

        return entry | COLLISION_BIT;

    }

    private static long read(long entry) {

        return entry & COLLISION_MASK;

    }

    private int nextValueTableSlot(int lastIdx) {

        assert(lastIdx < valueTable.length);

        int i = lastIdx;

        for (; i < valueTable.length; ++i) {

            if (valueTable[i] == null) {

                break;

            }

        }

        return i;

    }

    private int valueTableStore(V value, int lastIdx) {

        if (lastIdx > valueTable.length) {

            lastIdx = 0;

        }

        assert(lastIdx < valueTable.length);

        final int idx = nextValueTableSlot(lastIdx);

        assert(idx < valueTable.length);

        assert(valueTable[idx] == null);

        valueTable[idx] = value;

        return idx;

    }

    private void routeNonCollisions() {

        int lastIdx = 0;

        for (int i = 0; i < routeTable.length; ++i) {

            if (isNonColliding(routeTable[i])) {

                lastIdx = valueTableStore(loadMap.get(routeTable[i]).value, lastIdx);

                routeTable[i] = lastIdx++;

           }

        }

    }

    private void rollback(int idx, int length, long hashFunction) {

        assert(isColliding(routeTable[idx]));

        long key = read(routeTable[idx]);

        LinkedValue<V> v = loadMap.get(key); // boxing

        for (int i = 0; i < length; ++i) {

            final int valueIdx = getIndex(key, hashFunction);

            assert(valueIdx >= 0);

            assert(valueIdx < valueTable.length);

            assert(valueTable[valueIdx] != null);

            valueTable[valueIdx] = null;

            key = v.next;

            v = loadMap.get(v.next); // no boxing

        }

    }

    private boolean remap(int idx, long hashFunction) {

        assert(isColliding(routeTable[idx]));

        int completed = 0;

        long key = read(routeTable[idx]);

        LinkedValue<V> v = loadMap.get(key);

        while (key != 0) {

            final int valueIdx = getIndex(key, hashFunction);

            assert(valueIdx >= 0);

            assert(valueIdx < valueTable.length);

            if (valueTable[valueIdx] == null) {

                valueTable[valueIdx] = v.value;

                ++completed;

                key = v.next;

                v = loadMap.get(v.next);

                continue;

            }

            rollback(idx, completed, hashFunction);

            return false;

        }

        return true;

    }

    private boolean routeCollisions(int idx) {

        assert(isColliding(routeTable[idx]));

        boolean success = false;

        int attempts = 0;

        long randomHashFunction = 0;

        do {

            randomHashFunction = getRandomHashFunction();

            success = remap(idx, randomHashFunction);

            if (++attempts == MAX_REMAP_ATTEMPTS) {

                System.out.println("Failed number of attempts - restart: " + attempts);

                return false;

            }

        } while (!success);

        System.out.println("Number of remap attempts: " + attempts);

        routeTable[idx] = -1 - randomHashFunction;

        assert(-routeTable[idx] - 1 == randomHashFunction);

        return true;

    }

    private boolean routeCollisions() {

        for (int i = 0; i < routeTable.length; ++i) {

            if (isColliding(routeTable[i])) {

                if (!routeCollisions(i)) {

                    return false;

                }

            }

        }

        return true;

    }

    private static void clearLongTable(long[] table) {

        Arrays.fill(table, 0);

        for (int i = 0; i < table.length; ++i) {

            assert(table[i] == 0);

        }

    }

    private static <T extends Object> void clearReferenceTable(T[] table) {

        Arrays.fill(table, null);

        for (int i = 0; i < table.length; ++i) {

            assert(table[i] == null);

        }

    }

    private void unlinkChains() {

        for (long key : loadMap.keys()) {

            loadMap.get(key).next = 0;

        }

    }

    private void routeTableStore(long key, LinkedValue<V> value, int idx) {

        assert(idx >= 0);

        assert(idx < routeTable.length);

        long existing = read(routeTable[idx]);

        if (existing == 0) {

            routeTable[idx] = key;

            return;

        }

        ++collisions;

        routeTable[idx] = setColliding(existing);

        LinkedValue<V> existingValue = loadMap.get(existing);

        value.next = existingValue.next;

        existingValue.next = key;

    }

    private void mapKeys() {

        for (long key : loadMap.keys()) {

            routeTableStore(key, loadMap.get(key), getIndex(key, primaryHashFunction));

        }

    }

    private void validate() {

        for (long key : loadMap.keys()) {

            long element = routeTable[getIndex(key, primaryHashFunction)];

            int valueIdx = element < 0 ? getIndex(key, -element - 1) : (int)element;

            assert(valueIdx >= 0);

            assert(loadMap.get(key) == valueTable[valueIdx]);

        }

    }

    private void reset() {

        collisions = 0;

        clearLongTable(routeTable);

        clearReferenceTable(valueTable);

    }

    private int dimensionTableSize() {

        int size = loadMap.size() * sizeFactor;

        return (int)Log2.log2base10(Log2.log2(size));

    }

    @SuppressWarnings({"rawtypes","unchecked"})

    private void allocateTables() {

        int size = dimensionTableSize();

        this.tableLengthMask = size - 1;

        this.shift = W - tableExponent(size);

        this.shiftMask = Log2.log2base10(shift) - 1;

        routeTable = new long[size];

        valueTable = (V[])new Object[size];

        collisions = 0;

        retries = 0;

    }

    public void build() {

        start:

        while (true) {

            allocateTables();

            System.out.println("Table size " + routeTable.length);

            mapKeys();

            if (collisions > 0) {

                if (!routeCollisions()) {

                    unlinkChains();

                    if (++retries <= MAX_ATTEMPS_BEFORE_RESIZE) {

                      reset();

                    } else {

                      sizeFactor *= 2;

                    }

                    continue start;

                }

            }

            routeNonCollisions();

            return;

        }

    }

    public void rebuild() {

        sizeFactor = 1;

        build();

    }

    public int size() {

        return loadMap.size();

    }

}