/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.
 */
package org.apache.logging.log4j.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InvalidObjectException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.StreamCorruptedException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.Arrays;
import java.util.Collection;
import java.util.ConcurrentModificationException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

import org.apache.logging.log4j.status.StatusLogger;

/**
 * <em>Consider this class private.</em>
 * Array-based implementation of the {@code ReadOnlyStringMap} interface. Keys are held in a sorted array.
 * <p>
 * This is not a generic collection, but makes some trade-offs to optimize for the Log4j context data use case:
 * </p>
 * <ul>
 *   <li>Garbage-free iteration over key-value pairs with {@code BiConsumer} and {@code TriConsumer}.</li>
 *   <li>Fast copy. If the ThreadContextMap is also an instance of {@code SortedArrayStringMap}, the full thread context
 *     data can be transferred with two array copies and two field updates.</li>
 *   <li>Acceptable performance for small data sets. The current implementation stores keys in a sorted array, values
 *     are stored in a separate array at the same index.
 *     Worst-case performance of {@code get} and {@code containsKey} is O(log N),
 *     worst-case performance of {@code put} and {@code remove} is O(N log N).
 *     The expectation is that for the small values of {@code N} (less than 100) that are the vast majority of
 *     ThreadContext use cases, the constants dominate performance more than the asymptotic performance of the
 *     algorithms used.
 *     </li>
 *     <li>Compact representation.</li>
 * </ul>
 *
 * @since 2.7
 */
public class SortedArrayStringMap implements IndexedStringMap {

    /**
     * The default initial capacity.
     */
    private static final int DEFAULT_INITIAL_CAPACITY = 4;
    private static final long serialVersionUID = -5748905872274478116L;
    private static final int HASHVAL = 31;

    private static final TriConsumer<String, Object, StringMap> PUT_ALL = (key, value, contextData) -> contextData.putValue(key, value);

    /**
     * An empty array instance to share when the table is not inflated.
     */
    private static final String[] EMPTY = {};
    private static final String FROZEN = "Frozen collection cannot be modified";

    private transient String[] keys = EMPTY;
    private transient Object[] values = EMPTY;

    /**
     * The number of key-value mappings contained in this map.
     */
    private transient int size;

    private static final Method setObjectInputFilter;
    private static final Method getObjectInputFilter;
    private static final Method newObjectInputFilter;

    static {
        Method[] methods = ObjectInputStream.class.getMethods();
        Method setMethod = null;
        Method getMethod = null;
        for (final Method method : methods) {
            if (method.getName().equals("setObjectInputFilter")) {
                setMethod = method;
            } else if (method.getName().equals("getObjectInputFilter")) {
                getMethod = method;
            }
        }
        Method newMethod = null;
        try {
            if (setMethod != null) {
                final Class<?> clazz = Class.forName("org.apache.logging.log4j.internal.DefaultObjectInputFilter");
                methods = clazz.getMethods();
                for (final Method method : methods) {
                    if (method.getName().equals("newInstance") && Modifier.isStatic(method.getModifiers())) {
                        newMethod = method;
                        break;
                    }
                }
            }
        } catch (final ClassNotFoundException ex) {
            // Ignore the exception
        }
        newObjectInputFilter = newMethod;
        setObjectInputFilter = setMethod;
        getObjectInputFilter = getMethod;
    }

    /**
     * The next size value at which to resize (capacity * load factor).
     * @serial
     */
    // If table == EMPTY_TABLE then this is the initial capacity at which the
    // table will be created when inflated.
    private int threshold;
    private boolean immutable;
    private transient boolean iterating;

    public SortedArrayStringMap() {
        this(DEFAULT_INITIAL_CAPACITY);
    }

    public SortedArrayStringMap(final int initialCapacity) {
        if (initialCapacity < 0) {
            throw new IllegalArgumentException("Initial capacity must be at least zero but was " + initialCapacity);
        }
        threshold = ceilingNextPowerOfTwo(initialCapacity == 0 ? 1 : initialCapacity);
    }

    public SortedArrayStringMap(final ReadOnlyStringMap other) {
        if (other instanceof SortedArrayStringMap) {
            initFrom0((SortedArrayStringMap) other);
        } else if (other != null) {
            resize(ceilingNextPowerOfTwo(other.size()));
            other.forEach(PUT_ALL, this);
        }
    }

    public SortedArrayStringMap(final Map<String, ?> map) {
        resize(ceilingNextPowerOfTwo(map.size()));
        for (final Map.Entry<String, ?> entry : map.entrySet()) {
            putValue(entry.getKey(), entry.getValue());
        }
    }

