/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*
 * Modifications Copyright OpenSearch Contributors. See
 * GitHub history for details.
 */

package org.opensearch.common.collect;

import org.apache.lucene.util.mutable.MutableValueInt;

import java.lang.reflect.Array;
import java.util.AbstractMap;
import java.util.AbstractSet;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.Collection;
import java.util.Deque;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.stream.Stream;

/**
 * An immutable map whose writes result in a new copy of the map to be created.
 *
 * This is essentially a hash array mapped trie: inner nodes use a bitmap in
 * order to map hashes to slots by counting ones. In case of a collision (two
 * values having the same 32-bits hash), a leaf node is created which stores
 * and searches for values sequentially.
 *
 * Reads and writes both perform in logarithmic time. Null keys and values are
 * not supported.
 *
 * This structure might need to perform several object creations per write so
 * it is better suited for work-loads that are not too write-intensive.
 *
 * @see <a href="http://en.wikipedia.org/wiki/Hash_array_mapped_trie">the wikipedia page</a>
 *
 * @opensearch.internal
 */
public final class CopyOnWriteHashMap<K, V> extends AbstractMap<K, V> {

    private static final int TOTAL_HASH_BITS = 32;
    private static final Object[] EMPTY_ARRAY = new Object[0];

    private static final int HASH_BITS = 6;
    private static final int HASH_MASK = 0x3F;

    /**
     * Return a copy of the provided map.
     */
    public static <K, V> CopyOnWriteHashMap<K, V> copyOf(Map<? extends K, ? extends V> map) {
        if (map instanceof CopyOnWriteHashMap) {
            // no need to copy in that case
            @SuppressWarnings("unchecked")
            final CopyOnWriteHashMap<K, V> cowMap = (CopyOnWriteHashMap<K, V>) map;
            return cowMap;
        } else {
            return new CopyOnWriteHashMap<K, V>().copyAndPutAll(map);
        }
    }

    /**
     * Abstraction of a node, implemented by both inner and leaf nodes.
     *
     * @opensearch.internal
     */
    private abstract static class Node<K, V> {

        /**
         * Recursively get the key with the given hash.
         */
        abstract V get(Object key, int hash);

        /**
         * Recursively add a new entry to this node. <code>hashBits</code> is
         * the number of bits that are still set in the hash. When this value
         * reaches a number that is less than or equal to {@code 0}, a leaf
         * node needs to be created since it means that a collision occurred
         * on the 32 bits of the hash.
         */
        abstract Node<K, V> put(K key, int hash, int hashBits, V value, MutableValueInt newValue);

        /**
         * Recursively remove an entry from this node.
         */
        abstract Node<K, V> remove(Object key, int hash);

        /**
         * For the current node only, append entries that are stored on this
         * node to <code>entries</code> and sub nodes to <code>nodes</code>.
         */
        abstract void visit(Deque<Map.Entry<K, V>> entries, Deque<Node<K, V>> nodes);

        /**
         * Whether this node stores nothing under it.
         */
        abstract boolean isEmpty();

    }

    /**
     * A leaf of the tree where all hashes are equal. Values are added and retrieved in linear time.
     *
     * @opensearch.internal
     */
    private static class Leaf<K, V> extends Node<K, V> {

        private final K[] keys;
        private final V[] values;

        Leaf(K[] keys, V[] values) {
            this.keys = keys;
            this.values = values;
        }

        @SuppressWarnings("unchecked")
        Leaf() {
            this((K[]) EMPTY_ARRAY, (V[]) EMPTY_ARRAY);
        }

        @Override
        boolean isEmpty() {
            return keys.length == 0;
        }

        @Override
        void visit(Deque<Map.Entry<K, V>> entries, Deque<Node<K, V>> nodes) {
            for (int i = 0; i < keys.length; ++i) {
                entries.add(new AbstractMap.SimpleImmutableEntry<>(keys[i], values[i]));
            }
        }

        @Override
        V get(Object key, int hash) {
            for (int i = 0; i < keys.length; i++) {
                if (key.equals(keys[i])) {
                    return values[i];
                }
            }
            return null;

        }

        private static <T> T[] replace(T[] array, int index, T value) {
            final T[] copy = Arrays.copyOf(array, array.length);
            copy[index] = value;
            return copy;
        }

