/*
 * @author Ranides Atterwim <ranides@gmail.com>
 * @copyright Ranides Atterwim
 * @license WTFPL
 * @url http://ranides.net/projects/assira
 */

package net.ranides.assira.collection.maps;

import java.util.EnumSet;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

/**
 *
 * @author Ranides Atterwim <ranides@gmail.com>
 */
public final class MapCollectors {
    
    private MapCollectors() {
        /* utility class */
    }

    public static <R,K,V> Collector<R,?, Map<K,R>> unique(Function<? super R, ? extends K> key) {
        return MapCollectors.unique(key, v -> v);
    }
    
    public static <R,K,V> Collector<R,?, Map<K,V>> unique(Function<? super R, ? extends K> key, Function<R, V> val) {
        return MapCollectors.unique(key, val, () -> new HashMap<>());
    }
    
    public static <R,K,V> Collector<R,?, Map<K,V>> unique(Function<? super R, ? extends K> key, Function<? super R, ? extends V> val, Supplier<Map<K, V>> supplier) {
        return new UCollector<>(key, val, supplier);
    }
    
    public static <R,K,V> Collector<R,?, MultiMap<K,R>> multimap(Function<? super R, ? extends K> key) {
        return MapCollectors.multimap(key, v -> v);
    }
    
    public static <R,K,V> Collector<R,?, MultiMap<K,V>> multimap(Function<? super R, ? extends K> key, Function<R, V> val) {
        return MapCollectors.multimap(key, val, () -> new HashMultiMap<>());
    }
    
    public static <R,K,V> Collector<R,?, MultiMap<K,V>> multimap(Function<? super R, ? extends K> key, Function<? super R, ? extends V> val, Supplier<MultiMap<K, V>> supplier) {
        return new MCollector<>(key, val, supplier);
    }
    
    static IllegalArgumentException duplicate(Object key, Object previous, Object current) {
        return new IllegalArgumentException("Duplicate values (" + previous + "," + current + ") for key: " + key);
    }
    
    private static class UCollector<R,K,V> implements Collector<R,Map<K,V>,Map<K,V>> {
        
        private final Function<? super R, ? extends K> key;
        private final Function<? super R, ? extends V> val;
        private final Supplier<Map<K, V>> supplier;

        public UCollector(Function<? super R, ? extends K> key, Function<? super R, ? extends V> val, Supplier<Map<K, V>> supplier) {
            this.key = key;
            this.val = val;
            this.supplier = supplier;
        }
        
        @Override
        public Supplier<Map<K, V>> supplier() {
            return supplier;
        }

        @Override
        public BiConsumer<Map<K, V>, R> accumulator() {
            return this::$accumulate;
        }
        
        private void $accumulate(Map<K, V> map, R r) {
            K ik = key.apply(r);
            V iv = val.apply(r);
            if(map.containsKey(ik)) {
                throw duplicate(ik, map.get(ik), iv);
            }
            map.put(ik, iv);
        }

        @Override
        public BinaryOperator<Map<K, V>> combiner() {
            return this::$combine;
        }
        
        private Map<K, V> $combine(Map<K, V> target, Map<K, V> source) {
            for(Map.Entry<K,V> item : source.entrySet()) {
                K ik = item.getKey();
                V iv = item.getValue();
                if(target.containsKey(ik)) {
                    throw duplicate(ik, target.get(ik), iv);
                }
            }
            return target;
        }

        @Override
        public Function<Map<K, V>, Map<K, V>> finisher() {
            return m -> m;
        }

        @Override
        public Set<Collector.Characteristics> characteristics() {
            return EnumSet.of(Collector.Characteristics.IDENTITY_FINISH);
        }

    }
    
    private static class MCollector<R,K,V> implements Collector<R,MultiMap<K,V>,MultiMap<K,V>> {
        
        private final Function<? super R, ? extends K> key;
        private final Function<? super R, ? extends V> val;
        private final Supplier<MultiMap<K, V>> supplier;

        public MCollector(Function<? super R, ? extends K> key, Function<? super R, ? extends V> val, Supplier<MultiMap<K, V>> supplier) {
            this.key = key;
            this.val = val;
            this.supplier = supplier;
        }
        
        @Override
        public Supplier<MultiMap<K, V>> supplier() {
            return supplier;
        }

        @Override
        public BiConsumer<MultiMap<K, V>, R> accumulator() {
            return (map, r) -> map.put(key.apply(r), val.apply(r));
        }

        @Override
        public BinaryOperator<MultiMap<K, V>> combiner() {
            return (a,b) -> { a.putAll(b); return a; };
        }

        @Override
        public Function<MultiMap<K, V>, MultiMap<K, V>> finisher() {
            return m -> m;
        }

        @Override
        public Set<Collector.Characteristics> characteristics() {
            return EnumSet.of(Collector.Characteristics.IDENTITY_FINISH);
        }

    }
    
}
