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

import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Random;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.Supplier;

import javax.annotation.Nullable;

import com.google.common.collect.Iterables;

import org.mockito.Mockito;

import accord.utils.Gen;
import accord.utils.Gens;
import accord.utils.Invariants;
import accord.utils.RandomSource;

import org.apache.cassandra.concurrent.ExecutorFactory;
import org.apache.cassandra.concurrent.ScheduledExecutorPlus;
import org.apache.cassandra.concurrent.SequentialExecutorPlus;
import org.apache.cassandra.concurrent.SimulatedExecutorFactory;
import org.apache.cassandra.concurrent.Stage;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.compaction.ICompactionManager;
import org.apache.cassandra.db.repair.CassandraTableRepairManager;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.gms.ApplicationState;
import org.apache.cassandra.gms.EndpointState;
import org.apache.cassandra.gms.HeartBeatState;
import org.apache.cassandra.gms.IEndpointStateChangeSubscriber;
import org.apache.cassandra.gms.IFailureDetector;
import org.apache.cassandra.gms.IGossiper;
import org.apache.cassandra.gms.VersionedValue;
import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.locator.Locator;
import org.apache.cassandra.net.IVerbHandler;
import org.apache.cassandra.net.MessageDelivery;
import org.apache.cassandra.net.SimulatedMessageDelivery;
import org.apache.cassandra.net.SimulatedMessageDelivery.ActionSupplier;
import org.apache.cassandra.repair.IValidationManager;
import org.apache.cassandra.repair.SharedContext;
import org.apache.cassandra.repair.StreamExecutor;
import org.apache.cassandra.repair.TableRepairManager;
import org.apache.cassandra.repair.ValidationManager;
import org.apache.cassandra.service.ActiveRepairService;
import org.apache.cassandra.service.paxos.cleanup.PaxosRepairState;
import org.apache.cassandra.streaming.StreamEventHandler;
import org.apache.cassandra.streaming.StreamState;
import org.apache.cassandra.tcm.ClusterMetadata;
import org.apache.cassandra.tcm.ClusterMetadataService;
import org.apache.cassandra.tcm.MultiStepOperation;
import org.apache.cassandra.tcm.StubClusterMetadataService;
import org.apache.cassandra.tcm.Transformation;
import org.apache.cassandra.tcm.membership.Directory;
import org.apache.cassandra.tcm.membership.Location;
import org.apache.cassandra.tcm.membership.NodeAddresses;
import org.apache.cassandra.tcm.membership.NodeId;
import org.apache.cassandra.tcm.membership.NodeVersion;
import org.apache.cassandra.tcm.ownership.UniformRangePlacement;
import org.apache.cassandra.tcm.transformations.PrepareJoin;

import static org.apache.cassandra.utils.AccordGenerators.fromQT;

public class SimulatedMiniCluster
{
    private final RandomSource rs;
    private final Function<Node, IVerbHandler<?>> verbHandlerFactory;
    private final SimulatedExecutorFactory executorFactory;
    private final SequentialExecutorPlus orderedExecutor;
    private final ScheduledExecutorPlus unorderedScheduled;
    private final IFailureDetector failureDetector = Mockito.mock(IFailureDetector.class);
    private final Locator locator = Mockito.mock(Locator.class);
    private final MBeanWrapper mbean = Mockito.mock(MBeanWrapper.class);
    private final SimulatedGossip gossiper = new SimulatedGossip();
    private final List<Throwable> failures = new ArrayList<>();
    private final IPartitioner partitioner;
    private final Map<String, List<String>> dcsToRacks;
    private final List<String> dcs;
    private final int tokensPerInstance;
    private final Gen<Token> tokenGen;
    private ClusterMetadata current;
    private final Map<NodeId, Node> nodes = new LinkedHashMap<>();
    private final TreeSet<Token> knownTokens = new TreeSet<>(); // includes bootstraping nodes tokens (aka tokens not in the ring)

    private SimulatedMiniCluster(Builder builder)
    {
        this.rs = builder.rs;
        this.verbHandlerFactory = builder.verbHandlerFactory;
        this.executorFactory = new SimulatedExecutorFactory(rs, failures::add);
        this.orderedExecutor = executorFactory.configureSequential("ignore").build();
        this.unorderedScheduled = executorFactory.scheduled("ignored");
        this.partitioner = fromQT(CassandraGenerators.nonLocalPartitioners()).next(rs);
        this.dcsToRacks = createDcRackDetails(rs);
        this.dcs = new ArrayList<>(dcsToRacks.keySet());
        dcs.sort(Comparator.naturalOrder());
        this.tokensPerInstance = rs.nextBoolean() ? 1 : 4;
        this.tokenGen = fromQT(CassandraGenerators.token(partitioner)).filter(t -> !knownTokens.contains(t));
        // setup Directory with known dcs
        this.current = new ClusterMetadata(partitioner);
        ClusterMetadataService.unsetInstance();
        ClusterMetadataService.setInstance(StubClusterMetadataService.forTesting(current));
    }

