package com.wizzardo.tools.collections.lazy;

import java.util.List;
import java.util.Map;

/**
 * Created by wizzardo on 08.11.15.
 */
public class LazyGrouping<K, T, A, B extends LazyGroup<K, T, T>> extends AbstractLazy<A, B> {

    protected final Map<K, LazyGroup<K, T, T>> groups;

    LazyGrouping(Map<K, LazyGroup<K, T, T>> groups) {
        this.groups = groups;
    }

    public <V> Lazy<V, V> flatMap(Mapper<? super B, V> mapper) {
        LazyContinue<V> continueCommand = new LazyContinue<V>(this);
        then(new GroupCommand<K, V, B>(mapper, continueCommand));
        return continueCommand;
    }

    public Map<K, List<T>> toMap() {
        return toMap(new LazyGroupToListMapper<K, T>());
    }

    public <V> Map<K, V> toMap(Mapper<? super B, V> mapper) {
        ToMapCommand<K, V, B> toMap;
        then(toMap = new ToMapCommand<K, V, B>((Map<K, V>) groups, mapper));
        toMap.start();
        return toMap.get();
    }

    @Override
    public LazyGrouping<K, T, B, B> filter(final Filter<? super B> filter) {
        return this.then(new LazyGrouping<K, T, B, B>(groups) {
            @Override
            protected void process(B b) {
                if (filter.allow(b))
                    child.process(b);
            }
        });
    }

    @Override
    public LazyGrouping<K, T, B, B> each(final Consumer<? super B> consumer) {
        return this.then(new LazyGrouping<K, T, B, B>(groups) {
            @Override
            protected void process(B b) {
                consumer.consume(b);
                child.process(b);
            }
        });
    }

    public Lazy<B, T> merge() {
        return then(new LazyMerge<B, T>());
    }

    private static class GroupCommand<K, V, B extends LazyGroup<K, ?, ?>> extends Command<B, B> {
        private final Mapper<? super B, V> mapper;
        private final Command<V, V> continueCommand;

        public GroupCommand(Mapper<? super B, V> mapper, Command<V, V> continueCommand) {
            this.mapper = mapper;
            this.continueCommand = continueCommand;
        }

        @Override
        protected void process(B b) {
            mapper.map(b);
            getLast(b).then(new LazyOnEnd<V>(continueCommand));
        }

        @Override
        protected void onEnd() {
        }
    }

    private static class ToMapCommand<K, V, B extends LazyGroup<K, ?, ?>> extends FinishCommand<B, Map<K, V>> {
        private Map<K, V> groups;
        private final Mapper<? super B, V> mapper;

        public ToMapCommand(Map<K, V> groups, Mapper<? super B, V> mapper) {
            this.groups = groups;
            this.mapper = mapper;
        }

        @Override
        protected void process(final B b) {
            mapper.map(b);
            getLast(b).then(new LazyOnEnd<V>(new Command<V, V>() {
                @Override
                protected void process(V v) {
                    groups.put(b.getKey(), v);
                }
            }));
        }

        @Override
        protected Map<K, V> get() {
            return groups;
        }
    }

    private static class LazyContinue<T> extends Lazy<T, T> {
        private Command<?, ?> command;

        LazyContinue(Command<?, ?> command) {
            this.command = command;
        }

        @Override
        protected void start() {
            command.start();
        }

        @Override
        protected void stop() {
            command.stop();
        }

        @Override
        protected void process(T t) {
            child.process(t);
        }
    }

    private static class LazyOnEnd<T> extends Command<T, T> {
        private Command<T, ?> command;

        private LazyOnEnd(Command<T, ?> command) {
            this.command = command;
        }

        @Override
        protected void onEnd() {
            command.process(parent.get());
        }
    }
}

@@ -39,16 +39,48 @@
 .lazy;%0A%0A
+import java.util.LinkedHashMap;%0A
 import j
@@ -1114,38 +1114,74 @@
 ing%3CK, T, B, B%3E(
-groups
+new LinkedHashMap%3CK, LazyGroup%3CK, T, T%3E%3E()
 ) %7B%0A            
@@ -1268,16 +1268,65 @@
 llow(b))
+ %7B%0A                    groups.put(b.getKey(), b);
 %0A       
@@ -1348,32 +1348,50 @@
 ild.process(b);%0A
+                %7D%0A
             %7D%0A  

fix toMap after filtering