/*
 * 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.io.IOException;
import java.util.Random;

import com.google.common.collect.ImmutableMap;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Session;
import org.apache.cassandra.cql3.CQLTester;
import org.apache.cassandra.cql3.QueryOptions;
import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.exceptions.TransportException;
import org.apache.cassandra.transport.messages.OptionsMessage;
import org.apache.cassandra.transport.messages.QueryMessage;
import org.apache.cassandra.transport.messages.StartupMessage;

import static com.datastax.driver.core.ProtocolVersion.NEWEST_BETA;
import static com.datastax.driver.core.ProtocolVersion.NEWEST_SUPPORTED;
import static com.datastax.driver.core.ProtocolVersion.V1;
import static com.datastax.driver.core.ProtocolVersion.V2;
import static com.datastax.driver.core.ProtocolVersion.V3;
import static com.datastax.driver.core.ProtocolVersion.V4;
import static com.datastax.driver.core.ProtocolVersion.V5;
import static com.datastax.driver.core.ProtocolVersion.V6;
import static org.apache.cassandra.transport.messages.StartupMessage.CQL_VERSION;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;


public class ProtocolNegotiationTest extends CQLTester
{
    // to avoid JMX naming clashes between cluster metrics
    private int clusterId = 0;

    @BeforeClass
    public static void setup()
    {
        requireNetwork();
    }

    @Before
    public void initNetwork()
    {
        reinitializeNetwork();
    }

    @Test
    public void serverSupportsV3AndV4AndV5ByDefault()
    {
        // client can explicitly request either V3, V4 or V5
        testConnection(V3, V3);
        testConnection(V4, V4);
        testConnection(V5, V5);

        // if not specified, V5 is the default
        testConnection(null, V5);
        testConnection(NEWEST_SUPPORTED, V5);
    }

    @Test
    public void supportV6ConnectionWithBetaOption()
    {
        testConnection(V6, V6);
        testConnection(NEWEST_BETA, V6);
    }

    @Test
    public void olderVersionsAreUnsupported()
    {
        testConnection(V1, V4);
        testConnection(V2, V4);
    }

    @Test
    public void preNegotiationResponsesHaveCorrectStreamId()
    {
        ProtocolVersion.SUPPORTED.forEach(this::testStreamIdsAcrossNegotiation);
    }

    @Test
    public void validateReceivedMessageVersionMatchesNegotiated()
    {
        ProtocolVersion.SUPPORTED.forEach(this::validateMessageVersion);
    }

    private void testStreamIdsAcrossNegotiation(ProtocolVersion version)
    {
        long seed = System.currentTimeMillis();
        Random random = new Random(seed);
        SimpleClient.Builder builder = SimpleClient.builder(nativeAddr.getHostAddress(), nativePort);
        if (version.isBeta())
            builder.useBeta();
        else
            builder.protocolVersion(version);

        try (SimpleClient client = builder.build())
        {
            client.establishConnection();
            // Before STARTUP the client hasn't yet negotiated a protocol version.
            // All OPTIONS messages are received by the intial connection handler.
            OptionsMessage options = new OptionsMessage();
            for (int i = 0; i < 100; i++)
            {
                int streamId = random.nextInt(254) + 1;
                options.setStreamId(streamId);
                Message.Response response = client.execute(options);
                assertEquals(String.format("StreamId mismatch; version: %s, seed: %s, iter: %s, expected: %s, actual: %s",
                                           version, seed, i, streamId, response.getStreamId()),
                             streamId, response.getStreamId());
            }

            int streamId = random.nextInt(254) + 1;
            // STARTUP messages are handled by the initial connection handler
            StartupMessage startup = new StartupMessage(ImmutableMap.of(CQL_VERSION, QueryProcessor.CQL_VERSION.toString()));
            startup.setStreamId(streamId);
            Message.Response response = client.execute(startup);
            assertEquals(String.format("StreamId mismatch after negotiation; version: %s, expected: %s, actual %s",
                                       version, streamId, response.getStreamId()),
                         streamId, response.getStreamId());

            // Following STARTUP, the version specific handlers are fully responsible for processing messages
            QueryMessage query = new QueryMessage("SELECT * FROM system.local", QueryOptions.DEFAULT);
            query.setStreamId(streamId);
            response = client.execute(query);
            assertEquals(String.format("StreamId mismatch after negotiation; version: %s, expected: %s, actual %s",
                                       version, streamId, response.getStreamId()),
                         streamId, response.getStreamId());
        }
        catch (IOException e)
        {
            e.printStackTrace();
            fail("Error establishing connection");
        }
    }

    private void testConnection(com.datastax.driver.core.ProtocolVersion requestedVersion,
                                com.datastax.driver.core.ProtocolVersion expectedVersion)
    {
        boolean expectError = requestedVersion != null && requestedVersion != expectedVersion;
        Cluster.Builder builder = Cluster.builder()
                                         .addContactPoints(nativeAddr)
                                         .withClusterName("Test Cluster" + clusterId++)
                                         .withPort(nativePort);

        if (requestedVersion != null)
        {
            if (requestedVersion.toInt() > org.apache.cassandra.transport.ProtocolVersion.CURRENT.asInt())
                builder = builder.allowBetaProtocolVersion();
            else
                builder = builder.withProtocolVersion(requestedVersion);
        }

        Cluster cluster = builder.build();
        try (Session session = cluster.connect())
        {
            if (expectError)
                fail("Expected a protocol exception");
            session.execute("SELECT * FROM system.local");
        }
        catch (Exception e)
        {
            if (!expectError)
            {
                e.printStackTrace();
                fail("Did not expect any exception");
            }
            e.printStackTrace();
            assertTrue(e.getMessage().contains(String.format("Host does not support protocol version %s", requestedVersion)));
        } finally {
            cluster.closeAsync();
        }
    }

    private void validateMessageVersion(ProtocolVersion version)
    {
        SimpleClient.Builder builder = SimpleClient.builder(nativeAddr.getHostAddress(), nativePort)
                                                   .protocolVersion(version);
        if (version.isBeta())
            builder.useBeta();

        Random r = new Random();
        ProtocolVersion wrongVersion = version;
        while (wrongVersion.isSmallerThan(ProtocolVersion.MIN_SUPPORTED_VERSION) || wrongVersion == version)
            wrongVersion = ProtocolVersion.values()[r.nextInt(ProtocolVersion.values().length - 1)];

        try (SimpleClient client = builder.build().connect(false))
        {
            // The connection has been negotiated to use $version. Force the next message to be
            // encoded with a different version and it should trigger a ProtocolException
            final ProtocolVersion v = wrongVersion;
            QueryMessage query = new QueryMessage("SELECT * FROM system.local", QueryOptions.DEFAULT)
            {
                @Override
                public Envelope encode(ProtocolVersion originalVersion)
                {
                    return super.encode(v);
                }
            };
            try
            {
                client.execute(query);
                fail("Expected a protocol exception");
            }
            catch (RuntimeException e)
            {
                assertTrue(e.getCause() instanceof TransportException);
                assertTrue(e.getCause().getMessage().startsWith("Invalid message version"));
            }
        }
        catch (IOException e)
        {
            e.printStackTrace();
            fail("Error establishing connection");
        }
    }
}