        @Override
        Leaf<K, V> put(K key, int hash, int hashBits, V value, MutableValueInt newValue) {
            assert hashBits <= 0 : hashBits;
            int slot = -1;
            for (int i = 0; i < keys.length; i++) {
                if (key.equals(keys[i])) {
                    slot = i;
                    break;
                }
            }

            final K[] keys2;
            final V[] values2;

            if (slot < 0) {
                keys2 = appendElement(keys, key);
                values2 = appendElement(values, value);
                newValue.value = 1;
            } else {
                keys2 = replace(keys, slot, key);
                values2 = replace(values, slot, value);
            }

            return new Leaf<>(keys2, values2);
        }

        @Override
        Leaf<K, V> remove(Object key, int hash) {
            int slot = -1;
            for (int i = 0; i < keys.length; i++) {
                if (key.equals(keys[i])) {
                    slot = i;
                    break;
                }
            }
            if (slot < 0) {
                return this;
            }
            final K[] keys2 = removeArrayElement(keys, slot);
            final V[] values2 = removeArrayElement(values, slot);
            return new Leaf<>(keys2, values2);
        }
    }

    private static <T> T[] removeArrayElement(T[] array, int index) {
        final Object result = Array.newInstance(array.getClass().getComponentType(), array.length - 1);
        System.arraycopy(array, 0, result, 0, index);
        if (index < array.length - 1) {
            System.arraycopy(array, index + 1, result, index, array.length - index - 1);
        }

        return (T[]) result;
    }

    public static <T> T[] appendElement(final T[] array, final T element) {
        final T[] newArray = Arrays.copyOf(array, array.length + 1);
        newArray[newArray.length - 1] = element;
        return newArray;
    }

    public static <T> T[] insertElement(final T[] array, final T element, final int index) {
        final T[] result = Arrays.copyOf(array, array.length + 1);
        System.arraycopy(array, 0, result, 0, index);
        result[index] = element;
        if (index < array.length) {
            System.arraycopy(array, index, result, index + 1, array.length - index);
        }
        return result;
    }

    /**
     * An inner node in this trie. Inner nodes store up to 64 key-value pairs
     * and use a bitmap in order to associate hashes to them. For example, if
     * an inner node contains 5 values, then 5 bits will be set in the bitmap
     * and the ordinal of the bit set in this bit map will be the slot number.
     *
     * As a consequence, the number of slots in an inner node is equal to the
     * number of one bits in the bitmap.
     *
     * @opensearch.internal
     */
    private static class InnerNode<K, V> extends Node<K, V> {

        private final long mask; // the bitmap
        private final K[] keys;
        final Object[] subNodes; // subNodes[slot] is either a value or a sub node in case of a hash collision

        InnerNode(long mask, K[] keys, Object[] subNodes) {
            this.mask = mask;
            this.keys = keys;
            this.subNodes = subNodes;
            assert consistent();
        }

        // only used in assert
        private boolean consistent() {
            assert Long.bitCount(mask) == keys.length;
            assert Long.bitCount(mask) == subNodes.length;
            for (int i = 0; i < keys.length; ++i) {
                if (subNodes[i] instanceof Node) {
                    assert keys[i] == null;
                } else {
                    assert keys[i] != null;
                }
            }
            return true;
        }

        @Override
        boolean isEmpty() {
            return mask == 0;
        }

        @SuppressWarnings("unchecked")
        InnerNode() {
            this(0, (K[]) EMPTY_ARRAY, EMPTY_ARRAY);
        }

        @Override
        void visit(Deque<Map.Entry<K, V>> entries, Deque<Node<K, V>> nodes) {
            for (int i = 0; i < keys.length; ++i) {
                final Object sub = subNodes[i];
                if (sub instanceof Node) {
                    @SuppressWarnings("unchecked")
                    final Node<K, V> subNode = (Node<K, V>) sub;
                    assert keys[i] == null;
                    nodes.add(subNode);
                } else {
                    @SuppressWarnings("unchecked")
                    final V value = (V) sub;
                    entries.add(new AbstractMap.SimpleImmutableEntry<>(keys[i], value));
                }
            }
        }

        /**
         * For a given hash on 6 bits, its value is set if the bitmap has a one
         * at the corresponding index.
         */
        private boolean exists(int hash6) {
            return (mask & (1L << hash6)) != 0;
        }

        /**
         * For a given hash on 6 bits, the slot number is the number of one
         * bits on the right of the <code>hash6</code>-th bit.
         */
        private int slot(int hash6) {
            return Long.bitCount(mask & ((1L << hash6) - 1));
        }

