/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.cassandra.transport;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.Collections;
import java.util.function.Function;
import java.util.function.Predicate;

import com.google.common.collect.ImmutableMap;

import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.Mockito;

import org.apache.cassandra.auth.AuthenticatedUser;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.service.ClientState;
import org.apache.cassandra.transport.Server.ConnectionTracker;

import io.netty.channel.Channel;
import io.netty.channel.ChannelId;
import io.netty.channel.embedded.EmbeddedChannel;

import static org.junit.Assert.assertEquals;

public class ConnectionTrackerTest
{

    private static final Predicate<ServerConnection> PREDICATE_TRUE = (__) -> true;

    @BeforeClass
    public static void beforeClass()
    {
        // perform daemon initialization to intialize core components (like authenticator) that ConnectionTracker
        // depends on.
        DatabaseDescriptor.daemonInitialization();
    }

    private void registerConnections(ConnectionTracker tracker, int count)
    {
        registerConnections(tracker, count, (addr) -> null);
    }

    private void registerConnections(ConnectionTracker tracker, int count, Function<InetSocketAddress, AuthenticatedUser> authenticatedUserGenerator)
    {
        for (int i = 0; i < count; i++)
        {
            InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), 9040 + i);

            Channel c = new EmbeddedChannelWithSocketAddress(socketAddress);
            AuthenticatedUser user = authenticatedUserGenerator.apply(socketAddress);
            ServerConnection connection = createAuthenticatedConnection(c, user);
            c.attr(Connection.attributeKey).set(connection);