    public Node node(int id)
    {
        return node(new NodeId(id));
    }

    public Node node(NodeId id)
    {
        Node node = nodes.get(id);
        if (node == null)
            throw new AssertionError("Unable to find node for id " + id);
        return node;
    }

    public Node node(InetAddressAndPort address)
    {
        //TODO (performance): don't walk, keep index?
        for (Node node : nodes.values())
        {
            if (node.broadcastAddressAndPort.equals(address))
                return node;
        }
        throw new AssertionError("Unable to find node for address " + address);
    }

    private Collection<Token> nextUnknownTokens()
    {
        if (tokensPerInstance == 1) return Collections.singleton(tokenGen.next(rs));
        return Gens.lists(tokenGen).unique().ofSize(tokensPerInstance).next(rs);
    }

    public Node createNode()
    {
        if (nodes.isEmpty())
            return createFirstNode();

        NodeId id = new NodeId(nodes.size() + 1);
        UUID hostId = id.toUUID();
        Collection<Token> tokens = nextUnknownTokens();
        String dc = rs.pick(dcs);
        String rack = rs.pick(dcsToRacks.get(dc));
        Node node = new Node(id, hostId, address(id), tokens, dc, rack);
        register(node);
        return node;
    }

    public Node createNodeAndJoin()
    {
        if (nodes.isEmpty())
            return createFirstNode();

        NodeId id = new NodeId(nodes.size() + 1);
        UUID hostId = id.toUUID();
        Collection<Token> tokens = nextUnknownTokens();
        String dc = rs.pick(dcs);
        String rack = rs.pick(dcsToRacks.get(dc));
        Node node = new Node(id, hostId, address(id), tokens, dc, rack);
        registerAndJoin(node);
        return node;
    }

    private Node createFirstNode()
    {
        NodeId id = new NodeId(nodes.size() + 1);
        UUID hostId = id.toUUID();
        Collection<Token> tokens = nextUnknownTokens();
        String dc = dcs.get(0);
        String rack = dcsToRacks.get(dc).get(0);
        Node node = new Node(id, hostId, address(id), tokens, dc, rack);
        registerAndJoin(node);
        return node;
    }

    private void registerAndJoin(Node node)
    {
        register(node);
        prepareJoin(node.id);
        while (!current.inProgressSequences.isEmpty())
            bumpInProgress();
    }

    private void register(Node node)
    {
        nodes.put(node.id, node);
        knownTokens.addAll(node.tokens);
        registerWithSnitch(node);
        registerWithGossip(node);
        registerWithCMS(node);
    }

    private void registerWithCMS(Node node)
    {
        if (node.id.id() == 1)
        {
            // rebuild metadata from scratch
            Directory directory = Directory.EMPTY.with(new NodeAddresses(node.hostId, node.broadcastAddressAndPort, node.broadcastAddressAndPort, node.broadcastAddressAndPort), new Location(node.dc, node.rack));
            notifyMetadataChange(new ClusterMetadata(partitioner, directory));
        }
        else
        {
            notifyMetadataChange(current.transformer().register(new NodeAddresses(node.hostId, node.broadcastAddressAndPort, node.broadcastAddressAndPort, node.broadcastAddressAndPort),
                                                                new Location(node.dc, node.rack),
                                                                NodeVersion.CURRENT)
                                        .build().metadata);
        }
    }

    private void prepareJoin(NodeId id)
    {
        Node node = nodes.get(id);
        if (node == null)
            throw new IllegalArgumentException("Unknown " + id);
        PrepareJoin task = new PrepareJoin(id, new HashSet<>(node.tokens), new UniformRangePlacement(), true, false);
        notifyMetadataChange(process(task).metadata);
    }

    private void bumpInProgress()
    {
        if (current.inProgressSequences.isEmpty())
            throw new IllegalStateException("Attempted to bump epoch when nothing was pending");
        Iterator<MultiStepOperation<?>> it = current.inProgressSequences.iterator();
        Invariants.require(it.hasNext());
        notifyMetadataChange(process(it.next()).metadata);
    }

    protected void notifyMetadataChange(ClusterMetadata current)
    {
        this.current = current;
        ((StubClusterMetadataService) ClusterMetadataService.instance()).setMetadata(current);
    }

    private Transformation.Success process(Transformation transformation)
    {
        Transformation.Result result = transformation.execute(current);
        if (result.isRejected())
            throw new IllegalStateException("Unable to make TCM transition");
        return result.success();
    }