        @Override
        V get(Object key, int hash) {
            final int hash6 = hash & HASH_MASK;
            if (!exists(hash6)) {
                return null;
            }
            final int slot = slot(hash6);
            final Object sub = subNodes[slot];
            assert sub != null;
            if (sub instanceof Node) {
                assert keys[slot] == null; // keys don't make sense on inner nodes
                @SuppressWarnings("unchecked")
                final Node<K, V> subNode = (Node<K, V>) sub;
                return subNode.get(key, hash >>> HASH_BITS);
            } else {
                if (keys[slot].equals(key)) {
                    @SuppressWarnings("unchecked")
                    final V v = (V) sub;
                    return v;
                } else {
                    // we have an entry for this hash, but the value is different
                    return null;
                }
            }
        }

        private Node<K, V> newSubNode(int hashBits) {
            if (hashBits <= 0) {
                return new Leaf<K, V>();
            } else {
                return new InnerNode<K, V>();
            }
        }

        private InnerNode<K, V> putExisting(K key, int hash, int hashBits, int slot, V value, MutableValueInt newValue) {
            final K[] keys2 = Arrays.copyOf(keys, keys.length);
            final Object[] subNodes2 = Arrays.copyOf(subNodes, subNodes.length);

            final Object previousValue = subNodes2[slot];
            if (previousValue instanceof Node) {
                // insert recursively
                assert keys[slot] == null;
                subNodes2[slot] = ((Node<K, V>) previousValue).put(key, hash, hashBits, value, newValue);
            } else if (keys[slot].equals(key)) {
                // replace the existing entry
                subNodes2[slot] = value;
            } else {
                // hash collision
                final K previousKey = keys[slot];
                final int previousHash = previousKey.hashCode() >>> (TOTAL_HASH_BITS - hashBits);
                Node<K, V> subNode = newSubNode(hashBits);
                subNode = subNode.put(previousKey, previousHash, hashBits, (V) previousValue, newValue);
                subNode = subNode.put(key, hash, hashBits, value, newValue);
                keys2[slot] = null;
                subNodes2[slot] = subNode;
            }
            return new InnerNode<>(mask, keys2, subNodes2);
        }

        private InnerNode<K, V> putNew(K key, int hash6, int slot, V value) {
            final long mask2 = mask | (1L << hash6);
            final K[] keys2 = insertElement(keys, key, slot);
            final Object[] subNodes2 = insertElement(subNodes, value, slot);
            return new InnerNode<>(mask2, keys2, subNodes2);
        }

        @Override
        InnerNode<K, V> put(K key, int hash, int hashBits, V value, MutableValueInt newValue) {
            final int hash6 = hash & HASH_MASK;
            final int slot = slot(hash6);

            if (exists(hash6)) {
                hash >>>= HASH_BITS;
                hashBits -= HASH_BITS;
                return putExisting(key, hash, hashBits, slot, value, newValue);
            } else {
                newValue.value = 1;
                return putNew(key, hash6, slot, value);
            }
        }

        private InnerNode<K, V> removeSlot(int hash6, int slot) {
            final long mask2 = mask & ~(1L << hash6);
            final K[] keys2 = removeArrayElement(keys, slot);
            final Object[] subNodes2 = removeArrayElement(subNodes, slot);
            return new InnerNode<>(mask2, keys2, subNodes2);
        }

        @Override
        InnerNode<K, V> remove(Object key, int hash) {
            final int hash6 = hash & HASH_MASK;
            if (!exists(hash6)) {
                return this;
            }
            final int slot = slot(hash6);
            final Object previousValue = subNodes[slot];
            if (previousValue instanceof Node) {
                @SuppressWarnings("unchecked")
                final Node<K, V> subNode = (Node<K, V>) previousValue;
                final Node<K, V> removed = subNode.remove(key, hash >>> HASH_BITS);
                if (removed == subNode) {
                    // not in sub-nodes
                    return this;
                }
                if (removed.isEmpty()) {
                    return removeSlot(hash6, slot);
                }
                final K[] keys2 = Arrays.copyOf(keys, keys.length);
                final Object[] subNodes2 = Arrays.copyOf(subNodes, subNodes.length);
                subNodes2[slot] = removed;
                return new InnerNode<>(mask, keys2, subNodes2);
            } else if (keys[slot].equals(key)) {
                // remove entry
                return removeSlot(hash6, slot);
            } else {
                // hash collision, nothing to remove
                return this;
            }
        }

    }

    /**
     * Iterates over an entry
     *
     * @opensearch.internal
     */
    private static class EntryIterator<K, V> implements Iterator<Map.Entry<K, V>> {

        private final Deque<Map.Entry<K, V>> entries;
        private final Deque<Node<K, V>> nodes;

