package com.jenjinstudios.io.connection;

import com.jenjinstudios.io.ExecutionContext;
import com.jenjinstudios.io.ExecutionContextFactory;
import com.jenjinstudios.io.MessageIOFactory;

import java.io.IOException;
import java.net.Socket;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * Used to generate multiple connections from sockets.
 *
 * @author Caleb Brinkman
 */
public class MultiConnectionBuilder
{
    private final Collection<Consumer<ExecutionContext>> contextualTasks = new LinkedList<>();
    private final Collection<Consumer<Connection>> shutdownCallbacks = new LinkedList<>();
    private ExecutionContextFactory executionContextFactory;
    private BiConsumer<Connection, Throwable> errorCallback;
    private MessageIOFactory messageIOFactory;

    /**
     * Build a connection using the given socket.
     *
     * @param socket The socket.
     *
     * @return The built connection.
     *
     * @throws IOException If there is an exception when building the connection.
     */
    public Connection build(Socket socket) throws IOException {
        return new SingleConnectionBuilder()
              .withMessageIOFactory(messageIOFactory)
              .withErrorCallback(errorCallback)
              .withContextualTasks(contextualTasks)
              .withExecutionContext(executionContextFactory.createInstance())
              .withSocket(socket)
              .withShutdownCallbacks(shutdownCallbacks)
              .build();
    }

    /**
     * Build a connection with the given error callback function.
     *
     * @param callback The Consumer (accepting a Connection and Throwable) that will be invoked when an error is
     * encountered.
     *
     * @return This ConnectionBuilder.
     */
    public MultiConnectionBuilder withErrorCallback(BiConsumer<Connection, Throwable> callback) {
        this.errorCallback = callback;
        return this;
    }

    /**
     * Use the given MessageIOFactory to create MessageReader and MessageWriter instances from Java Input and Output
     * streams.
     *
     * @param factory The MessageIOFactory.
     *
     * @return This ConnectionBuilder.
     */
    public MultiConnectionBuilder withMessageIOFactory(MessageIOFactory factory) {
        if (messageIOFactory == null) {
            this.messageIOFactory = factory;
        } else {
            throw new IllegalStateException("MessageIOFactory already set");
        }
        return this;
    }

    /**
     * Build a connection with the given ExecutionContext.
     *
     * @param context The context in which the Connection will execute messages.
     *
     * @return This ConnectionBuilder
     */
    public MultiConnectionBuilder withExecutionContextFactory(ExecutionContextFactory context) {
        if (executionContextFactory == null) {
            executionContextFactory = context;
        } else {
            throw new IllegalStateException("Execution context is already set");
        }
        return this;
    }

    /**
     * Build a connection that includes the given contextual task to be executed synchronously with message execution.
     *
     * @param task The task to be executed; a Consumer accepting an ExecutionContext.
     *
     * @return This ConnectionBuilder.
     */
    public MultiConnectionBuilder withContextualTask(Consumer<ExecutionContext> task) {
        contextualTasks.add(task);
        return this;
    }

    /**
     * Build a connection that includes the given contextual task to be executed synchronously with message execution.
     *
     * @param tasks The tasks to be executed; Consumers accepting an ExecutionContext.
     *
     * @return This ConnectionBuilder.
     */
    @SafeVarargs
    public final MultiConnectionBuilder withContextualTasks(Consumer<ExecutionContext>... tasks) {
        Collections.addAll(contextualTasks, tasks);
        return this;
    }

    /**
     * Build a connection that includes the given contextual task to be executed synchronously with message execution.
     *
     * @param tasks The tasks to be executed; Consumers accepting an ExecutionContext.
     *
     * @return This ConnectionBuilder.
     */
    public MultiConnectionBuilder withContextualTasks(Collection<Consumer<ExecutionContext>> tasks) {
        contextualTasks.addAll(tasks);
        return this;
    }

    /**
     * Build a connection that will execute the given callback on shutdown.
     *
     * @param callback The callback to be invoked.
     *
     * @return This ReusableConnectionBuilder.
     */
    public MultiConnectionBuilder withShutdownCallback(Consumer<Connection> callback) {
        shutdownCallbacks.add(callback);
        return this;
    }
}

@@ -67,123 +67,9 @@
 .io.
-ExecutionContext;%0Aimport com.jenjinstudios.io.ExecutionContextFactory;%0Aimport com.jenjinstudios.io.MessageIOFactory
+*
 ;%0A%0Ai
@@ -590,32 +590,142 @@
 LinkedList%3C%3E();%0A
+    private MessageReaderFactory messageReaderFactory;%0A    private MessageWriterFactory messageWriterFactory;%0A
     private Exec
@@ -833,55 +833,8 @@
 ack;
-%0A    private MessageIOFactory messageIOFactory;
 %0A%0A  
@@ -1197,34 +1197,38 @@
     .withMessage
-IO
+Reader
 Factory(messageI
@@ -1222,26 +1222,92 @@
 tory(message
-IO
+ReaderFactory)%0A              .withMessageWriterFactory(messageWriter
 Factory)%0A   
@@ -2375,24 +2375,25 @@
      if 
+(
 (message
 IOFactor
@@ -2380,26 +2380,30 @@
 if ((message
-IO
+Reader
 Factory == n
@@ -2411,146 +2411,435 @@
 ll) 
-%7B%0A            this.messageIOFactory = factory;%0A        %7D else %7B%0A            throw new IllegalStateException(%22MessageIOFactory already set%22
+&& (messageWriterFactory == null)) %7B%0A            this.messageReaderFactory = factory;%0A            this.messageWriterFactory = factory;%0A        %7D else %7B%0A            if (messageReaderFactory != null) %7B%0A                throw new IllegalStateException(%22MessageReaderFactory already set: %22 + messageReaderFactory);%0A            %7D%0A            throw new IllegalStateException(%22MessageWriterFactory already set: %22 + messageWriterFactory
 );%0A 

Refactor to use Message(Reader|Writer)Factory