/*
 * 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.locator;

import java.io.IOException;
import java.util.Collections;
import java.util.Map;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;

import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.stubbing.Answer;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.cassandra.tcm.ClusterMetadata;
import org.apache.cassandra.tcm.ClusterMetadataService;
import org.apache.cassandra.tcm.StubClusterMetadataService;
import org.apache.cassandra.tcm.membership.Directory;
import org.apache.cassandra.utils.Pair;

import static org.apache.cassandra.locator.Ec2MultiRegionSnitch.PRIVATE_IP_QUERY;
import static org.apache.cassandra.locator.Ec2MultiRegionSnitch.PUBLIC_IP_QUERY;
import static org.apache.cassandra.locator.Ec2Snitch.EC2_NAMING_LEGACY;
import static org.junit.Assert.assertEquals;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

public class Ec2SnitchTest
{
    private final SnitchProperties legacySnitchProps = new SnitchProperties()
    {
        public String get(String propertyName, String defaultValue)
        {
            return propertyName.equals("ec2_naming_scheme") ? EC2_NAMING_LEGACY : super.get(propertyName, defaultValue);
        }
    };

    @BeforeClass
    public static void setup() throws Exception
    {
        DatabaseDescriptor.daemonInitialization();
    }

    @Before
    public void resetCMS()
    {
        ClusterMetadataService.unsetInstance();
        ClusterMetadataService.setInstance(StubClusterMetadataService.forTesting());
    }

    @Test
    public void testLegacyRac() throws Exception
    {
        Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
        doReturn("us-east-1d").when(connectorMock).apiCall(anyString());
        doReturn(legacySnitchProps).when(connectorMock).getProperties();

        Ec2LocationProvider locationProvider = new Ec2LocationProvider(connectorMock);
        testLegacyRacInternal(locationProvider);

        Ec2Snitch snitch = new Ec2Snitch(connectorMock);
        testLegacyRacInternal(snitch);
    }

    @Test
    public void testMultiregionLegacyRac() throws Exception
    {
        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
        doReturn(legacySnitchProps).when(spy).getProperties();
        doAnswer((Answer<String>) invocation -> {
            String query = invocation.getArgument(0);
            return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-1d";
        }).when(spy).apiCall(anyString());

        Ec2LocationProvider locationProvider = new Ec2LocationProvider(spy);
        testLegacyRacInternal(locationProvider);

        Ec2Snitch snitch = new Ec2MultiRegionSnitch(spy);
        testLegacyRacInternal(snitch);
    }

    @Test
    public void testLegacyNewRegions() throws Exception
    {
        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
        doReturn(legacySnitchProps).when(spy).getProperties();
        doReturn("us-east-2d").when(spy).apiCall(anyString());
        testLegacyNewRegionsInternal(new Ec2LocationProvider(spy));
        testLegacyNewRegionsInternal(new Ec2Snitch(spy));
    }

    @Test
    public void testLegacyMultiRegionNewRegions() throws Exception
    {
        Ec2MetadataServiceConnector spy = spy(Ec2MetadataServiceConnector.create(legacySnitchProps));
        doReturn(legacySnitchProps).when(spy).getProperties();
        doAnswer((Answer<String>) invocation -> {
            String query = invocation.getArgument(0);
            return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d";
        }).when(spy).apiCall(anyString());

        testLegacyNewRegionsInternal(new Ec2LocationProvider(spy));
        testLegacyNewRegionsInternal(new Ec2MultiRegionSnitch(spy));
    }

    @Test
    public void testFullNamingScheme() throws Exception
    {
        Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
        when(connectorMock.apiCall(anyString())).thenReturn("us-east-2d");
        when(connectorMock.getProperties()).thenReturn(new SnitchProperties());

        Ec2LocationProvider locationProvider = new Ec2LocationProvider(connectorMock);
        assertEquals("us-east-2", locationProvider.initialLocation().datacenter);
        assertEquals("us-east-2d", locationProvider.initialLocation().rack);

        Ec2Snitch snitch = new Ec2Snitch(connectorMock);
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east-2", snitch.getLocalDatacenter());
        assertEquals("us-east-2d", snitch.getLocalRack());

        Ec2MetadataServiceConnector multiRegionConnectorMock = mock(Ec2MetadataServiceConnector.class);
        when(multiRegionConnectorMock.getProperties()).thenReturn(new SnitchProperties());
        when(multiRegionConnectorMock.apiCall(anyString())).then((Answer<String>) invocation -> {
            String query = invocation.getArgument(0);
            return (PUBLIC_IP_QUERY.equals(query) || PRIVATE_IP_QUERY.equals(query)) ? "127.0.0.1" : "us-east-2d";
        });

        Ec2LocationProvider mrLocationProvider = new Ec2LocationProvider(connectorMock);
        assertEquals("us-east-2", mrLocationProvider.initialLocation().datacenter);
        assertEquals("us-east-2d", mrLocationProvider.initialLocation().rack);

        Ec2MultiRegionSnitch mrSnitch = new Ec2MultiRegionSnitch(multiRegionConnectorMock);
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east-2", mrSnitch.getLocalDatacenter());
        assertEquals("us-east-2d", mrSnitch.getLocalRack());
    }

    @Test
    public void validateDatacenter_RequiresLegacy_CorrectAmazonName() throws IOException
    {
        // legacy scheme / standard dc / ignore racks
        assertValid(true, "us-east-1");
    }

    @Test
    public void validateDatacenter_RequiresLegacy_LegacyName() throws IOException
    {
        // legacy scheme / legacy dc / ignore racks
        assertValid(true, "us-east");
    }

    @Test
    public void validate_RequiresLegacy_HappyPath() throws IOException
    {
        // legacy scheme / legacy dc / legacy rack
        assertValid(true, "us-east", "1a");
    }

    @Test
    public void validate_RequiresLegacy_HappyPathWithDCSuffix() throws IOException
    {
        // legacy scheme / custom suffix dc / legacy rack
        assertValid(true, "us-east_CUSTOM_SUFFIX", "1a");
    }

    @Test
    public void validateRack_RequiresAmazonName_CorrectAmazonName() throws IOException
    {
        // standard scheme / ignore dc / non legacy rack
        assertValid(false, "", "us-east-1a");
    }

    @Test(expected = ConfigurationException.class)
    public void validateRack_RequiresAmazonName_LegacyName() throws IOException
    {
        // standard scheme / ignore dc / legacy rack
        assertValid(false, "", "1a");
    }

    @Test
    public void validate_RequiresAmazonName_HappyPath() throws IOException
    {
        // standard scheme / standard dc / standard rack
        assertValid(false, "us-east-1", "us-east-1a");
    }

    @Test
    public void validate_RequiresAmazonName_HappyPathWithDCSuffix() throws IOException
    {
        // standard scheme / custom suffix dc / standard rack
        assertValid(false, "us-east-1_CUSTOM_SUFFIX", "us-east-1a");
    }

    /**
     * Validate upgrades in legacy mode for regions that didn't change name between the standard and legacy modes.
     */
    @Test
    public void validate_RequiresLegacy_DCValidStandardAndLegacy() throws IOException
    {
        // legacy scheme / standard dc / legacy racks
        assertValid(true, "us-west-2", "2a", "2b");
    }

    /**
     * Check that racks names are enough to detect a mismatch in naming conventions.
     */
    @Test(expected = ConfigurationException.class)
    public void validate_RequiresLegacy_RackInvalidForLegacy() throws Exception
    {
        // legacy scheme / standard dc / standard rack
        assertValid(true, "us-west-2", "us-west-2a");
    }

    private static void assertValid(boolean useLegacy, String datacenter, String...racks) throws IOException
    {
        SnitchProperties properties = new SnitchProperties(Pair.create(Ec2LocationProvider.SNITCH_PROP_NAMING_SCHEME,
                                                                       useLegacy ? Ec2LocationProvider.EC2_NAMING_LEGACY
                                                                                 : Ec2LocationProvider.EC2_NAMING_STANDARD));
        Ec2MetadataServiceConnector connectorMock = mock(Ec2MetadataServiceConnector.class);
        when(connectorMock.getProperties()).thenReturn(properties);
        Ec2LocationProvider provider = new Ec2LocationProvider(connectorMock);

        Multimap<String, InetAddressAndPort> dcRacks = HashMultimap.create();
        for (String rack : racks)
            dcRacks.put(rack, null); // endpoints are used
        Directory directory = new Directory()
        {
            @Override
            public Map<String, Multimap<String, InetAddressAndPort>> allDatacenterRacks()
            {
                return Collections.singletonMap(datacenter, dcRacks);
            }

            @Override
            public Multimap<String, InetAddressAndPort> datacenterRacks(String datacenter)
            {
                return dcRacks;
            }
        };
        ClusterMetadata metadata = new ClusterMetadata(Murmur3Partitioner.instance, directory);

        provider.validate(metadata);
    }

    private void testLegacyRacInternal(Ec2LocationProvider provider) throws Exception
    {
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east", provider.initialLocation().datacenter);
        assertEquals("1d", provider.initialLocation().rack);
    }

    private void testLegacyRacInternal(Ec2Snitch snitch) throws Exception
    {
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east", snitch.getLocalDatacenter());
        assertEquals("1d", snitch.getLocalRack());
    }

    private void testLegacyNewRegionsInternal(Ec2LocationProvider provider) throws Exception
    {
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east-2", provider.initialLocation().datacenter);
        assertEquals("2d", provider.initialLocation().rack);
    }

    private void testLegacyNewRegionsInternal(Ec2Snitch snitch) throws Exception
    {
        // for registering a new node, location is obtained from the cloud metadata service
        assertEquals("us-east-2", snitch.getLocalDatacenter());
        assertEquals("2d", snitch.getLocalRack());
    }
}
