package net.ranides.assira.collection.maps;

import lombok.AllArgsConstructor;

import net.ranides.assira.collection.sets.ASet;
import net.ranides.assira.generic.Ref;
import net.ranides.assira.generic.Pair;

import java.io.Serializable;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Set;

/**
 * Specialized CrossMap implementation where both keys are enum types.
 *
 * @param <K1> type of 1st key
 * @param <K2> type of 2nd key
 * @param <V> type of value
 */
public class EnumCrossMap<K1 extends Enum<K1>, K2 extends Enum<K2>, V> implements CrossMap<K1, K2, V>, Serializable {

    private static final long serialVersionUID = 1L;

    private static final Cache<Class<? extends Enum<?>>, Object> KEYS_CACHE = Cache.getInstance();

    private final Class<K1> type1;
    private final Class<K2> type2;
    private final K1[] universe1;
    private final K2[] universe2;
    private final Object[][] values;

    private int size;

    /**
     * Creates new empty CrossMap with specified types as keys.
     *
     * @param type1 type1
     * @param type2 type2
     */
    @SuppressWarnings("unchecked")
    public EnumCrossMap(Class<K1> type1, Class<K2> type2) {
        this.type1 = type1;
        this.type2 = type2;
        this.universe1 = (K1[])getSharedKeys(type1);
        this.universe2 = (K2[])getSharedKeys(type2);
        this.values = new Object[universe1.length][universe2.length];
        this.size = 0;
    }

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

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

    @Override
    public boolean containsKey(K1 key1, K2 key2) {
        return getOptional(key1, key2).isPresent();
    }

    @Override
    public boolean containsValue(V value) {
        for(int i1=0, n1=universe1.length; i1<n1; i1++) {
            for(int i2=0, n2=universe2.length; i2<n2; i2++) {
                if(Objects.equals(values[i1][i2], value)) {
                    return true;
                }
            }
        }
        return false;
    }

    @Override
    public V get(K1 key1, K2 key2) {
        return getValueAt(key1.ordinal(), key2.ordinal());
    }

    @Override
    public Ref<V> getOptional(K1 key1, K2 key2) {
        return Ref.ofNullable(get(key1, key2));
    }

    @Override
    public V put(K1 key1, K2 key2, V value) {
        return setValueAt(key1.ordinal(), key2.ordinal(), value);
    }

    @Override
    public V remove(K1 key1, K2 key2) {
        return setValueAt(key1.ordinal(), key2.ordinal(), null);
    }

    @SuppressWarnings("unchecked")
    protected V getValueAt(int index1, int index2) {
        return (V)values[index1][index2];
    }

    protected V setValueAt(int index1, int index2, V value) {
        V prev = getValueAt(index1, index2);

        if(value == null) {
            if(prev != null) {
                size--;
            }
        } else {
            if(prev == null) {
                size++;
            }
        }

        values[index1][index2] = value;
        return prev;
    }

    @Override
    public void putAll(CrossMap<K1, K2, ? extends V> values) {
        values.forEachPair((pair, value) -> put(pair.getFirst(), pair.getSecond(), value));
    }

    @Override
    public void clear() {
        for(int i1=0, n1=universe1.length; i1<n1; i1++) {
            for(int i2=0, n2=universe2.length; i2<n2; i2++) {
                values[i1][i2] = null;
            }
        }
        size = 0;
    }

    @Override
    public MultiMap<K2, V> reduceFirst() {
        OpenMultiMap<K2, V> result = new OpenMultiMap<>();
        forEach((key1, key2, value) -> result.put(key2, value));
        return result;
    }

    @Override
    public MultiMap<K1, V> reduceSecond() {
        OpenMultiMap<K1, V> result = new OpenMultiMap<>();
        forEach((key1, key2, value) -> result.put(key1, value));
        return result;
    }

    @Override
    public MultiMap<K1, K2> reduceValues() {
        OpenMultiMap<K1, K2> result = new OpenMultiMap<>();
        forEach((key1, key2, value) -> result.put(key1, key2));
        return result;
    }

    @Override
    public Set<Pair<K1, K2>> keySet() {
        return new Keys();
    }

    @Override
    public Collection<V> values() {
        return new Values();
    }

    @Override
    public Set<Map.Entry<Pair<K1, K2>, V>> entrySet() {
        return new Entries();
    }

    @Override
    public String toString() {
        return entrySet().toString();
    }

