/*
 * Copyright 2014 Real Logic Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package uk.co.real_logic.aeron.driver;

import uk.co.real_logic.aeron.common.collections.Int2ObjectHashMap;
import uk.co.real_logic.aeron.common.concurrent.AtomicBuffer;
import uk.co.real_logic.aeron.common.protocol.DataHeaderFlyweight;
import uk.co.real_logic.aeron.common.protocol.HeaderFlyweight;
import uk.co.real_logic.aeron.driver.exceptions.UnknownSubscriptionException;

import java.net.InetSocketAddress;

/**
 * Handling of dispatching data frames to {@link DriverConnection}s streams.
 * <p>
 * All methods should be called via {@link Receiver} thread
 */
public class DataFrameDispatcher
{
    private static final String INIT_IN_PROGRESS = "Connection initialisation in progress";

    private final Int2ObjectHashMap<String> initialisationInProgressMap = new Int2ObjectHashMap<>();
    private final Int2ObjectHashMap<Int2ObjectHashMap<DriverConnection>> connectionsByStreamIdMap = new Int2ObjectHashMap<>();
    private final DriverConductorProxy conductorProxy;
    private final ReceiveChannelEndpoint channelEndpoint;

    public DataFrameDispatcher(final DriverConductorProxy conductorProxy, final ReceiveChannelEndpoint channelEndpoint)
    {
        this.conductorProxy = conductorProxy;
        this.channelEndpoint = channelEndpoint;
    }

    public void addSubscription(final int streamId)
    {
        Int2ObjectHashMap<DriverConnection> connectionBySessionIdMap = connectionsByStreamIdMap.get(streamId);

        if (null == connectionBySessionIdMap)
        {
            connectionBySessionIdMap = new Int2ObjectHashMap<>();
            connectionsByStreamIdMap.put(streamId, connectionBySessionIdMap);
        }
    }

    public void removeSubscription(final int streamId)
    {
        final Int2ObjectHashMap<DriverConnection> connectionBySessionIdMap = connectionsByStreamIdMap.remove(streamId);

        if (null == connectionBySessionIdMap)
        {
            throw new UnknownSubscriptionException("No connectionBySessionIdMap registered on " + streamId);
        }

        for (final DriverConnection connection : connectionBySessionIdMap.values())
        {
            connection.disableStatusMessages();
        }
    }

    public void addConnection(final DriverConnection connection)
    {
        final Int2ObjectHashMap<DriverConnection> connectionBySessionIdMap = connectionsByStreamIdMap.get(connection.streamId());

        if (null == connectionBySessionIdMap)
        {
            throw new IllegalStateException("No connectionBySessionIdMap registered on " + connection.streamId());
        }

        connectionBySessionIdMap.put(connection.sessionId(), connection);
        initialisationInProgressMap.remove(connection.sessionId());

        connection.enableStatusMessages();
    }

    public void removeConnection(final DriverConnection connection)
    {
        final Int2ObjectHashMap<DriverConnection> connectionBySessionIdMap = connectionsByStreamIdMap.get(connection.streamId());

        if (null != connectionBySessionIdMap)
        {
            connection.disableStatusMessages();
            connectionBySessionIdMap.remove(connection.sessionId());
            initialisationInProgressMap.remove(connection.sessionId());
        }
    }

    public void onDataFrame(
        final DataHeaderFlyweight dataHeader, final AtomicBuffer buffer, final int length, final InetSocketAddress srcAddress)
    {
        final int streamId = dataHeader.streamId();
        final Int2ObjectHashMap<DriverConnection> connectionBySessionIdMap = connectionsByStreamIdMap.get(streamId);

        if (null != connectionBySessionIdMap)
        {
            final int sessionId = dataHeader.sessionId();
            final int termId = dataHeader.termId();
            final DriverConnection connection = connectionBySessionIdMap.get(sessionId);

            if (null != connection)
            {
                if (length > DataHeaderFlyweight.HEADER_LENGTH || dataHeader.headerType() == HeaderFlyweight.HDR_TYPE_PAD)
                {
                    connection.insertIntoTerm(termId, dataHeader.termOffset(), buffer, length);
                }
            }
            else if (null == initialisationInProgressMap.get(sessionId))
            {
                createConnection(srcAddress, streamId, sessionId, termId);
            }
        }
    }

    private void createConnection(final InetSocketAddress srcAddress, final int streamId, final int sessionId, final int termId)
    {
        final UdpTransport transport = channelEndpoint.udpTransport();
        final InetSocketAddress controlAddress =
            transport.isMulticast() ? transport.udpChannel().remoteControl() : srcAddress;

        initialisationInProgressMap.put(sessionId, INIT_IN_PROGRESS);
        conductorProxy.createConnection(sessionId, streamId, termId, controlAddress, channelEndpoint);
    }
}

@@ -628,16 +628,82 @@
 river;%0A%0A
+import uk.co.real_logic.aeron.common.collections.BiInt2ObjectMap;%0A
 import u
@@ -1359,32 +1359,34 @@
   private final 
+Bi
 Int2ObjectHashMa
@@ -1375,28 +1375,24 @@
 BiInt2Object
-Hash
 Map%3CString%3E 
@@ -1417,32 +1417,34 @@
 ogressMap = new 
+Bi
 Int2ObjectHashMa
@@ -1429,36 +1429,32 @@
 new BiInt2Object
-Hash
 Map%3C%3E();%0A    pri
@@ -2881,32 +2881,139 @@
 nnection)%0A    %7B%0A
+        final int sessionId = connection.sessionId();%0A        final int streamId = connection.streamId();%0A%0A
         final In
@@ -3102,35 +3102,24 @@
 Map.get(
-connection.
 streamId
 ());%0A%0A  
@@ -3102,34 +3102,32 @@
 Map.get(streamId
-()
 );%0A%0A        if (
@@ -3261,35 +3261,24 @@
  on %22 + 
-connection.
 streamId
 ());%0A   
@@ -3269,18 +3269,16 @@
 streamId
-()
 );%0A     
@@ -3320,27 +3320,16 @@
 Map.put(
-connection.
 sessionI
@@ -3329,18 +3329,16 @@
 essionId
-()
 , connec
@@ -3387,38 +3387,35 @@
 .remove(
-connection.sessionId()
+sessionId, streamId
 );%0A%0A    
@@ -3526,32 +3526,139 @@
 nnection)%0A    %7B%0A
+        final int sessionId = connection.sessionId();%0A        final int streamId = connection.streamId();%0A%0A
         final In
@@ -3747,35 +3747,24 @@
 Map.get(
-connection.
 streamId
 ());%0A%0A  
@@ -3755,18 +3755,16 @@
 streamId
-()
 );%0A%0A    
@@ -3903,35 +3903,24 @@
 dMap.remove(
-connection.
 sessionId())
@@ -3908,34 +3908,32 @@
 remove(sessionId
-()
 );%0A            i
@@ -3966,38 +3966,35 @@
 .remove(
-connection.sessionId()
+sessionId, streamId
 );%0A     
@@ -4982,16 +4982,26 @@
 essionId
+, streamId
 ))%0A     
@@ -5512,16 +5512,26 @@
 ssionId,
+ streamId,
  INIT_IN

fix map for initialisation in progress to account for session Id and stream Id.