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

import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.LongSupplier;

import javax.annotation.Nullable;

import accord.utils.Gens;
import accord.utils.RandomSource;

import org.apache.cassandra.exceptions.RequestFailure;
import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.utils.concurrent.AsyncPromise;
import org.apache.cassandra.utils.concurrent.Future;

public class SimulatedMessageDelivery implements MessageDelivery
{
    public enum Action { DELIVER, DELIVER_WITH_FAILURE, DROP, DROP_PARTITIONED, FAILURE }

    public interface ActionSupplier
    {
        Action get(InetAddressAndPort self, Message<?> message, InetAddressAndPort to);
    }

    public interface NetworkDelaySupplier
    {
        @Nullable
        Duration jitter(Message<?> message, InetAddressAndPort to);
    }

    public static NetworkDelaySupplier noDelay()
    {
        return (i1, i2) -> null;
    }

    public static NetworkDelaySupplier randomDelay(RandomSource rs)
    {
        class Connection
        {
            final InetAddressAndPort from, to;

            private Connection(InetAddressAndPort from, InetAddressAndPort to)
            {
                this.from = from;
                this.to = to;
            }

            @Override
            public boolean equals(Object o)
            {
                if (this == o) return true;
                if (o == null || getClass() != o.getClass()) return false;
                Connection that = (Connection) o;
                return from.equals(that.from) && to.equals(that.to);
            }

            @Override
            public int hashCode()
            {
                return Objects.hash(from, to);
            }

            @Override
            public String toString()
            {
                return "Connection{" + "from=" + from + ", to=" + to + '}';
            }
        }
        final Map<Connection, LongSupplier> networkLatencies = new HashMap<>();
        return (msg, to) -> {
            InetAddressAndPort from = msg.from();
            long delayNanos = networkLatencies.computeIfAbsent(new Connection(from, to), ignore -> {
                long min = TimeUnit.MICROSECONDS.toNanos(500);
                long maxSmall = TimeUnit.MILLISECONDS.toNanos(5);
                long max = TimeUnit.SECONDS.toNanos(5);
                LongSupplier small = () -> rs.nextLong(min, maxSmall);
                LongSupplier large = () -> rs.nextLong(maxSmall, max);
                return Gens.bools().biasedRepeatingRuns(rs.nextInt(1, 11) / 100.0D, rs.nextInt(3, 15))
                           .mapToLong(b -> b ? large.getAsLong() : small.getAsLong())
                           .asLongSupplier(rs.fork());
            }).getAsLong();
            return Duration.ofNanos(delayNanos);
        };
    }

    public interface Scheduler
    {
        void schedule(Runnable command, long delay, TimeUnit unit);
    }

    public interface DropListener
    {
        void onDrop(Action action, InetAddressAndPort from, Message<?> msg);
    }

    private final InetAddressAndPort self;
    private final ActionSupplier actions;
    private final NetworkDelaySupplier networkDelay;
    private final BiConsumer<InetAddressAndPort, Message<?>> reciever;
    private final DropListener onDropped;
    private final Scheduler scheduler;
    private final Consumer<Throwable> onError;
    private final Map<CallbackKey, CallbackContext> callbacks = new HashMap<>();
    private enum Status { Up, Down }
    private Status status = Status.Up;

    public SimulatedMessageDelivery(InetAddressAndPort self,
                                    ActionSupplier actions,
                                    NetworkDelaySupplier networkDelay,
                                    BiConsumer<InetAddressAndPort, Message<?>> reciever,
                                    DropListener onDropped,
                                    Scheduler scheduler,
                                    Consumer<Throwable> onError)
    {
        this.self = self;
        this.actions = actions;
        this.networkDelay = networkDelay;
        this.reciever = reciever;
        this.onDropped = onDropped;
        this.scheduler = scheduler;
        this.onError = onError;
    }

    public void stop()
    {
        callbacks.clear();
        status = Status.Down;
    }

    @Override
    public <REQ> void send(Message<REQ> message, InetAddressAndPort to)
    {
        message = message.withFrom(self);
        maybeEnqueue(message, to, null);
    }

