package org.shedlang.compiler.util;

import org.pcollections.*;

import java.util.Map;
import java.util.stream.Collector;

public class PCollectors {
    private PCollectors() {
    }

    public static <K, V> Collector<Map.Entry<K, V>, PMapCollector<K, V>, PMap<K, V>> toMap() {
        return Collector.of(
            PMapCollector::new,
            PMapCollector::add,
            PMapCollector::addAll,
            PMapCollector::toMap
        );
    }

    private static class PMapCollector<K, V> {
        private PMap<K, V> map = Empty.map();

        void add(Map.Entry<K, V> entry) {
            map = map.plus(entry.getKey(), entry.getValue());
        }

        PMapCollector<K, V> addAll(PMapCollector<K, V> other) {
            map = map.plusAll(other.map);
            return this;
        }

        PMap<K, V> toMap() {
            return map;
        }
    }

    public static <K, V> Collector<Map.Entry<K, V>, POrderedMapCollector<K, V>, POrderedMap<K, V>> toOrderedMap() {
        return Collector.of(
            POrderedMapCollector::new,
            POrderedMapCollector::add,
            POrderedMapCollector::addAll,
            POrderedMapCollector::toOrderedMap
        );
    }

    private static class POrderedMapCollector<K, V> {
        private POrderedMap<K, V> map = POrderedMap.empty();

        void add(Map.Entry<K, V> entry) {
            map = map.plus(entry.getKey(), entry.getValue());
        }

        POrderedMapCollector<K, V> addAll(POrderedMapCollector<K, V> other) {
            map = map.plusAll(other.map);
            return this;
        }

        POrderedMap<K, V> toOrderedMap() {
            return map;
        }
    }

    public static <K, V> Collector<Map.Entry<K, V>, OrderedPMapCollector<K, V>, OrderedPMap<K, V>> toOrderedPMap() {
        return Collector.of(
            OrderedPMapCollector::new,
            OrderedPMapCollector::add,
            OrderedPMapCollector::addAll,
            OrderedPMapCollector::toOrderedMap
        );
    }

    private static class OrderedPMapCollector<K, V> {
        private OrderedPMap<K, V> map = OrderedPMap.empty();

        void add(Map.Entry<K, V> entry) {
            map = map.plus(entry.getKey(), entry.getValue());
        }

        OrderedPMapCollector<K, V> addAll(OrderedPMapCollector<K, V> other) {
            map = map.plusAll(other.map);
            return this;
        }

        OrderedPMap<K, V> toOrderedMap() {
            return map;
        }
    }

    public static <T> Collector<T, OrderedPSetCollector<T>, OrderedPSet<T>> toOrderedSet() {
        return Collector.of(
            OrderedPSetCollector::new,
            OrderedPSetCollector::add,
            OrderedPSetCollector::addAll,
            OrderedPSetCollector::toSet
        );
    }

    private static class OrderedPSetCollector<T> {
        private OrderedPSet<T> set = OrderedPSet.empty();

        void add(T value) {
            set = set.plus(value);
        }

        OrderedPSetCollector<T> addAll(OrderedPSetCollector<T> other) {
            set = set.plusAll(other.set);
            return this;
        }

        OrderedPSet<T> toSet() {
            return set;
        }
    }

    public static <T> Collector<T, PVectorCollector<T>, PVector<T>> toVector() {
        return Collector.of(
            PVectorCollector::new,
            PVectorCollector::add,
            PVectorCollector::addAll,
            PVectorCollector::toVector
        );
    }

    private static class PVectorCollector<T> {
        private PVector<T> vector = Empty.vector();

        void add(T value) {
            vector = vector.plus(value);
        }

        PVectorCollector<T> addAll(PVectorCollector<T> other) {
            vector = vector.plusAll(other.vector);
            return this;
        }

        PVector<T> toVector() {
            return vector;
        }
    }
}