/*
 * 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 accord.impl.mock;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiFunction;
import java.util.function.LongFunction;
import java.util.function.ToLongFunction;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import accord.NetworkFilter;
import accord.api.Agent;
import accord.api.AsyncExecutor;
import accord.api.Journal;
import accord.api.MessageSink;
import accord.coordinate.CoordinationAdapter;
import accord.impl.DefaultTimeouts;
import accord.impl.InMemoryCommandStores;
import accord.impl.IntKey;
import accord.impl.DefaultLocalListeners;
import accord.impl.basic.InMemoryJournal;
import accord.impl.progresslog.DefaultProgressLogs;
import accord.impl.DefaultRemoteListeners;
import accord.impl.SizeOfIntersectionSorter;
import accord.impl.TestAgent;
import accord.local.DurableBefore;
import accord.local.Node;
import accord.local.Node.Id;
import accord.local.ShardDistributor;
import accord.local.TimeService;
import accord.local.UniqueTimeService;
import accord.messages.Callback;
import accord.messages.Reply;
import accord.messages.Request;
import accord.primitives.Ranges;
import accord.primitives.TxnId;
import accord.topology.Topology;
import accord.topology.TopologyUtils;
import accord.utils.DefaultRandom;
import accord.utils.Invariants;
import accord.utils.RandomSource;
import accord.utils.ThreadPoolScheduler;
import accord.utils.async.Cancellable;

import static accord.Utils.id;
import static accord.Utils.idList;
import static accord.local.TimeService.elapsedWrapperFromNonMonotonicSource;
import static accord.primitives.Routable.Domain.Key;
import static accord.primitives.Txn.Kind.Write;
import static accord.utils.async.AsyncChainUtils.awaitUninterruptibly;

public class MockCluster implements Network, AutoCloseable, Iterable<Node>
{
    private static final Logger logger = LoggerFactory.getLogger(MockCluster.class);

    private final RandomSource random;
    private final Config config;
    private final TimeService time;
    private final Map<Id, Node> nodes = new ConcurrentHashMap<>();
    private final BiFunction<Id, Network, MessageSink> messageSinkFactory;
    private int nextNodeId = 1;
    public NetworkFilter networkFilter = new NetworkFilter();

    private long nextMessageId = 0;
    Map<Long, Callback> callbacks = new ConcurrentHashMap<>();
    private final LongFunction<Topology> fetchTopology;

    private MockCluster(Builder builder)
    {
        this.config = new Config(builder);
        this.random = new DefaultRandom(config.seed);
        this.time = builder.time;
        this.messageSinkFactory = builder.messageSinkFactory;
        this.fetchTopology = builder.fetchTopology;

        init(builder.topology);
    }

    @Override
    public Iterator<Node> iterator()
    {
        return nodes.values().iterator();
    }

    @Override
    public void close()
    {
        nodes.values().forEach(n -> {
            ((ThreadPoolScheduler)n.scheduler()).stop();
            n.shutdown();
        });

    }

    private synchronized Id nextNodeId()
    {
        return id(nextNodeId++);
    }

    private synchronized long nextMessageId()
    {
        return nextMessageId++;
    }

    private Node createNode(Id id, Topology topology)
    {
        MockStore store = new MockStore();
        MessageSink messageSink = messageSinkFactory.apply(id, this);
        MockTopologyService configurationService = new MockTopologyService(fetchTopology, topology);
        Agent agent = new TestAgent(time);
        Journal journal = new InMemoryJournal(id, random.fork());
        ThreadPoolScheduler scheduler = new ThreadPoolScheduler();
        Node node = new Node(id,
                             messageSink,
                             configurationService,
                             time, new UniqueTimeService.AtomicUniqueTime(time),
                             () -> store,
                             new ShardDistributor.EvenSplit(8, ignore -> new IntKey.Splitter()),
                             agent,
                             random.fork(),
                             scheduler,
                             SizeOfIntersectionSorter.SUPPLIER,
                             DefaultRemoteListeners::new,
                             time -> new DefaultTimeouts(time, Runnable::run),
                             DefaultProgressLogs::new,
                             DefaultLocalListeners.Factory::new,
                             InMemoryCommandStores.SingleThread::new,
                             new CoordinationAdapter.DefaultFactory(),
                             DurableBefore.NOOP_PERSISTER,
                             journal);
        journal.start(node);
        awaitUninterruptibly(node.unsafeStart());
        node.durability().start();
        node.topology().reportTopology(topology);
        return node;
    }

    private void init(Topology topology)
    {
        List<Id> ids = new ArrayList<>(config.initialNodes);
        for (int i=0; i<config.initialNodes; i++)
        {
            Id nextId = nextNodeId();
            ids.add(nextId);
        }
        if (topology == null)
        {
            Ranges ranges = TopologyUtils.initialRanges(config.initialNodes, config.maxKey);
            topology = TopologyUtils.initialTopology(ids, ranges, config.replication);
        }
        for (int i=0; i<config.initialNodes; i++)
        {
            Id id = ids.get(i);
            Node node = createNode(id, topology);
            nodes.put(id, node);
        }
    }

    @Override
    public Cancellable send(Id from, Id to, Request request, AsyncExecutor executor, Callback callback)
    {
        Node node = nodes.get(to);
        if (node == null)
        {
            logger.info("dropping message to unknown node {}: {} from {}", to, request, from);
            return () -> {};
        }

        if (networkFilter.shouldDiscard(from, to, request))
        {
            // TODO (desired, testing): more flexible timeouts
            if (callback != null)
                callback.onFailure(to, null);
            logger.info("discarding filtered message from {} to {}: {}", from, to, request);
            return null;
        }

        long messageId = nextMessageId();
        if (callback != null)
        {
            callbacks.put(messageId, callback);
            node.scheduler().once(() -> {
                if (callbacks.remove(messageId, callback))
                    callback.onFailure(to, null);
                }, 2L, TimeUnit.SECONDS);
        }

        logger.info("processing message[{}] from {} to {}: {}", messageId, from, to, request);
        node.receive(request, from, Network.replyCtxFor(messageId));
        return () -> callbacks.remove(messageId);
    }

    @Override
    public void reply(Id from, Id replyingToNode, long replyingToMessage, Reply reply)
    {
        Node node = nodes.get(replyingToNode);
        if (node == null)
        {
            logger.info("dropping reply to unknown node {}: {} from {}", replyingToNode, reply, from);
            return;
        }

        Callback sc = reply.isFinal() ? callbacks.remove(replyingToMessage) : callbacks.get(replyingToMessage);

        if (networkFilter.shouldDiscard(from, replyingToNode, reply))
        {
            logger.info("discarding filtered reply from {} to {}: {}", from, reply, reply);
            if (sc != null)
                sc.onFailure(from, null);
            return;
        }

        if (sc == null)
        {
            logger.warn("Callback not found for reply from {} to {}: {} (msgid: {})", from, replyingToNode, reply, replyingToMessage);
            return;
        }

        logger.info("processing reply[{}] from {} to {}: {}", replyingToMessage, from, replyingToNode, reply);
        node.scheduler().now(() -> {
            try
            {
                if (reply instanceof Reply.FailureReply) sc.onFailure(from, ((Reply.FailureReply) reply).failure);
                else sc.onSuccess(from, reply);
            }
            catch (Throwable t)
            {
                sc.onFailure(from, t);
            }
        });
    }

    public Node get(Id id)
    {
        Node node = nodes.get(id);
        if (node == null)
            throw new NoSuchElementException("No node exists with id " + id);
        return node;
    }

    public Node get(int i)
    {
        return get(id(i));
    }

    public List<Node> nodes(Iterable<Id> ids)
    {
        List<Node> rlist = new ArrayList<>();
        for (Id id : ids)
            rlist.add(get(id));
        return rlist;
    }

    public Iterable<Node> nodes(int... ids)
    {
        assert ids.length > 0;
        return nodes(idList(ids));
    }

    public static class Config
    {
        private final long seed;
        private final int initialNodes;
        private final int replication;
        private final int maxKey;

        private Config(MockCluster.Builder builder)
        {
            this.seed = builder.seed;
            this.initialNodes = builder.initialNodes;
            this.replication = builder.replication;
            this.maxKey = builder.maxKey;
        }
    }

    public static class Builder
    {
        private long seed = 0;
        private int initialNodes = 3;
        private int replication = 3;
        private int maxKey = 10000;
        private Topology topology = null;
        private TimeService time = TimeService.ofNonMonotonic(System::currentTimeMillis, TimeUnit.MILLISECONDS);
        private BiFunction<Id, Network, MessageSink> messageSinkFactory = SimpleMessageSink::new;
        private LongFunction<Topology> fetchTopology = ignore -> null;

        public Builder seed(long seed)
        {
            this.seed = seed;
            return this;
        }

        public Builder nodes(int initialNodes)
        {
            this.initialNodes = initialNodes;
            return this;
        }

        public Builder replication(int replication)
        {
            this.replication = replication;
            return this;
        }

        public Builder maxKey(int max)
        {
            this.maxKey = max;
            return this;
        }

        public Builder time(TimeService time)
        {
            this.time = time;
            return this;
        }

        public Builder topology(Topology topology)
        {
            this.topology = topology;
            return this;
        }

        public Builder messageSink(BiFunction<Id, Network, MessageSink> factory)
        {
            this.messageSinkFactory = factory;
            return this;
        }

        public Builder setFetchTopology(LongFunction<Topology> fetchTopology)
        {
            this.fetchTopology = fetchTopology;
            return this;
        }

        public MockCluster build()
        {
            Invariants.requireArgument(initialNodes > 0);
            Invariants.requireArgument(replication > 0);
            Invariants.requireArgument(maxKey >= 0);
            return new MockCluster(this);
        }
    }

    public static Builder builder()
    {
        return new Builder();
    }

    public static class Clock implements TimeService
    {
        private final AtomicLong now;
        private final ToLongFunction<TimeUnit> elapsed;

        public Clock(long now)
        {
            this.now = new AtomicLong(now);
            this.elapsed = elapsedWrapperFromNonMonotonicSource(TimeUnit.MILLISECONDS, this::now);
        }

        public long increment(long by)
        {
            return now.addAndGet(by);
        }

        public long increment()
        {
            return increment(1);
        }

        public long now()
        {
            return now.get();
        }

        @Override
        public long elapsed(TimeUnit unit)
        {
            return elapsed.applyAsLong(unit);
        }


        public TxnId idForNode(long epoch, Id id)
        {
            return new TxnId(epoch, now.get(), 0, Write, Key, id);
        }

        public TxnId idForNode(long epoch, int id)
        {
            return idForNode(epoch, new Id(id));
        }
    }
}