    @Override
    public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb)
    {
        message = message.withFrom(self);
        maybeEnqueue(message, to, cb);
    }

    @Override
    public <REQ, RSP> void sendWithCallback(Message<REQ> message, InetAddressAndPort to, RequestCallback<RSP> cb, ConnectionType specifyConnection)
    {
        message = message.withFrom(self);
        maybeEnqueue(message, to, cb);
    }

    @Override
    public <REQ, RSP> Future<Message<RSP>> sendWithResult(Message<REQ> message, InetAddressAndPort to)
    {
        AsyncPromise<Message<RSP>> promise = new AsyncPromise<>();
        sendWithCallback(message, to, new RequestCallback<RSP>()
        {
            @Override
            public void onResponse(Message<RSP> msg)
            {
                promise.trySuccess(msg);
            }

            @Override
            public void onFailure(InetAddressAndPort from, RequestFailure failure)
            {
                promise.tryFailure(new MessagingService.FailureResponseException(from, failure));
            }

            @Override
            public boolean invokeOnFailure()
            {
                return true;
            }
        });
        return promise;
    }

    @Override
    public <V> void respond(V response, Message<?> message)
    {
        send(message.responseWith(response), message.respondTo());
    }

    private <REQ, RSP> void maybeEnqueue(Message<REQ> message, InetAddressAndPort to, @Nullable RequestCallback<RSP> callback)
    {
        if (status != Status.Up)
            return;
        CallbackContext cb;
        if (callback != null)
        {
            CallbackKey key = new CallbackKey(message.id(), to);
            if (callbacks.containsKey(key))
                throw new AssertionError("Message id " + message.id() + " to " + to + " already has a callback");
            cb = new CallbackContext(callback);
            callbacks.put(key, cb);
        }
        else
        {
            cb = null;
        }
        Action action = actions.get(self, message, to);
        switch (action)
        {
            case DELIVER:
                deliver(message, to);
                break;
            case DROP:
            case DROP_PARTITIONED:
                onDropped.onDrop(action, to, message);
                break;
            case DELIVER_WITH_FAILURE:
                deliver(message, to);
            case FAILURE:
                if (action == Action.FAILURE)
                    onDropped.onDrop(action, to, message);
                if (callback != null)
                    scheduler.schedule(() -> callback.onFailure(to, RequestFailure.UNKNOWN),
                                       message.verb().expiresAfterNanos(), TimeUnit.NANOSECONDS);
                return;
            default:
                throw new UnsupportedOperationException("Unknown action type: " + action);
        }
        if (cb != null)
        {
            scheduler.schedule(() -> {
                CallbackContext ctx = callbacks.remove(new CallbackKey(message.id(), to));
                if (ctx != null)
                {
                    assert ctx == cb;
                    try
                    {
                        ctx.onFailure(to, RequestFailure.TIMEOUT);
                    }
                    catch (Throwable t)
                    {
                        onError.accept(t);
                    }
                }
            }, message.verb().expiresAfterNanos(), TimeUnit.NANOSECONDS);
        }
    }

    private void deliver(Message<?> message, InetAddressAndPort to)
    {
        Duration delay = networkDelay.jitter(message, to);
        if (delay == null) reciever.accept(to, message);
        else               scheduler.schedule(() -> reciever.accept(to, message), delay.toNanos(), TimeUnit.NANOSECONDS);
    }

    @SuppressWarnings("rawtypes")
    public SimulatedMessageReceiver receiver(IVerbHandler onMessage)
    {
        return new SimulatedMessageReceiver(onMessage);
    }

    public class SimulatedMessageReceiver
    {
        @SuppressWarnings("rawtypes")
        final IVerbHandler onMessage;

        @SuppressWarnings("rawtypes")
        public SimulatedMessageReceiver(IVerbHandler onMessage)
        {
            this.onMessage = onMessage;
        }

        public void recieve(Message<?> msg)
        {
            if (status != Status.Up)
                return;
            if (msg.verb().isManagedResponse())
            {
                CallbackKey key = new CallbackKey(msg.id(), msg.from());
                if (callbacks.containsKey(key))
                {
                    CallbackContext callback = callbacks.remove(key);
                    if (callback == null)
                        return;
                    try
                    {
                        if (msg.isFailureResponse())
                            callback.onFailure(msg.from(), (RequestFailure) msg.payload);
                        else callback.onResponse(msg);
                    }
                    catch (Throwable t)
                    {
                        onError.accept(t);
                    }
                }
            }
            else
            {
                try
                {
                    //noinspection unchecked
                    onMessage.doVerb(msg);
                }
                catch (Throwable t)
                {
                    onError.accept(t);
                }
            }
        }
    }

    @SuppressWarnings("rawtypes")
    public static class SimpleVerbHandler implements IVerbHandler
    {
        private final Map<Verb, IVerbHandler<?>> handlers;

        public SimpleVerbHandler(Map<Verb, IVerbHandler<?>> handlers)
        {
            this.handlers = handlers;
        }

        @Override
        public void doVerb(Message msg) throws IOException
        {
            IVerbHandler<?> handler = handlers.get(msg.verb());
            if (handler == null)
                throw new AssertionError("Unexpected verb: " + msg.verb());
            //noinspection unchecked
            handler.doVerb(msg);
        }
    }

    private static class CallbackContext
    {
        @SuppressWarnings("rawtypes")
        final RequestCallback callback;

        @SuppressWarnings("rawtypes")
        private CallbackContext(RequestCallback callback)
        {
            this.callback = Objects.requireNonNull(callback);
        }

        @SuppressWarnings({ "rawtypes", "unchecked" })
        public void onResponse(Message msg)
        {
            callback.onResponse(msg);
        }

        public void onFailure(InetAddressAndPort from, RequestFailure failure)
        {
            if (callback.invokeOnFailure()) callback.onFailure(from, failure);
        }
    }

    private static class CallbackKey
    {
        private final long id;
        private final InetAddressAndPort peer;

        private CallbackKey(long id, InetAddressAndPort peer)
        {
            this.id = id;
            this.peer = peer;
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            CallbackKey that = (CallbackKey) o;
            return id == that.id && peer.equals(that.peer);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(id, peer);
        }

        @Override
        public String toString()
        {
            return "CallbackKey{" +
                   "id=" + id +
                   ", peer=" + peer +
                   '}';
        }
    }
}