    /**
     * For unit testing.
     * @return The threshold.
     */
    public int getThreshold() {
        return threshold;
    }

    private void assertNotFrozen() {
        if (immutable) {
            throw new UnsupportedOperationException(FROZEN);
        }
    }

    private void assertNoConcurrentModification() {
        if (iterating) {
            throw new ConcurrentModificationException();
        }
    }

    @Override
    public void clear() {
        if (keys == EMPTY) {
            return;
        }
        assertNotFrozen();
        assertNoConcurrentModification();

        Arrays.fill(keys, 0, size, null);
        Arrays.fill(values, 0, size, null);
        size = 0;
    }

    @Override
    public boolean containsKey(final String key) {
        return indexOfKey(key) >= 0;
    }

    @Override
    public Map<String, String> toMap() {
        final Map<String, String> result = new HashMap<>(size());
        for (int i = 0; i < size(); i++) {
            final Object value = getValueAt(i);
            result.put(getKeyAt(i), value == null ? null : String.valueOf(value));
        }
        return result;
    }

    @Override
    public void freeze() {
        immutable = true;
    }

    @Override
    public boolean isFrozen() {
        return immutable;
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> V getValue(final String key) {
        final int index = indexOfKey(key);
        if (index < 0) {
            return null;
        }
        return (V) values[index];
    }

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

    @Override
    public int indexOfKey(final String key) {
        if (keys == EMPTY) {
            return -1;
        }
        if (key == null) { // null key is located at the start of the array
            return nullKeyIndex(); // insert at index zero
        }
        final int start = size > 0 && keys[0] == null ? 1 : 0;
        return Arrays.binarySearch(keys, start, size, key);
    }

    private int nullKeyIndex() {
        return size > 0 && keys[0] == null ? 0 : ~0;
    }

    @Override
    public void putValue(final String key, final Object value) {
        assertNotFrozen();
        assertNoConcurrentModification();

        if (keys == EMPTY) {
            inflateTable(threshold);
        }
        final int index = indexOfKey(key);
        if (index >= 0) {
            keys[index] = key;
            values[index] = value;
        } else { // not found, so insert.
            insertAt(~index, key, value);
        }
    }

    private void insertAt(final int index, final String key, final Object value) {
        ensureCapacity();
        System.arraycopy(keys, index, keys, index + 1, size - index);
        System.arraycopy(values, index, values, index + 1, size - index);
        keys[index] = key;
        values[index] = value;
        size++;
    }

    @Override
    public void putAll(final ReadOnlyStringMap source) {
        if (source == this || source == null || source.isEmpty()) {
            // this.putAll(this) does not modify this collection
            // this.putAll(null) does not modify this collection
            // this.putAll(empty ReadOnlyStringMap) does not modify this collection
            return;
        }
        assertNotFrozen();
        assertNoConcurrentModification();

        if (source instanceof SortedArrayStringMap) {
            if (this.size == 0) {
                initFrom0((SortedArrayStringMap) source);
            } else {
                merge((SortedArrayStringMap) source);
            }
        } else if (source != null) {
            source.forEach(PUT_ALL, this);
        }
    }

    private void initFrom0(final SortedArrayStringMap other) {
        if (keys.length < other.size) {
            keys = new String[other.threshold];
            values = new Object[other.threshold];
        }
        System.arraycopy(other.keys, 0, keys, 0, other.size);
        System.arraycopy(other.values, 0, values, 0, other.size);

        size = other.size;
        threshold = other.threshold;
    }

    private void merge(final SortedArrayStringMap other) {
        final String[] myKeys = keys;
        final Object[] myVals = values;
        final int newSize = other.size + this.size;
        threshold = ceilingNextPowerOfTwo(newSize);
        if (keys.length < threshold) {
            keys = new String[threshold];
            values = new Object[threshold];
        }
        // move largest collection to the beginning of this data structure, smallest to the end
        boolean overwrite = true;
        if (other.size() > size()) {
            // move original key-values to end
            System.arraycopy(myKeys, 0, keys, other.size, this.size);
            System.arraycopy(myVals, 0, values, other.size, this.size);

            // insert additional key-value pairs at the beginning
            System.arraycopy(other.keys, 0, keys, 0, other.size);
            System.arraycopy(other.values, 0, values, 0, other.size);
            size = other.size;

            // loop over original keys and insert (drop values for same key)
            overwrite = false;
        } else {
            System.arraycopy(myKeys, 0, keys, 0, this.size);
            System.arraycopy(myVals, 0, values, 0, this.size);

            // move additional key-value pairs to end
            System.arraycopy(other.keys, 0, keys, this.size, other.size);
            System.arraycopy(other.values, 0, values, this.size, other.size);

            // new values are at the end, will be processed below. Overwrite is true.
        }
        for (int i = size; i < newSize; i++) {
            final int index = indexOfKey(keys[i]);
            if (index < 0) { // key does not exist
                insertAt(~index, keys[i], values[i]);
            } else if (overwrite) { // existing key: only overwrite when looping over the new values
                keys[index] = keys[i];
                values[index] = values[i];
            }
        }
        // prevent memory leak: null out references
        Arrays.fill(keys, size, newSize, null);
        Arrays.fill(values, size, newSize, null);
    }

    private void ensureCapacity() {
        if (size >= threshold) {
            resize(threshold * 2);
        }
    }

    private void resize(final int newCapacity) {
        final String[] oldKeys = keys;
        final Object[] oldValues = values;

        keys = new String[newCapacity];
        values = new Object[newCapacity];

        System.arraycopy(oldKeys, 0, keys, 0, size);
        System.arraycopy(oldValues, 0, values, 0, size);

        threshold = newCapacity;
    }

    /**
     * Inflates the table.
     */
    private void inflateTable(final int toSize) {
        threshold = toSize;
        keys = new String[toSize];
        values = new Object[toSize];
    }

    @Override
    public void remove(final String key) {
        if (keys == EMPTY) {
            return;
        }

        final int index = indexOfKey(key);
        if (index >= 0) {
            assertNotFrozen();
            assertNoConcurrentModification();

            System.arraycopy(keys, index + 1, keys, index, size - 1 - index);
            System.arraycopy(values, index + 1, values, index, size - 1 - index);
            keys[size - 1] = null;
            values[size - 1] = null;
            size--;
        }
    }

    @Override
    public String getKeyAt(final int index) {
        if (index < 0 || index >= size) {
            return null;
        }
        return keys[index];
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> V getValueAt(final int index) {
        if (index < 0 || index >= size) {
            return null;
        }
        return (V) values[index];
    }

    @Override
    public int size() {
        return size;
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V> void forEach(final BiConsumer<String, ? super V> action) {
        iterating = true;
        try {
            for (int i = 0; i < size; i++) {
                action.accept(keys[i], (V) values[i]);
            }
        } finally {
            iterating = false;
        }
    }

    @SuppressWarnings("unchecked")
    @Override
    public <V, T> void forEach(final TriConsumer<String, ? super V, T> action, final T state) {
        iterating = true;
        try {
            for (int i = 0; i < size; i++) {
                action.accept(keys[i], (V) values[i], state);
            }
        } finally {
            iterating = false;
        }
    }

    @Override
    public boolean equals(final Object obj) {
        if (obj == this) {
            return true;
        }
        if (!(obj instanceof SortedArrayStringMap)) {
            return false;
        }
        final SortedArrayStringMap other = (SortedArrayStringMap) obj;
        if (this.size() != other.size()) {
            return false;
        }
        for (int i = 0; i < size(); i++) {
            if (!Objects.equals(keys[i], other.keys[i])) {
                return false;
            }
            if (!Objects.equals(values[i], other.values[i])) {
                return false;
            }
        }
        return true;
    }

    @Override
    public int hashCode() {
        int result = 37;
        result = HASHVAL * result + size;
        result = HASHVAL * result + hashCode(keys, size);
        result = HASHVAL * result + hashCode(values, size);
        return result;
    }

    private static int hashCode(final Object[] values, final int length) {
        int result = 1;
        for (int i = 0; i < length; i++) {
            result = HASHVAL * result + (values[i] == null ? 0 : values[i].hashCode());
        }
        return result;
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder(256);
        sb.append('{');
        for (int i = 0; i < size; i++) {
            if (i > 0) {
                sb.append(", ");
            }
            sb.append(keys[i]).append('=');
            sb.append(values[i] == this ? "(this map)" : values[i]);
        }
        sb.append('}');
        return sb.toString();
    }

    /**
     * Save the state of the {@code SortedArrayStringMap} instance to a stream (i.e.,
     * serialize it).
     *
     * @param s The ObjectOutputStream.
     * @throws IOException if there is an error serializing the object to the stream.
     *
     * @serialData The <i>capacity</i> of the SortedArrayStringMap (the length of the
     *             bucket array) is emitted (int), followed by the
     *             <i>size</i> (an int, the number of key-value
     *             mappings), followed by the key (Object) and value (Object)
     *             for each key-value mapping.  The key-value mappings are
     *             emitted in no particular order.
     */
    private void writeObject(final java.io.ObjectOutputStream s) throws IOException {
        // Write out the threshold, and any hidden stuff
        s.defaultWriteObject();

        // Write out number of buckets
        if (keys == EMPTY) {
            s.writeInt(ceilingNextPowerOfTwo(threshold));
        } else {
            s.writeInt(keys.length);
        }

        // Write out size (number of Mappings)
        s.writeInt(size);

        // Write out keys and values (alternating)
        if (size > 0) {
            for (int i = 0; i < size; i++) {
                s.writeObject(keys[i]);
                try {
                    s.writeObject(marshall(values[i]));
                } catch (final Exception e) {
                    handleSerializationException(e, i, keys[i]);
                    s.writeObject(null);
                }
            }
        }
    }

    private static byte[] marshall(final Object obj) throws IOException {
        if (obj == null) {
            return null;
        }
        final ByteArrayOutputStream bout = new ByteArrayOutputStream();
        try (final ObjectOutputStream oos = new ObjectOutputStream(bout)) {
            oos.writeObject(obj);
            oos.flush();
            return bout.toByteArray();
        }
    }

    @SuppressWarnings("BanSerializableRead")
    private static Object unmarshall(final byte[] data, final ObjectInputStream inputStream)
            throws IOException, ClassNotFoundException {
        final ByteArrayInputStream bin = new ByteArrayInputStream(data);
        Collection<String> allowedClasses = null;
        final ObjectInputStream ois;
        if (inputStream instanceof FilteredObjectInputStream) {
            allowedClasses = ((FilteredObjectInputStream) inputStream).getAllowedClasses();
            ois = new FilteredObjectInputStream(bin, allowedClasses);
        } else {
            try {
                final Object obj = getObjectInputFilter.invoke(inputStream);
                final Object filter = newObjectInputFilter.invoke(null, obj);
                ois = new ObjectInputStream(bin);
                setObjectInputFilter.invoke(ois, filter);
            } catch (final IllegalAccessException | InvocationTargetException ex) {
                throw new StreamCorruptedException("Unable to set ObjectInputFilter on stream");
            }
        }
        try {
            return ois.readObject();
        } finally {
            ois.close();
        }
    }

    /**
     * Calculate the next power of 2, greater than or equal to x.
     * <p>
     * From Hacker's Delight, Chapter 3, Harry S. Warren Jr.
     *
     * @param x Value to round up
     * @return The next power of 2 from x inclusive
     */
    private static int ceilingNextPowerOfTwo(final int x) {
        final int BITS_PER_INT = 32;
        return 1 << (BITS_PER_INT - Integer.numberOfLeadingZeros(x - 1));
    }

    /**
     * Reconstitute the {@code SortedArrayStringMap} instance from a stream (i.e.,
     * deserialize it).
     * @param s The ObjectInputStream.
     * @throws IOException If there is an error reading the input stream.
     * @throws ClassNotFoundException if the class to be instantiated could not be found.
     */
    private void readObject(final java.io.ObjectInputStream s)  throws IOException, ClassNotFoundException {
        if (!(s instanceof FilteredObjectInputStream) && setObjectInputFilter == null) {
            throw new IllegalArgumentException("readObject requires a FilteredObjectInputStream or an ObjectInputStream that accepts an ObjectInputFilter");
        }
        // Read in the threshold (ignored), and any hidden stuff
        s.defaultReadObject();

        // set other fields that need values
        keys = EMPTY;
        values = EMPTY;

        // Read in number of buckets
        final int capacity = s.readInt();
        if (capacity < 0) {
            throw new InvalidObjectException("Illegal capacity: " + capacity);
        }

        // Read number of mappings
        final int mappings = s.readInt();
        if (mappings < 0) {
            throw new InvalidObjectException("Illegal mappings count: " + mappings);
        }

        // allocate the bucket array;
        if (mappings > 0) {
            inflateTable(capacity);
        } else {
            threshold = capacity;
        }

        // Read the keys and values, and put the mappings in the arrays
        for (int i = 0; i < mappings; i++) {
            keys[i] = (String) s.readObject();
            try {
                final byte[] marshalledObject = (byte[]) s.readObject();
                values[i] = marshalledObject == null ? null : unmarshall(marshalledObject, s);
            } catch (final Exception | LinkageError error) {
                handleSerializationException(error, i, keys[i]);
                values[i] = null;
            }
        }
        size = mappings;
    }

    private void handleSerializationException(final Throwable t, final int i, final String key) {
        StatusLogger.getLogger().warn("Ignoring {} for key[{}] ('{}')", String.valueOf(t), i, keys[i]);
    }
}