/*
 * Copyright (c) 2019 the Eclipse Milo Authors
 *
 * This program and the accompanying materials are made
 * available under the terms of the Eclipse Public License 2.0
 * which is available at https://www.eclipse.org/legal/epl-2.0/
 *
 * SPDX-License-Identifier: EPL-2.0
 */

package org.eclipse.milo.opcua.stack.server.transport;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.LinkedList;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import com.google.common.collect.ConcurrentHashMultiset;
import com.google.common.collect.Maps;
import com.google.common.collect.Multiset;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.ipfilter.AbstractRemoteAddressFilter;
import org.eclipse.milo.opcua.stack.core.Stack;
import org.eclipse.milo.opcua.stack.server.UaStackServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A shared, stack-wide, one-per-application-regardless-of-how-many-server-instances-you-have handler that is added to
 * the beginning of every server pipeline to handle rate limiting and connection limits.
 * <p>
 * Any configuration changes must be made before {@link UaStackServer#startup()} is called for the first time,
 * application-wide. Once the instance has been created further configuration changes will have no effect.
 */
@ChannelHandler.Sharable
public class RateLimitingHandler extends AbstractRemoteAddressFilter<InetSocketAddress> {

    /**
     * Cumulative count of all connection rejections for the lifetime of the server.
     */
    @SuppressWarnings("WeakerAccess")
    public static final AtomicLong CUMULATIVE_CONNECTIONS_REJECTED = new AtomicLong(0L);

    /**
     * Get the shared {@link RateLimitingHandler} instance.
     * <p>
     * The values of {@link Stack.ConnectionLimits#RATE_LIMIT_ENABLED},
     * {@link Stack.ConnectionLimits#RATE_LIMIT_MAX_ATTEMPTS},
     * {@link Stack.ConnectionLimits#RATE_LIMIT_WINDOW_MS},
     * {@link Stack.ConnectionLimits#RATE_LIMIT_MAX_CONNECTIONS},
     * and {@link Stack.ConnectionLimits#RATE_LIMIT_MAX_CONNECTIONS_PER_ADDRESS}
     * will be locked in whenever the first invocation of this method occurs.
     *
     * @return the shared {@link RateLimitingHandler} instance.
     */
    public static RateLimitingHandler getInstance() {
        return InstanceHolder.INSTANCE;
    }

    private static class InstanceHolder {

        private static final RateLimitingHandler INSTANCE = new RateLimitingHandler(
            Stack.ConnectionLimits.RATE_LIMIT_ENABLED,
            Stack.ConnectionLimits.RATE_LIMIT_MAX_ATTEMPTS,
            Stack.ConnectionLimits.RATE_LIMIT_WINDOW_MS,
            Stack.ConnectionLimits.RATE_LIMIT_MAX_CONNECTIONS,
            Stack.ConnectionLimits.RATE_LIMIT_MAX_CONNECTIONS_PER_ADDRESS
        );

    }

    private final Logger logger = LoggerFactory.getLogger(getClass());

    private final Multiset<InetAddress> connections = ConcurrentHashMultiset.create();
    private final ConcurrentMap<InetAddress, LinkedList<Long>> timestamps = Maps.newConcurrentMap();

    private final boolean enabled;
    private final int maxAttempts;
    private final int rateLimitWindowMs;
    private final int maxConnections;
    private final int maxConnectionsPerAddress;

    private RateLimitingHandler(
        boolean enabled,
        int maxAttempts,
        int rateLimitWindowMs,
        int maxConnections,
        int maxConnectionsPerAddress) {

        this.enabled = enabled;
        this.maxAttempts = maxAttempts;
        this.rateLimitWindowMs = rateLimitWindowMs;
        this.maxConnections = maxConnections;
        this.maxConnectionsPerAddress = maxConnectionsPerAddress;

        logger.debug(String.format(
            "enabled=%s, maxAttempts=%s, rateLimitWindowMs=%s, maxConnections=%s, maxConnectionsPerAddress=%s",
            enabled, maxAttempts, rateLimitWindowMs, maxConnections, maxConnectionsPerAddress));
    }

    @Override
    protected synchronized boolean accept(ChannelHandlerContext ctx, InetSocketAddress isa) {
        final InetAddress address = isa.getAddress();

        if (!enabled || address.isLoopbackAddress()) {
            return true;
        }

        LinkedList<Long> attempts = timestamps.computeIfAbsent(address, ia -> new LinkedList<>());

        long now = System.currentTimeMillis();

        if (attempts.size() >= maxAttempts) {
            // count the number of previous connections from this address
            // that have occurred within the rate limit window.
            int attemptsInWindow = 0;
            for (Long ts : attempts) {
                if (now - ts < rateLimitWindowMs) {
                    attemptsInWindow++;
                }
            }

            attempts.addLast(now);
            while (attempts.size() > maxAttempts) {
                attempts.removeFirst();
            }

            int connectionsTotal = connections.size();
            int connectionsFromAddress = connections.count(address);

            boolean accept = attemptsInWindow < maxAttempts &&
                connectionsTotal < maxConnections &&
                connectionsFromAddress < maxConnectionsPerAddress;

            if (accept) {
                logger.debug(String.format(
                    "Accepting connection from %s. " +
                        "window=%sms, attemptsInWindow=%s, connectionsTotal=%s, connectionsFromAddress=%s",
                    isa, rateLimitWindowMs, attemptsInWindow, connectionsTotal, connectionsFromAddress));
            } else {
                logger.debug(String.format(
                    "Rejecting connection from %s. " +
                        "window=%sms, attemptsInWindow=%s, connectionsTotal=%s, connectionsFromAddress=%s",
                    isa, rateLimitWindowMs, attemptsInWindow, connectionsTotal, connectionsFromAddress));

                long cumulativeConnectionsRejected = CUMULATIVE_CONNECTIONS_REJECTED.incrementAndGet();
                logger.debug("cumulativeConnectionsRejected=" + cumulativeConnectionsRejected);
            }

            return accept;
        } else {
            attempts.addLast(now);

            return true;
        }
    }

    @Override
    protected void channelAccepted(ChannelHandlerContext ctx, InetSocketAddress remoteAddress) {
        final InetAddress address = remoteAddress.getAddress();

        if (!enabled || address.isLoopbackAddress()) {
            return;
        }

        connections.add(address);

        ctx.channel().closeFuture().addListener((ChannelFutureListener) future -> {
            connections.remove(address);

            if (connections.count(address) == 0) {
                logger.debug("Scheduling timestamp removal for " + address);

                ctx.executor().schedule(
                    () -> {
                        // If there's still no connections from the remote address after the rate limit window remove
                        // the timestamps.
                        // Removing them before the window elapses would allow a remote address to connect/disconnect
                        // at an effectively unlimited rate.
                        if (connections.count(address) == 0) {
                            timestamps.remove(address);
                            logger.debug("Removed timestamps for " + address);
                        }
                    },
                    rateLimitWindowMs,
                    TimeUnit.MILLISECONDS
                );
            }
        });
    }

}