            tracker.addConnection(c, connection);
        }
    }

    @Test
    public void testCountConnectedClients()
    {
        final String ODD_USER = "odd";
        final String EVEN_USER = "even";
        ConnectionTracker connectionTracker = new ConnectionTracker(() -> true);
        registerConnections(connectionTracker, 6, (addr) -> {
            // if port divisible by 10, return 0, this will be treated as an anonymous user.
            if (addr.getPort() % 10 == 0)
            {
                return null;
            }
            return new AuthenticatedUser(addr.getPort() % 2 == 0 ? EVEN_USER : ODD_USER);
        });

        // Validate that countConnectedClients by predicate returns the proper count.

        // Should be 6 total connections when all matching this predicate
        assertEquals(6, connectionTracker.countConnectedClients(PREDICATE_TRUE));

        // Three connections should have a port > 9042.
        assertEquals(3, connectionTracker.countConnectedClients((conn) -> {
            int port = conn.getClientState().getRemoteAddress().getPort();
            return port > 9042;
        }));

        // Zero connections using ssl.
        assertEquals(0, connectionTracker.countConnectedClients(ServerConnection::isSSL));

        // Verify countConnectedClientsByUser appropriately counts by user.

        // Expect 2 users authenticated as 'even' {2, 4}, 3 as 'odd' {1, 3, 5}, one anonymous {0}.
        assertEquals(ImmutableMap.of(EVEN_USER, 2, ODD_USER, 3, "anonymous", 1), connectionTracker.countConnectedClientsByUser());
    }

    @Test
    public void testShouldNotCountChannelsMissingConnectionAttribute()
    {
        ConnectionTracker connectionTracker = new ConnectionTracker(() -> true);
        registerConnections(connectionTracker, 10);

        // Given a ServerConnection registered that lacks a 'Connection.attributeKey' attribute, it should
        // be skipped over when counting client connections.
        InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), 9042);
        Channel c = new EmbeddedChannelWithSocketAddress(socketAddress);
        Connection connection = new ServerConnection(c, ProtocolVersion.V5, NO_OP_TRACKER);
        // intentionally omit setting Connection attribute.
        connectionTracker.addConnection(c, connection);

        // Expect the 10 existing connections, omitting the one added without a connection attribute.
        assertEquals(10, connectionTracker.countConnectedClients(PREDICATE_TRUE));
        assertEquals(Collections.singletonMap("anonymous", 10), connectionTracker.countConnectedClientsByUser());
    }

    @Test
    public void testShouldNotCountNonServerConnection()
    {
        ConnectionTracker connectionTracker = new ConnectionTracker(() -> true);
        registerConnections(connectionTracker, 10);

        // Given a Connection registered that is not a 'ServerConnection' it should be skipped over when counting client
        // connections.
        InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), 9042);
        Channel c = new EmbeddedChannelWithSocketAddress(socketAddress);
        Connection connection = new Connection(c, ProtocolVersion.V5, NO_OP_TRACKER);
        c.attr(Connection.attributeKey).set(connection);
        connectionTracker.addConnection(c, connection);

        // Expect the 10 existing connections, omitting the one that isn't a ServerConnection.
        assertEquals(10, connectionTracker.countConnectedClients(PREDICATE_TRUE));
        assertEquals(Collections.singletonMap("anonymous", 10), connectionTracker.countConnectedClientsByUser());
    }

    @Test(expected = IllegalArgumentException.class)
    public void testShouldThrowThrowable()
    {
        ConnectionTracker connectionTracker = new ConnectionTracker(() -> true);
        registerConnections(connectionTracker, 10);

        // Mock a ServerConnection such that when you attempt to access its connection attribute, an unhandled throwable
        // is raised.
        InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), 9042);
        Channel c = new EmbeddedChannelWithSocketAddress(socketAddress);
        Channel spyChannel = Mockito.spy(c);
        Connection connection = new ServerConnection(c, ProtocolVersion.V5, NO_OP_TRACKER);
        c.attr(Connection.attributeKey).set(connection);
        Mockito.when(spyChannel.attr(Connection.attributeKey)).thenThrow(new IllegalArgumentException("Unexpected behavior"));
        connectionTracker.addConnection(spyChannel, connection);

        // When calling countConnectedClients, we expect that expection to be propagated as we can't get a correct calculation.
        connectionTracker.countConnectedClients(PREDICATE_TRUE);
    }

    /**
     * @return a {@link ServerConnection} whose {@link ClientState} has the given {@link AuthenticatedUser}
     */
    private ServerConnection createAuthenticatedConnection(Channel channel, AuthenticatedUser user)
    {
        ServerConnection connection = new ServerConnection(channel, ProtocolVersion.V5, new Connection.Tracker()
        {
            @Override
            public void addConnection(Channel ch, Connection connection)
            {

            }

            @Override
            public boolean isRunning()
            {
                return false;
            }
        });

        if (user == null)
        {
            return connection;
        }

        ClientState state = connection.getClientState();

        ClientState spyState = Mockito.spy(state);
        Mockito.when(spyState.getUser()).thenReturn(user);

        ServerConnection spyConnection = Mockito.spy(connection);
        Mockito.when(spyConnection.getClientState()).thenReturn(spyState);

        return spyConnection;
    }

    /**
     * An embedded channel that uses the given {@link SocketAddress}'s string as its channel id, and {@link #remoteAddress()}
     * returns the provided address.
     *
     * This is needed because {@link ServerConnection} expects a {@link Channel} with a {@link SocketAddress}
     */
    private static class EmbeddedChannelWithSocketAddress extends EmbeddedChannel
    {
        private final SocketAddress socketAddress;

        EmbeddedChannelWithSocketAddress(SocketAddress socketAddress)
        {
            super(createChannelId(socketAddress));
            this.socketAddress = socketAddress;
        }

        private static ChannelId createChannelId(SocketAddress socketAddress)
        {
            return new ChannelId()
            {
                @Override
                public String asShortText()
                {
                    return socketAddress.toString();
                }

                @Override
                public String asLongText()
                {
                    return asShortText();
                }

                @Override
                public int compareTo(ChannelId o)
                {
                    return this.asShortText().compareTo(this.asShortText());
                }
            };
        }

        @Override
        public SocketAddress remoteAddress()
        {
            return this.socketAddress;
        }
    }

    private final Connection.Tracker NO_OP_TRACKER =  new Connection.Tracker()
    {
        @Override
        public void addConnection(Channel ch, Connection connection)
        {
        }

        @Override
        public boolean isRunning()
        {
            return true;
        }
    };
}