    private static Object getSharedKeys(Class<? extends Enum<?>> type) {
        return KEYS_CACHE.get(type, Class::getEnumConstants);
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private Ref<Map.Entry<Pair<K1, K2>,V>> castEntry(Object o) {
        if (o instanceof Map.Entry<?, ?>) {
            Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
            return castPair(entry.getKey()).map(x -> (Map.Entry)o);
        }
        return Ref.empty();
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private Ref<Pair<K1, K2>> castPair(Object o) {
            if (o instanceof Pair<?, ?>) {
                Pair<?, ?> pair = (Pair<?, ?>) o;
                Object key1 = pair.getFirst();
                Object key2 = pair.getSecond();
                if(type1.isInstance(key1) && type2.isInstance(key2)) {
                    return Ref.of((Pair)pair);
            }
        }
        return Ref.empty();
    }

    private class Keys extends ASet<Pair<K1, K2>> {
        @Override
        public Iterator<Pair<K1, K2>> iterator() {
            return new KeyIterator();
        }

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

        @Override
        public boolean contains(Object value) {
            return castPair(value)
                .map(pair -> EnumCrossMap.this.containsKey(pair.getFirst(), pair.getSecond()))
                .orElse(false);
        }

        @Override
        public boolean remove(Object value) {
            return castPair(value)
                .map(pair -> EnumCrossMap.this.remove(pair.getFirst(), pair.getSecond()) != null)
                .orElse(false);
        }


    }

    @SuppressWarnings("unchecked")
    private class Values extends ASet<V> {
        @Override
        public Iterator<V> iterator() {
            return new ValueIterator();
        }

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

        @Override
        public boolean contains(Object value) {
            return containsValue((V)value);
        }
    }

    private class Entries extends ASet<Map.Entry<Pair<K1, K2>, V>> {
        @Override
        public Iterator<Map.Entry<Pair<K1, K2>, V>> iterator() {
            return new EntryIterator();
        }

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

        @Override
        public boolean contains(Object o) {
            return castEntry(o)
                .filter(entry -> {
                    Pair<K1, K2> pair = entry.getKey();
                    V prev = EnumCrossMap.this.get(pair.getFirst(), pair.getSecond());
                    return Objects.equals(prev, entry.getValue());
                })
                .isPresent();
        }

        @Override
        public boolean remove(Object o) {
            return castEntry(o)
                .filter(entry -> {
                    Pair<K1, K2> pair = entry.getKey();
                    V prev = EnumCrossMap.this.get(pair.getFirst(), pair.getSecond());
                    return Objects.equals(prev, entry.getValue());
                })
                .map(entry -> {
                    Pair<K1, K2> pair = entry.getKey();
                    EnumCrossMap.this.remove(pair.getFirst(), pair.getSecond());
                    return true;
                })
                .orElse(false);
        }

    }

    private class ValueIterator extends CrossIterator<V> {
        @Override
        public V produce() {
            return getValueAt(index1, index2);
        }
    }

    private class KeyIterator extends CrossIterator<Pair<K1, K2>> {
        @Override
        public Pair<K1, K2> produce() {
            return Pair.of(universe1[index1], universe2[index2]);
        }
    }

    private class EntryIterator extends CrossIterator<Map.Entry<Pair<K1, K2>, V>> {
        @Override
        public Map.Entry<Pair<K1, K2>, V> produce() {
            return new CrossEntry(index1, index2);
        }
    }

    @AllArgsConstructor
    private class CrossEntry implements Map.Entry<Pair<K1, K2>, V> {
        private final int index1;
        private final int index2;

        @Override
        public Pair<K1, K2> getKey() {
            return Pair.of(universe1[index1], universe2[index2]);
        }

        @Override
        public V getValue() {
            return getValueAt(index1, index2);
        }

        @Override
        public V setValue(V value) {
            V prev = getValueAt(index1, index2);
            values[index1][index2] = value;
            return prev;
        }
    }

    private abstract class CrossIterator<Q> implements Iterator<Q> {
        protected int visited;
        protected int index1;
        protected int index2;

        public abstract Q produce();

        @Override
        public boolean hasNext() {
            return visited < size;
        }

        @Override
        public Q next() {
            while (inc()) {}
            visited++;
            return produce();
        }

        private boolean inc() {
            index1++;
            if(index1 >= universe1.length) {
                index1 = 0;
                index2++;
            }
            if(index2 >= universe2.length) {
                throw new NoSuchElementException();
            }
            return values[index1][index2] != null;
        }

        @Override
        public void remove() {
            if(values[index1][index2] != null) {
                values[index1][index2] = null;
                size--;
            }
        }
    }

}