        EntryIterator(Node<K, V> node) {
            entries = new ArrayDeque<>();
            nodes = new ArrayDeque<>();
            node.visit(entries, nodes);
        }

        @Override
        public boolean hasNext() {
            return !entries.isEmpty() || !nodes.isEmpty();
        }

        @Override
        public Map.Entry<K, V> next() {
            while (entries.isEmpty()) {
                if (nodes.isEmpty()) {
                    throw new NoSuchElementException();
                }
                final Node<K, V> nextNode = nodes.pop();
                nextNode.visit(entries, nodes);
            }
            return entries.pop();
        }

        @Override
        public final void remove() {
            throw new UnsupportedOperationException();
        }

    }

    private final InnerNode<K, V> root;
    private final int size;

    /**
     * Create a new empty map.
     */
    public CopyOnWriteHashMap() {
        this(new InnerNode<K, V>(), 0);
    }

    private CopyOnWriteHashMap(InnerNode<K, V> root, int size) {
        this.root = root;
        this.size = size;
    }

    @Override
    public boolean containsKey(Object key) {
        // works fine since null values are not supported
        return get(key) != null;
    }

    @Override
    public V get(Object key) {
        if (key == null) {
            throw new IllegalArgumentException("null keys are not supported");
        }
        final int hash = key.hashCode();
        return root.get(key, hash);
    }

    @Override
    public int size() {
        assert size != 0 || root.isEmpty();
        return size;
    }

    /**
     * Associate <code>key</code> with <code>value</code> and return a new copy
     * of the hash table. The current hash table is not modified.
     */
    public CopyOnWriteHashMap<K, V> copyAndPut(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("null keys are not supported");
        }
        if (value == null) {
            throw new IllegalArgumentException("null values are not supported");
        }
        final int hash = key.hashCode();
        final MutableValueInt newValue = new MutableValueInt();
        final InnerNode<K, V> newRoot = root.put(key, hash, TOTAL_HASH_BITS, value, newValue);
        final int newSize = size + newValue.value;
        return new CopyOnWriteHashMap<>(newRoot, newSize);
    }

    /**
     * Same as {@link #copyAndPut(Object, Object)} but for an arbitrary number of entries.
     */
    public CopyOnWriteHashMap<K, V> copyAndPutAll(Map<? extends K, ? extends V> other) {
        return copyAndPutAll(other.entrySet());
    }

    public <K1 extends K, V1 extends V> CopyOnWriteHashMap<K, V> copyAndPutAll(Iterable<Entry<K1, V1>> entries) {
        CopyOnWriteHashMap<K, V> result = this;
        for (Entry<K1, V1> entry : entries) {
            result = result.copyAndPut(entry.getKey(), entry.getValue());
        }
        return result;
    }

    public <K1 extends K, V1 extends V> CopyOnWriteHashMap<K, V> copyAndPutAll(Stream<Entry<K1, V1>> entries) {
        return copyAndPutAll(entries::iterator);
    }

    /**
     * Remove the given key from this map. The current hash table is not modified.
     */
    public CopyOnWriteHashMap<K, V> copyAndRemove(Object key) {
        if (key == null) {
            throw new IllegalArgumentException("null keys are not supported");
        }
        final int hash = key.hashCode();
        final InnerNode<K, V> newRoot = root.remove(key, hash);
        if (root == newRoot) {
            return this;
        } else {
            return new CopyOnWriteHashMap<>(newRoot, size - 1);
        }
    }

    /**
     * Same as {@link #copyAndRemove(Object)} but for an arbitrary number of entries.
     */
    public CopyOnWriteHashMap<K, V> copyAndRemoveAll(Collection<?> keys) {
        CopyOnWriteHashMap<K, V> result = this;
        for (Object key : keys) {
            result = result.copyAndRemove(key);
        }
        return result;
    }

    @Override
    public Set<Map.Entry<K, V>> entrySet() {
        return new AbstractSet<Map.Entry<K, V>>() {

            @Override
            public Iterator<java.util.Map.Entry<K, V>> iterator() {
                return new EntryIterator<>(root);
            }

            @Override
            public boolean contains(Object o) {
                if (o == null || !(o instanceof Map.Entry)) {
                    return false;
                }
                Map.Entry<?, ?> entry = (java.util.Map.Entry<?, ?>) o;
                return entry.getValue().equals(CopyOnWriteHashMap.this.get(entry.getKey()));
            }

            @Override
            public int size() {
                return CopyOnWriteHashMap.this.size();
            }
        };
    }

}