/*
 * Copyright 2020 The Netty Project
 *
 * The Netty Project 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:
 *
 *   https://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 io.netty.incubator.codec.quic;

import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.channel.MessageSizeEstimator;
import io.netty.channel.socket.DatagramPacket;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;
import org.jetbrains.annotations.Nullable;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.HashSet;
import java.util.Queue;
import java.util.Set;
import java.util.function.Consumer;

import static io.netty.incubator.codec.quic.Quiche.allocateNativeOrder;

/**
 * Abstract base class for QUIC codecs.
 */
abstract class QuicheQuicCodec extends ChannelDuplexHandler {
    private static final InternalLogger LOGGER = InternalLoggerFactory.getInstance(QuicheQuicCodec.class);
    private final ConnectionIdChannelMap connectionIdToChannel = new ConnectionIdChannelMap();
    private final Set<QuicheQuicChannel> channels = new HashSet<>();
    private final Queue<QuicheQuicChannel> needsFireChannelReadComplete = new ArrayDeque<>();
    private final Queue<QuicheQuicChannel> delayedRemoval = new ArrayDeque<>();

    private final Consumer<QuicheQuicChannel> freeTask = this::removeChannel;
    private final FlushStrategy flushStrategy;
    private final int localConnIdLength;
    private final QuicheConfig config;

    private MessageSizeEstimator.Handle estimatorHandle;
    private QuicHeaderParser headerParser;
    private QuicHeaderParser.QuicHeaderProcessor parserCallback;
    private int pendingBytes;
    private int pendingPackets;
    private boolean inChannelReadComplete;
    private boolean delayRemoval;

    // This buffer is used to copy InetSocketAddress to sockaddr_storage and so pass it down the JNI layer.
    private ByteBuf senderSockaddrMemory;
    private ByteBuf recipientSockaddrMemory;

    QuicheQuicCodec(QuicheConfig config, int localConnIdLength, FlushStrategy flushStrategy) {
        this.config = config;
        this.localConnIdLength = localConnIdLength;
        this.flushStrategy = flushStrategy;
    }

    @Override
    public final boolean isSharable() {
        return false;
    }

    @Nullable
    protected final QuicheQuicChannel getChannel(ByteBuffer key) {
        return connectionIdToChannel.get(key);
    }

    private void addMapping(QuicheQuicChannel channel, ByteBuffer id) {
        QuicheQuicChannel ch = connectionIdToChannel.put(id, channel);
        assert ch == null;
    }

    private void removeMapping(QuicheQuicChannel channel, ByteBuffer id) {
        QuicheQuicChannel ch = connectionIdToChannel.remove(id);
        assert ch == channel;
    }

    private void processDelayedRemoval() {
        for (;;) {
            // Now remove all channels that we marked for removal.
            QuicheQuicChannel toBeRemoved = delayedRemoval.poll();
            if (toBeRemoved == null) {
                break;
            }
            removeChannel(toBeRemoved);
        }
    }

    private void removeChannel(QuicheQuicChannel channel) {
        if (delayRemoval) {
            boolean added = delayedRemoval.offer(channel);
            assert added;
        } else {
            boolean removed = channels.remove(channel);
            if (removed) {
                for (ByteBuffer id : channel.sourceConnectionIds()) {
                    QuicheQuicChannel ch = connectionIdToChannel.remove(id);
                    assert ch == channel;
                }
            }
        }
    }

    protected final void addChannel(QuicheQuicChannel channel) {
        boolean added = channels.add(channel);
        assert added;
        for (ByteBuffer id : channel.sourceConnectionIds()) {
            QuicheQuicChannel ch = connectionIdToChannel.put(id.duplicate(), channel);
            assert ch == null;
        }
    }

    @Override
    public final void handlerAdded(ChannelHandlerContext ctx) {
        senderSockaddrMemory = allocateNativeOrder(Quiche.SIZEOF_SOCKADDR_STORAGE);
        recipientSockaddrMemory = allocateNativeOrder(Quiche.SIZEOF_SOCKADDR_STORAGE);
        headerParser = new QuicHeaderParser(localConnIdLength);
        parserCallback = new QuicCodecHeaderProcessor(ctx);
        estimatorHandle = ctx.channel().config().getMessageSizeEstimator().newHandle();
        handlerAdded(ctx, localConnIdLength);
    }