    private Transformation.Success process(MultiStepOperation<?> transformation)
    {
        Transformation.Result result = transformation.applyTo(current);
        if (result.isRejected())
            throw new IllegalStateException("Unable to make TCM transition");
        return result.success();
    }

    private static InetAddressAndPort address(NodeId id)
    {
        try
        {
            return InetAddressAndPort.getByAddress(ByteArrayUtil.bytes(id.id()));
        }
        catch (UnknownHostException e)
        {
            throw new AssertionError("Unable to create address for id " + id, e);
        }
    }

    private static Map<String, List<String>> createDcRackDetails(RandomSource rs)
    {
        int numDCs = rs.nextInt(1, 4);
        Map<String, List<String>> map = new LinkedHashMap<>();
        for (int i = 0; i < numDCs; i++)
        {
            String name = "DC" + (i + 1);
            int numRacks = rs.nextInt(1, 10);
            List<String> racks = Gens.lists(Gens.strings().ascii().ofLength(5).map(s -> "R" + s)).unique().ofSize(numRacks).next(rs);
            racks.sort(Comparator.naturalOrder());
            map.put(name, racks);
        }
        return map;
    }

    public boolean hasWork()
    {
        return executorFactory.hasWork();
    }

    public boolean processAny()
    {
        return executorFactory.processAny();
    }

    public boolean processOne()
    {
        return executorFactory.processOne();
    }

    public void processAll()
    {
        executorFactory.processAll();;
    }

    public void simulateStages(Stage... stages)
    {
        for (Stage stage : stages)
        {
            switch (stage)
            {
                case GOSSIP:
                case ANTI_ENTROPY:
                case MIGRATION:
                case MISC:
                case TRACING:
                case FETCH_METADATA:
                    stage.unsafeSetExecutor(orderedExecutor);
                    break;
                default:
                    stage.unsafeSetExecutor(unorderedScheduled);
            }
        }
    }

    private void registerWithSnitch(Node node)
    {
        Mockito.when(locator.location(Mockito.eq(node.broadcastAddressAndPort))).thenReturn(new Location(node.dc, node.rack));
    }

    private void registerWithGossip(Node node)
    {
        VersionedValue.VersionedValueFactory valueFactory = node.valueFactory;
        EndpointState state = new EndpointState(new HeartBeatState(42, 42));
        state.addApplicationState(ApplicationState.STATUS, valueFactory.normal(node.tokens));
        state.addApplicationState(ApplicationState.STATUS_WITH_PORT, valueFactory.normal(node.tokens));
        state.addApplicationState(ApplicationState.HOST_ID, valueFactory.hostId(node.hostId));
        state.addApplicationState(ApplicationState.TOKENS, valueFactory.tokens(node.tokens));
        state.addApplicationState(ApplicationState.DC, valueFactory.datacenter(node.dc));
        state.addApplicationState(ApplicationState.RACK, valueFactory.rack(node.rack));
        state.addApplicationState(ApplicationState.RELEASE_VERSION, valueFactory.releaseVersion());

        gossiper.endpoints.put(node.broadcastAddressAndPort, state);
    }

    public static class Builder
    {
        private final RandomSource rs;
        private final Function<Node, IVerbHandler<?>> verbHandlerFactory;

        public Builder(RandomSource rs, Function<Node, IVerbHandler<?>> verbHandlerFactory)
        {
            this.rs = rs;
            this.verbHandlerFactory = verbHandlerFactory;
        }

        public SimulatedMiniCluster build()
        {
            return new SimulatedMiniCluster(this);
        }
    }

     private enum NodeStatus { Init, Registered, Joining, Joined, Leaving, Removed}

    public class Node implements SharedContext
    {
        private final ICompactionManager compactionManager = Mockito.mock(ICompactionManager.class);
        private final NodeId id;
        private final UUID hostId;
        private final InetAddressAndPort broadcastAddressAndPort;
        private final Collection<Token> tokens;
        private final String dc, rack;
        private final VersionedValue.VersionedValueFactory valueFactory;
        private final SimulatedMessageDelivery messaging;
        private final SimulatedMessageDelivery.SimulatedMessageReceiver receiver;
        private final ActiveRepairService activeRepairService;
        private final PaxosRepairState paxosRepairState;
        private final IValidationManager validationManager;
        private final StreamExecutor streamExecutor;
        private NodeStatus status = NodeStatus.Init;
        private ActionSupplier messagingActions = (self, msg, to) -> SimulatedMessageDelivery.Action.DELIVER;

