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

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;

import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.config.StartupChecksConfiguration;
import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.api.IIsolatedExecutor;
import org.apache.cassandra.distributed.shared.WithProperties;
import org.apache.cassandra.exceptions.StartupException;
import org.apache.cassandra.io.util.File;
import org.apache.cassandra.service.DataResurrectionCheck.Heartbeat;
import org.apache.cassandra.service.StartupCheck;
import org.apache.cassandra.utils.Clock.Global;

import static java.lang.String.format;
import static java.util.concurrent.TimeUnit.MINUTES;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.apache.cassandra.config.CassandraRelevantProperties.CHECK_DATA_RESURRECTION_HEARTBEAT_PERIOD;
import static org.apache.cassandra.config.StartupChecksConfiguration.ENABLED_PROPERTY;
import static org.apache.cassandra.distributed.Cluster.build;
import static org.apache.cassandra.distributed.api.Feature.NATIVE_PROTOCOL;
import static org.apache.cassandra.io.util.FileUtils.createTempFile;
import static org.apache.cassandra.service.DataResurrectionCheck.EXCLUDED_KEYSPACES_CONFIG_PROPERTY;
import static org.apache.cassandra.service.DataResurrectionCheck.EXCLUDED_TABLES_CONFIG_PROPERTY;
import static org.apache.cassandra.service.DataResurrectionCheck.HEARTBEAT_FILE_CONFIG_PROPERTY;
import static org.apache.cassandra.service.DataResurrectionCheck.MINIMUM_THRESHOLD_CONFIG_PROPERTY;
import static org.awaitility.Awaitility.await;
import static org.hamcrest.Matchers.containsString;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;

public class DataResurrectionCheckTest extends TestBaseImpl
{
    public File heartbeatFile;

    @Before
    public void setupTest()
    {
        heartbeatFile = createTempFile("cassandra-heartbeat-" + UUID.randomUUID(), "");
    }

    @After
    public void teardownTest()
    {
        heartbeatFile.delete();
    }

    @Test
    public void testDataResurrectionCheck() throws Exception
    {
        // set it to 1 hour so check will be not updated after it is written, for test purposes
        try (WithProperties properties = new WithProperties().set(CHECK_DATA_RESURRECTION_HEARTBEAT_PERIOD, 3600000))
        {
            // start the node with the check enabled, it will just pass fine as there are not any user tables yet
            // and system tables are young enough
            try (Cluster cluster = build().withNodes(1)
                                          .withConfig(config -> config.with(NATIVE_PROTOCOL)
                                                                      .set("startup_checks",
                                                                           getStartupChecksConfig(ENABLED_PROPERTY, "true",
                                                                                                  HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath())))
                                          .start())
            {
                IInvokableInstance instance = cluster.get(1);

                checkHeartbeat();

                for (String ks : new String[]{ "ks1", "ks2", "ks3" })
                {
                    cluster.schemaChange("CREATE KEYSPACE " + ks + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};");
                    cluster.schemaChange(format("CREATE TABLE %s.tb1 (pk text PRIMARY KEY) WITH gc_grace_seconds = 10", ks));
                    cluster.schemaChange(format("CREATE TABLE %s.tb2 (pk text PRIMARY KEY)", ks));
                }

                AtomicReference<Throwable> throwable = new AtomicReference<>();
                // periodically execute check on a running instance and wait until the exception is thrown on all keyspaces
                // wait for all violations by Awaitility as due to nature how tables were created,
                // they will not expire on their gc_grace_period exactly at the same time
                await().timeout(1, MINUTES)
                       .pollInterval(5, SECONDS)
                       .until(() -> {
                           Throwable t = executeChecksOnInstance(instance,
                                                                 HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath());
                           if (t == null)
                               return false;
                           String message = t.getMessage();
                           if (!message.contains("ks1") || !message.contains("ks2") || !message.contains("ks3"))
                           {
                               return false;
                           }
                           throwable.set(t);
                           return true;
                       });

                assertThat(throwable.get().getMessage(), containsString("Invalid tables"));
                // returned tables in output are not in any particular order
                // it is how they are returned from system tables
                assertThat(throwable.get().getMessage(), containsString("ks1.tb1"));
                assertThat(throwable.get().getMessage(), containsString("ks2.tb1"));
                assertThat(throwable.get().getMessage(), containsString("ks3.tb1"));

                // exclude failing keyspaces which already expired on their gc_grace_seconds, so we will pass the check
                assertNull(executeChecksOnInstance(instance,
                                                   EXCLUDED_KEYSPACES_CONFIG_PROPERTY, "ks1,ks2,ks3",
                                                   HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath()));

                // exclude failing tables which already expired on their gc_grace_seconds, so we will pass the check
                assertNull(executeChecksOnInstance(instance,
                                                   EXCLUDED_TABLES_CONFIG_PROPERTY, "ks1.tb1,ks2.tb1,ks3.tb1",
                                                   HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath()));

                // exclude failing tables, but not all of them,
                // so check detects only one table violates the check
                Throwable t = executeChecksOnInstance(instance,
                                                      EXCLUDED_TABLES_CONFIG_PROPERTY, "ks1.tb1,ks2.tb1",
                                                      HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath());

                assertNotNull(t);
                assertThat(t.getMessage(), containsString("Invalid tables: ks3.tb1"));

                // shadow table exclusion with keyspace exclusion, we have not excluded ks3.tb1, but we excluded whole ks3
                assertNull(executeChecksOnInstance(instance,
                                                   EXCLUDED_TABLES_CONFIG_PROPERTY, "ks1.tb1,ks2.tb1",
                                                   EXCLUDED_KEYSPACES_CONFIG_PROPERTY, "ks3",
                                                   HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath()));
            }
        }
    }