    /**
     * See {@link io.netty.channel.ChannelHandler#handlerAdded(ChannelHandlerContext)}.
     */
    protected void handlerAdded(ChannelHandlerContext ctx, int localConnIdLength) {
        // NOOP.
    }

    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) {
        try {
            // Use a copy of the array as closing the channel may cause an unwritable event that could also
            // remove channels.
            for (QuicheQuicChannel ch : channels.toArray(new QuicheQuicChannel[0])) {
                ch.forceClose();
            }
            if (pendingPackets > 0) {
                flushNow(ctx);
            }
        } finally {
            channels.clear();
            connectionIdToChannel.clear();
            needsFireChannelReadComplete.clear();
            delayedRemoval.clear();

            config.free();
            if (senderSockaddrMemory != null) {
                senderSockaddrMemory.release();
            }
            if (recipientSockaddrMemory != null) {
                recipientSockaddrMemory.release();
            }
            if (headerParser != null) {
                headerParser.close();
                headerParser = null;
            }
        }
    }

    @Override
    public final void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        DatagramPacket packet = (DatagramPacket) msg;
        try {
            ByteBuf buffer = ((DatagramPacket) msg).content();
            if (!buffer.isDirect()) {
                // We need a direct buffer as otherwise we can not access the memoryAddress.
                // Let's do a copy to direct memory.
                ByteBuf direct = ctx.alloc().directBuffer(buffer.readableBytes());
                try {
                    direct.writeBytes(buffer, buffer.readerIndex(), buffer.readableBytes());
                    handleQuicPacket(packet.sender(), packet.recipient(), direct);
                } finally {
                    direct.release();
                }
            } else {
                handleQuicPacket(packet.sender(), packet.recipient(), buffer);
            }
        } finally {
            packet.release();
        }
    }

    private void handleQuicPacket(InetSocketAddress sender, InetSocketAddress recipient, ByteBuf buffer) {
        try {
            headerParser.parse(sender, recipient, buffer, parserCallback);
        } catch (Exception e) {
            LOGGER.debug("Error while processing QUIC packet", e);
        }
    }

    /**
     * Handle a QUIC packet and return {@link QuicheQuicChannel} that is mapped to the id.
     *
     * @param ctx the {@link ChannelHandlerContext}.
     * @param sender the {@link InetSocketAddress} of the sender of the QUIC packet
     * @param recipient the {@link InetSocketAddress} of the recipient of the QUIC packet
     * @param type the type of the packet.
     * @param version the QUIC version
     * @param scid the source connection id.
     * @param dcid the destination connection id
     * @param token the token
     * @param senderSockaddrMemory the {@link ByteBuf} that can be used for the sender {@code struct sockaddr).
     * @param recipientSockaddrMemory the {@link ByteBuf} that can be used for the recipient {@code struct sockaddr).
     * @param freeTask the {@link Consumer} that will be called once native memory of the {@link QuicheQuicChannel} is
     *                  freed and so the mappings should be deleted to the ids.
     * @param localConnIdLength the length of the local connection ids.
     * @param config the {@link QuicheConfig} that is used.
     * @return the {@link QuicheQuicChannel} that is mapped to the id.
     * @throws Exception  thrown if there is an error during processing.
     */
    @Nullable
    protected abstract QuicheQuicChannel quicPacketRead(ChannelHandlerContext ctx, InetSocketAddress sender,
                                                        InetSocketAddress recipient, QuicPacketType type, long version,
                                                        ByteBuf scid, ByteBuf dcid, ByteBuf token,
                                                        ByteBuf senderSockaddrMemory, ByteBuf recipientSockaddrMemory,
                                                        Consumer<QuicheQuicChannel> freeTask,
                                                        int localConnIdLength, QuicheConfig config) throws Exception;

    @Override
    public final void channelReadComplete(ChannelHandlerContext ctx) {
        inChannelReadComplete = true;
        try {
            for (;;) {
                QuicheQuicChannel channel = needsFireChannelReadComplete.poll();
                if (channel == null) {
                    break;
                }
                channel.recvComplete();
            }
        } finally {
            inChannelReadComplete = false;
            if (pendingPackets > 0) {
                flushNow(ctx);
            }
        }
    }

    @Override
    public final void channelWritabilityChanged(ChannelHandlerContext ctx) {
        if (ctx.channel().isWritable()) {
            // Ensure we delay removal from the channels Set as otherwise we will might see an exception
            // due modifications while iteration.
            delayRemoval = true;
            try {
                for (QuicheQuicChannel channel : channels) {
                    // TODO: Be a bit smarter about this.
                    channel.writable();
                }
            } finally {
                // We are done with the loop, reset the flag and process the removals from the channels Set.
                delayRemoval = false;
                processDelayedRemoval();
            }
        } else {
            // As we batch flushes we need to ensure we at least try to flush a batch once the channel becomes
            // unwritable. Otherwise we may end up with buffering too much writes and so waste memory.
            ctx.flush();
        }

        ctx.fireChannelWritabilityChanged();
    }

    @Override
    public final void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)  {
        pendingPackets ++;
        int size = estimatorHandle.size(msg);
        if (size > 0) {
            pendingBytes += size;
        }
        try {
            ctx.write(msg, promise);
        } finally {
            flushIfNeeded(ctx);
        }
    }

    @Override
    public final void flush(ChannelHandlerContext ctx) {
        // If we are in the channelReadComplete(...) method we might be able to delay the flush(...) until we finish
        // processing all channels.
        if (inChannelReadComplete) {
            flushIfNeeded(ctx);
        } else if (pendingPackets > 0) {
            flushNow(ctx);
        }
    }

    @Override
    public void connect(ChannelHandlerContext ctx, SocketAddress remoteAddress, SocketAddress localAddress,
                        ChannelPromise promise) throws Exception {
        if (remoteAddress instanceof QuicheQuicChannelAddress) {
            QuicheQuicChannelAddress addr = (QuicheQuicChannelAddress) remoteAddress;
            QuicheQuicChannel channel = addr.channel;
            connectQuicChannel(channel, remoteAddress, localAddress,
                    senderSockaddrMemory, recipientSockaddrMemory, freeTask, localConnIdLength, config, promise);
        } else {
            ctx.connect(remoteAddress, localAddress, promise);
        }
    }

    /**
     * Connects the given {@link QuicheQuicChannel}.
     *
     * @param channel                   the {@link QuicheQuicChannel} to connect.
     * @param remoteAddress             the remote {@link SocketAddress}.
     * @param localAddress              the local  {@link SocketAddress}
     * @param senderSockaddrMemory      the {@link ByteBuf} that can be used for the sender {@code struct sockaddr).
     * @param recipientSockaddrMemory   the {@link ByteBuf} that can be used for the recipient {@code struct sockaddr).
     * @param freeTask                  the {@link Consumer} that will be called once native memory of the
     *                                  {@link QuicheQuicChannel} is freed and so the mappings should be deleted to
     *                                  the ids.
     * @param localConnIdLength         the length of the local connection ids.
     * @param config                    the {@link QuicheConfig} that is used.
     * @param promise                   the {@link ChannelPromise} to notify once the connect is done.
     */
    protected abstract void connectQuicChannel(QuicheQuicChannel channel, SocketAddress remoteAddress,
                                               SocketAddress localAddress, ByteBuf senderSockaddrMemory,
                                               ByteBuf recipientSockaddrMemory, Consumer<QuicheQuicChannel> freeTask,
                                               int localConnIdLength, QuicheConfig config, ChannelPromise promise);


    private void flushIfNeeded(ChannelHandlerContext ctx) {
        // Check if we should force a flush() and so ensure the packets are delivered in a timely
        // manner and also make room in the outboundbuffer again that belongs to the underlying channel.
        if (flushStrategy.shouldFlushNow(pendingPackets, pendingBytes)) {
            flushNow(ctx);
        }
    }

    private void flushNow(ChannelHandlerContext ctx) {
        pendingBytes = 0;
        pendingPackets = 0;
        ctx.flush();
    }

    private final class QuicCodecHeaderProcessor implements QuicHeaderParser.QuicHeaderProcessor {

        private final ChannelHandlerContext ctx;

        QuicCodecHeaderProcessor(ChannelHandlerContext ctx) {
            this.ctx = ctx;
        }

        @Override
        public void process(InetSocketAddress sender, InetSocketAddress recipient, ByteBuf buffer, QuicPacketType type,
                            long version, ByteBuf scid, ByteBuf dcid, ByteBuf token) throws Exception {
            QuicheQuicChannel channel = quicPacketRead(ctx, sender, recipient,
                    type, version, scid,
                    dcid, token, senderSockaddrMemory, recipientSockaddrMemory, freeTask, localConnIdLength, config);
            if (channel != null) {
                // Add to queue first, we might be able to safe some flushes and consolidate them
                // in channelReadComplete(...) this way.
                if (channel.markInFireChannelReadCompleteQueue()) {
                    needsFireChannelReadComplete.add(channel);
                }
                channel.recv(sender, recipient, buffer);
                for (ByteBuffer retiredSourceConnectionId : channel.retiredSourceConnectionId()) {
                    removeMapping(channel, retiredSourceConnectionId);
                }
                for (ByteBuffer newSourceConnectionId : channel.newSourceConnectionIds()) {
                    addMapping(channel, newSourceConnectionId);
                }
            }
        }
    }
}