        public Node(NodeId id, UUID hostId, InetAddressAndPort broadcastAddressAndPort, Collection<Token> tokens, String dc, String rack)
        {
            this.id = id;
            this.hostId = hostId;
            this.broadcastAddressAndPort = broadcastAddressAndPort;
            this.tokens = tokens;
            this.dc = dc;
            this.rack = rack;

            IPartitioner partitioner = Iterables.getFirst(tokens, null).getPartitioner();
            this.valueFactory = new VersionedValue.VersionedValueFactory(partitioner);
            this.messaging = new SimulatedMessageDelivery(broadcastAddressAndPort,
                                                          messagingActions::get,
                                                          SimulatedMessageDelivery.randomDelay(rs),
                                                          (to, msg) -> unorderedScheduled.submit(() -> node(to).receiver.recieve(msg)),
                                                          (action, to, msg) -> {},
                                                          unorderedScheduled::schedule,
                                                          failures::add);
            this.activeRepairService = new ActiveRepairService(this);
            this.paxosRepairState = new PaxosRepairState(this);
            this.validationManager = (cfs, validator) -> unorderedScheduled.submit(() -> {
                try
                {
                    ValidationManager.doValidation(cfs, validator);
                }
                catch (Throwable e)
                {
                    validator.fail(e);
                }
            });
            this.streamExecutor = plan -> {
                long delayNanos = rs.nextLong(TimeUnit.SECONDS.toNanos(5), TimeUnit.MINUTES.toNanos(10));
                unorderedScheduled.schedule(() -> {
                    StreamState success = new StreamState(plan.planId(), plan.streamOperation(), Collections.emptySet());
                    for (StreamEventHandler handler : plan.handlers())
                        handler.onSuccess(success);
                }, delayNanos, TimeUnit.NANOSECONDS);
                return null;
            };

            // setup last as "this" is leaking, so make sure all final fields are defined first
            this.receiver = messaging.receiver(verbHandlerFactory.apply(this));
        }

        public NodeId id()
        {
            return id;
        }

        public UUID hostId()
        {
            return hostId;
        }

        public void messagingActions(ActionSupplier messagingActions)
        {
            this.messagingActions = Objects.requireNonNull(messagingActions);
        }

        @Override
        public InetAddressAndPort broadcastAddressAndPort()
        {
            return broadcastAddressAndPort;
        }

        @Override
        public Supplier<Random> random()
        {
            return () -> rs.fork().asJdkRandom();
        }

        @Override
        public Clock clock()
        {
            return executorFactory;
        }

        @Override
        public ExecutorFactory executorFactory()
        {
            return executorFactory;
        }

        @Override
        public MBeanWrapper mbean()
        {
            return mbean;
        }

        @Override
        public ScheduledExecutorPlus optionalTasks()
        {
            return unorderedScheduled;
        }

        @Override
        public ScheduledExecutorPlus nonPeriodicTasks()
        {
            return unorderedScheduled;
        }

        @Override
        public ScheduledExecutorPlus scheduledTasks()
        {
            return unorderedScheduled;
        }


        @Override
        public IFailureDetector failureDetector()
        {
            return failureDetector;
        }

        @Override
        public Locator locator()
        {
            return locator;
        }

        @Override
        public IGossiper gossiper()
        {
            return gossiper;
        }

        @Override
        public MessageDelivery messaging()
        {
            return messaging;
        }

        @Override
        public ActiveRepairService repair()
        {
            return activeRepairService;
        }

        @Override
        public PaxosRepairState paxosRepairState()
        {
            return paxosRepairState;
        }

        @Override
        public ICompactionManager compactionManager()
        {
            return compactionManager;
        }

        @Override
        public IValidationManager validationManager()
        {
            return validationManager;
        }

        @Override
        public TableRepairManager repairManager(ColumnFamilyStore store)
        {
            return new CassandraTableRepairManager(store, this)
            {
                @Override
                public void snapshot(String name, Collection<Range<Token>> ranges, boolean force)
                {
                    // no-op
                }
            };
        }

        @Override
        public StreamExecutor streamExecutor()
        {
            return streamExecutor;
        }
    }

    private class SimulatedGossip implements IGossiper
    {
        private final Map<InetAddressAndPort, EndpointState> endpoints = new HashMap<>();

        @Override
        public void register(IEndpointStateChangeSubscriber subscriber)
        {

        }

        @Override
        public void unregister(IEndpointStateChangeSubscriber subscriber)
        {

        }

        @Nullable
        @Override
        public EndpointState getEndpointStateForEndpoint(InetAddressAndPort ep)
        {
            return endpoints.get(ep);
        }

        @Override
        public void notifyFailureDetector(Map<InetAddressAndPort, EndpointState> remoteEpStateMap)
        {

        }

        @Override
        public void applyStateLocally(Map<InetAddressAndPort, EndpointState> epStateMap)
        {
            // If we were testing paxos this would be wrong...
            // CASSANDRA-18917 added support for simulating Gossip, but gossip issues were found so couldn't merge that patch...
            // For the paxos repair, since we don't care about paxos messages, this is ok to no-op for now, but if paxos cleanup
            // ever was to be tested this logic would need to be implemented
        }
    }
}
