package com.digitalpetri.enip.cip.services;

import java.util.List;
import java.util.function.BiConsumer;

import com.digitalpetri.enip.cip.CipResponseException;
import com.digitalpetri.enip.cip.epath.EPath.PaddedEPath;
import com.digitalpetri.enip.cip.epath.LogicalSegment.ClassId;
import com.digitalpetri.enip.cip.epath.LogicalSegment.InstanceId;
import com.digitalpetri.enip.cip.structs.MessageRouterRequest;
import com.digitalpetri.enip.cip.structs.MessageRouterResponse;
import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;

import static com.google.common.collect.Lists.newArrayList;
import static java.util.Collections.synchronizedList;

public class MultipleServicePacketService implements CipService<Void> {

    public static final int SERVICE_CODE = 0x0A;

    private static final PaddedEPath MESSAGE_ROUTER_PATH = new PaddedEPath(
            new ClassId(0x02),
            new InstanceId(0x01));

    private final List<CipService<?>> services;
    private final List<BiConsumer<?, Throwable>> consumers;

    public MultipleServicePacketService(List<CipService<?>> services, List<BiConsumer<?, Throwable>> consumers) {
        assert (services.size() == consumers.size());

        this.services = synchronizedList(newArrayList(services));
        this.consumers = synchronizedList(newArrayList(consumers));
    }

    @Override
    public void encodeRequest(ByteBuf buffer) {
        MessageRouterRequest request = new MessageRouterRequest(
                SERVICE_CODE,
                MESSAGE_ROUTER_PATH,
                this::encode);

        MessageRouterRequest.encode(request, buffer);
    }

    @Override
    public Void decodeResponse(ByteBuf buffer) throws CipResponseException, PartialResponseException {
        MessageRouterResponse response = MessageRouterResponse.decode(buffer);

        if (response.getGeneralStatus() == 0x00 || response.getGeneralStatus() == 0x1E) {
            try {
                List<Object[]> partials = newArrayList();

                ByteBuf[] serviceData = decode(response.getData());

                for (int i = 0; i < serviceData.length; i++) {
                    CipService<?> service = services.get(i);

                    @SuppressWarnings("unchecked")
                    BiConsumer<Object, Throwable> consumer =
                            (BiConsumer<Object, Throwable>) consumers.get(i);

                    try {
                        consumer.accept(service.decodeResponse(serviceData[i]), null);
                    } catch (PartialResponseException prx) {
                        partials.add(new Object[]{service, consumer});
                    } catch (Throwable t) {
                        consumer.accept(null, t);
                    } finally {
                        ReferenceCountUtil.release(serviceData[i]);
                    }
                }

                if (!partials.isEmpty()) {
                    services.clear();
                    consumers.clear();

                    for (Object[] oa : partials) {
                        CipService<?> service = (CipService<?>) oa[0];

                        @SuppressWarnings("unchecked")
                        BiConsumer<Object, Throwable> consumer =
                                (BiConsumer<Object, Throwable>) oa[1];

                        services.add(service);
                        consumers.add(consumer);
                    }

                    throw PartialResponseException.INSTANCE;
                }

                return null;
            } finally {
                ReferenceCountUtil.release(response.getData());
            }
        } else {
            throw new CipResponseException(response.getGeneralStatus(), response.getAdditionalStatus());
        }
    }

    private void encode(ByteBuf buffer) {
        int serviceCount = services.size();

        buffer.writeShort(serviceCount);

        int[] offsets = new int[serviceCount];
        int offsetsStartIndex = buffer.writerIndex();
        buffer.writeZero(serviceCount * 2);

        for (int i = 0; i < serviceCount; i++) {
            offsets[i] = buffer.writerIndex() - offsetsStartIndex + 2;
            services.get(i).encodeRequest(buffer);
        }

        buffer.markWriterIndex();
        buffer.writerIndex(offsetsStartIndex);
        for (int offset : offsets) {
            buffer.writeShort(offset);
        }
        buffer.resetWriterIndex();
    }

    private ByteBuf[] decode(ByteBuf buffer) {
        int dataStartIndex = buffer.readerIndex();
        int serviceCount = buffer.readUnsignedShort();

        int[] offsets = new int[serviceCount];
        for (int i = 0; i < serviceCount; i++) {
            offsets[i] = buffer.readUnsignedShort();
        }

        ByteBuf[] serviceData = new ByteBuf[serviceCount];
        for (int i = 0; i < serviceCount; i++) {
            int offset = offsets[i];

            int length = (i + 1 < serviceCount) ?
                    offsets[i + 1] - offset :
                    buffer.readableBytes();

            serviceData[i] = buffer.slice(dataStartIndex + offsets[i], length).retain();

            buffer.skipBytes(length);
        }

        return serviceData;
    }

}

@@ -1033,16 +1033,139 @@
 umers;%0A%0A
+    private final List%3CCipService%3C?%3E%3E currentServices;%0A    private final List%3CBiConsumer%3C?, Throwable%3E%3E currentConsumers;%0A%0A
     publ
@@ -1266,24 +1266,24 @@
 onsumers) %7B%0A
-
         asse
@@ -1454,24 +1454,173 @@
 consumers));
+%0A%0A        this.currentServices = synchronizedList(newArrayList(services));%0A        this.currentConsumers = synchronizedList(newArrayList(consumers));
 %0A    %7D%0A%0A    
@@ -2442,17 +2442,24 @@
 rvice = 
-s
+currentS
 ervices.
@@ -2640,16 +2640,23 @@
 able%3E) c
+urrentC
 onsumers
@@ -3169,9 +3169,8 @@
 if (
-!
 part
@@ -3211,47 +3211,412 @@
     
-services.clear();%0A                    c
+// Reset to the original state.%0A                    currentServices.clear();%0A                    currentServices.addAll(services);%0A%0A                    currentConsumers.clear();%0A                    currentConsumers.addAll(consumers);%0A                %7D else %7B%0A                    // Keep sending only services that aren't done yet.%0A                    currentServices.clear();%0A                    currentC
 onsu
@@ -3961,33 +3961,40 @@
                 
-s
+currentS
 ervices.add(serv
@@ -4020,24 +4020,31 @@
            c
+urrentC
 onsumers.add
@@ -4497,17 +4497,24 @@
 Count = 
-s
+currentS
 ervices.
@@ -4838,25 +4838,32 @@
             
-s
+currentS
 ervices.get(

Make sure MultipleServicePacketService is re-usable