    @Test
    public void testMinimumThreshold() throws Exception
    {
        // set it to 1 second so check will be updated very frequently
        // but when gc_grace_seconds is very low, that will still not fail the check,
        // because the minimum threshold is way bigger
        try (WithProperties properties = new WithProperties().set(CHECK_DATA_RESURRECTION_HEARTBEAT_PERIOD, 1000))
        {
            try (Cluster cluster = build().withNodes(1)
                                          .withConfig(config -> config.with(NATIVE_PROTOCOL)
                                                                      .set("startup_checks",
                                                                           getStartupChecksConfig(ENABLED_PROPERTY, "true",
                                                                                                  // set it to 1h to not fail even gc_grace_seconds is 0
                                                                                                  // and heartbeating period is 1 second
                                                                                                  MINIMUM_THRESHOLD_CONFIG_PROPERTY, "1h",
                                                                                                  HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath())))
                                          .start())
            {
                IInvokableInstance instance = cluster.get(1);

                checkHeartbeat();

                cluster.schemaChange("CREATE KEYSPACE ks1 WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1};");
                cluster.schemaChange("CREATE TABLE ks1.tb1 (pk text PRIMARY KEY) WITH gc_grace_seconds = 0");

                // just check 5 times with some poll that we will not get any
                // violations, even gc_grace_seconds is 0.
                // that is thanks to setting minimum_threshold to higher value, e.g. 1h

                for (int i = 0; i < 5; i++)
                {
                    // we do not get any exceptions / violations during whole time
                    // even heartbeating is going on and gc_grace_seconds is 0
                    Assert.assertNull(executeChecksOnInstance(instance,
                                                              MINIMUM_THRESHOLD_CONFIG_PROPERTY, "1h",
                                                              HEARTBEAT_FILE_CONFIG_PROPERTY, heartbeatFile.absolutePath()));
                    Thread.sleep(5000); // sleep 5s
                }
            }
        }
    }

    private Throwable executeChecksOnInstance(IInvokableInstance instance, final String... config)
    {
        assert config.length % 2 == 0;
        return instance.callsOnInstance((IIsolatedExecutor.SerializableCallable<Throwable>) () ->
        {
            try
            {
                StartupChecksConfiguration checksConfiguration = DatabaseDescriptor.getStartupChecksConfiguration();
                checksConfiguration.enable("check_data_resurrection");
                Map<String, Object> currentConfig = checksConfiguration.getConfig("check_data_resurrection");
                if (currentConfig != null)
                {
                    boolean wasEnabled = checksConfiguration.isEnabled("check_data_resurrection");
                    currentConfig.clear();
                    currentConfig.put(ENABLED_PROPERTY, wasEnabled);
                }

                for (int i = 0; i < config.length - 1; i = i + 2)
                    checksConfiguration.set("check_data_resurrection", config[i], config[i + 1]);

                StartupCheck check = checksConfiguration.getCheck("check_data_resurrection");
                check.execute(checksConfiguration);
                return null;
            }
            catch (StartupException e)
            {
                return e;
            }
        }).call();
    }

    private Map<String, Map<String, Object>> getStartupChecksConfig(String... configs)
    {
        return new HashMap<>()
        {{
            put("check_data_resurrection",
                new HashMap<>()
                {{
                    for (int i = 0; i < configs.length - 1; i = i + 2)
                        put(configs[i], configs[i + 1]);
                }});
        }};
    }

    private void checkHeartbeat() throws Exception
    {
        assertTrue(heartbeatFile.exists());
        Heartbeat heartbeat = Heartbeat.deserializeFromJsonFile(heartbeatFile);
        assertNotNull(heartbeat.lastHeartbeat);
        assertTrue(heartbeat.lastHeartbeat.toEpochMilli() < Global.currentTimeMillis());
    }
}
