/*
 * Copyright (C) 2025-2026  Tim Zimmermann <tim@aglais.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.aglais.zms.server;

import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.channel.epoll.*;
import io.netty.channel.group.ChannelGroup;
import io.netty.channel.group.DefaultChannelGroup;
import io.netty.channel.kqueue.*;
import io.netty.channel.nio.NioIoHandler;
import io.netty.channel.socket.nio.NioDatagramChannel;
import io.netty.channel.socket.nio.NioServerDomainSocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.unix.DomainSocketAddress;
import io.netty.handler.codec.http.HttpContentCompressor;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.codec.http2.Http2FrameCodecBuilder;
import io.netty.handler.codec.http2.Http2MultiplexHandler;
import io.netty.handler.codec.http2.Http2SecurityUtil;
import io.netty.handler.codec.http2.Http2StreamFrameToHttpObjectCodec;
import io.netty.handler.codec.http3.Http3;
import io.netty.handler.codec.http3.Http3FrameToHttpObjectCodec;
import io.netty.handler.codec.http3.Http3ServerConnectionHandler;
import io.netty.handler.codec.quic.QuicChannel;
import io.netty.handler.codec.quic.QuicSslContext;
import io.netty.handler.codec.quic.QuicSslContextBuilder;
import io.netty.handler.codec.quic.QuicStreamChannel;
import io.netty.handler.ssl.*;
import io.netty.util.concurrent.GlobalEventExecutor;
import joptsimple.OptionParser;
import joptsimple.OptionSet;
import org.aglais.zms.db.core.Database;
import org.aglais.zms.server.codec.quic.SecureQuicTokenHandler;
import org.aglais.zms.server.routing.AglaisHttpRouteHandler;
import org.aglais.zms.util.ConfigUtil;
import org.aglais.zms.util.Log;
import sun.misc.Signal;

import javax.net.ssl.SSLException;
import java.io.File;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.UnixDomainSocketAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import static org.aglais.zms.api.core.APIConstants.PROJECT_NAME;

/**
 * Main server class
 */
public class AglaisZMS {

    private static final String LOG_TAG = AglaisZMS.class.getSimpleName();

    /**
     * Create a {@link SslContext} from the given listener configuration
     *
     * @param listener The listener to load a certificate for
     */
    private static SslContext loadSslCertificate(ServerConfig.Listener listener) {
        SslProvider provider = SslProvider.isAlpnSupported(SslProvider.OPENSSL) ? SslProvider.OPENSSL : SslProvider.JDK;

        try {
            return SslContextBuilder.forServer(new File(listener.ssl().certificate()),
                            new File(listener.ssl().certificateKey()))
                    .sslProvider(provider)
                    .ciphers(Http2SecurityUtil.CIPHERS, SupportedCipherSuiteFilter.INSTANCE)
                    .applicationProtocolConfig(new ApplicationProtocolConfig(
                            ApplicationProtocolConfig.Protocol.ALPN,
                            ApplicationProtocolConfig.SelectorFailureBehavior.NO_ADVERTISE,
                            ApplicationProtocolConfig.SelectedListenerFailureBehavior.ACCEPT,
                            ApplicationProtocolNames.HTTP_2,
                            ApplicationProtocolNames.HTTP_1_1
                    ))
                    .build();
        } catch (SSLException e) {
            Log.e(LOG_TAG, "Failed to load SSL certificate", e);
        }

        return null;
    }

    /**
     * Create a {@link QuicSslContext} from the given listener configuration
     *
     * @param listener The listener to load a certificate for
     */
    private static QuicSslContext loadQuicSslCertificate(ServerConfig.Listener listener) {
        try {
            return QuicSslContextBuilder.forServer(new File(listener.ssl().certificateKey()), null,
                            new File(listener.ssl().certificate()))
                    .applicationProtocols(Http3.supportedApplicationProtocols())
                    .build();
        } catch (Exception e) {
            Log.e(LOG_TAG, "Failed to load SSL certificate", e);
        }

        return null;
    }

    /**
     * Creates a handler for HTTP/3 connections
     *
     * @param sslCtx      The {@link QuicSslContext} for handling SSL
     * @param db          The {@link Database} instance to use for authentication
     * @param config      The {@link ServerConfig} to get database credentials from
     * @param allChannels {@link ChannelGroup} containing all channels with user attributes
     */
    private static ChannelHandler createHttp3Handler(QuicSslContext sslCtx, Database db, ServerConfig config,
                                                     ChannelGroup allChannels) {
        AglaisHttpAuthenticator authenticator = new AglaisHttpAuthenticator(true, db, config.trustedProxies);
        AglaisHttpRouteHandler handler = new AglaisHttpRouteHandler(true, Database.newInstance(config.database.type(),
                config.database.host(), config.database.name(), config.database.username(), config.database.password()),
                allChannels);

        return Http3.newQuicServerCodecBuilder()
                .sslContext(sslCtx)
                .maxIdleTimeout(60, TimeUnit.SECONDS)
                .initialMaxData(Integer.MAX_VALUE)
                .initialMaxStreamDataBidirectionalLocal(Integer.MAX_VALUE)
                .initialMaxStreamDataBidirectionalRemote(Integer.MAX_VALUE)
                .initialMaxStreamsBidirectional(100)
                .tokenHandler(new SecureQuicTokenHandler())
                .handler(new ChannelInitializer<QuicChannel>() {
                    @Override
                    protected void initChannel(QuicChannel ch) {
                        ch.pipeline().addLast(new Http3ServerConnectionHandler(
                                new ChannelInitializer<QuicStreamChannel>() {
                                    @Override
                                    protected void initChannel(QuicStreamChannel ch) {
                                        ChannelPipeline p = ch.pipeline();
                                        p.addLast(new Http3FrameToHttpObjectCodec(true));
                                        p.addLast(new HttpContentCompressor());
                                        p.addLast(authenticator);
                                        p.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
                                        p.addLast(handler);
                                    }
                                }));
                    }
                })
                .build();
    }

    /**
     * Register handlers for HTTP/2 connections
     *
     * @param p             The {@link ChannelPipeline} to register handlers on
     * @param db            The {@link Database} instance to use for authentication
     * @param config        The {@link ServerConfig} to get database credentials from
     * @param allChannels   {@link ChannelGroup} containing all channels with user attributes
     * @param advertiseQuic True if support for QUIC should be advertised using the Alt-Svc header
     */
    private static void registerHttp2Handlers(ChannelPipeline p, Database db, ServerConfig config,
                                              ChannelGroup allChannels, boolean advertiseQuic) {
        AglaisHttpAuthenticator authenticator = new AglaisHttpAuthenticator(advertiseQuic, db, config.trustedProxies);
        AglaisHttpRouteHandler handler = new AglaisHttpRouteHandler(advertiseQuic,
                Database.newInstance(config.database.type(), config.database.host(), config.database.name(),
                        config.database.username(), config.database.password()), allChannels);

        p.addLast(Http2FrameCodecBuilder.forServer().build());
        p.addLast(new Http2MultiplexHandler(new ChannelInitializer<>() {
            @Override
            protected void initChannel(Channel ch) {
                ChannelPipeline p = ch.pipeline();
                p.addLast(new Http2StreamFrameToHttpObjectCodec(true));
                p.addLast(new HttpContentCompressor());
                p.addLast(authenticator);
                p.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
                p.addLast(handler);
            }
        }));
    }

    /**
     * Register handlers for HTTP/1.x connections
     *
     * @param p             The {@link ChannelPipeline} to register handlers on
     * @param db            The {@link Database} instance to use for authentication
     * @param config        The {@link ServerConfig} to get database credentials from
     * @param allChannels   {@link ChannelGroup} containing all channels with user attributes
     * @param advertiseQuic True if support for QUIC should be advertised using the Alt-Svc header
     */
    private static void registerHttp1Handlers(ChannelPipeline p, Database db, ServerConfig config,
                                              ChannelGroup allChannels, boolean advertiseQuic) {
        p.addLast(new HttpServerCodec());
        p.addLast(new HttpContentCompressor());
        p.addLast(new AglaisHttpAuthenticator(advertiseQuic, db, config.trustedProxies));
        p.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
        p.addLast(new AglaisHttpRouteHandler(advertiseQuic, Database.newInstance(config.database.type(),
                config.database.host(), config.database.name(), config.database.username(), config.database.password()),
                allChannels));
    }

    /**
     * Register a negotiation handler for negotiating HTTP version.
     * {@link #registerHttp2Handlers} will be called for HTTP/2 connections and {@link #registerHttp1Handlers} for
     * HTTP/1.1 connections.
     *
     * @param p             The {@link ChannelPipeline} to register handlers on
     * @param db            The {@link Database} instance to use for authentication
     * @param config        The {@link ServerConfig} to get database credentials from
     * @param allChannels   {@link ChannelGroup} containing all channels with user attributes
     * @param advertiseQuic True if support for QUIC should be advertised using the Alt-Svc header
     */
    private static void registerNegotiationHandler(ChannelPipeline p, Database db, ServerConfig config,
                                                   ChannelGroup allChannels, boolean advertiseQuic) {
        p.addLast(new ApplicationProtocolNegotiationHandler(ApplicationProtocolNames.HTTP_1_1) {
            @Override
            protected void configurePipeline(ChannelHandlerContext ctx, String protocol) {
                if (ApplicationProtocolNames.HTTP_2.equals(protocol)) {
                    registerHttp2Handlers(ctx.pipeline(), db, config, allChannels, advertiseQuic);
                } else if (ApplicationProtocolNames.HTTP_1_1.equals(protocol)) {
                    registerHttp1Handlers(ctx.pipeline(), db, config, allChannels, advertiseQuic);
                } else {
                    throw new IllegalStateException("Unknown protocol: " + protocol);
                }
            }
        });
    }

    /**
     * Application entry point
     *
     * @param args Command line arguments
     */
    public static void main(String[] args) throws Exception {
        OptionParser parser = new OptionParser();

        parser.accepts("config")
                .withRequiredArg()
                .ofType(String.class)
                .describedAs("Config file location")
                .defaultsTo("server.json");
        parser.accepts("help")
                .forHelp();

        OptionSet options = parser.parse(args);

        // Print only help if requested
        if (options.has("help")) {
            parser.printHelpOn(System.out);
            return;
        }

        ServerConfig config = ConfigUtil.loadConfig((String) options.valueOf("config"), ServerConfig.class,
                ServerConfig::new);
        if (config.database == null) {
            Log.e(LOG_TAG, "Database settings missing from config");
            return;
        }

        // Check connectivity
        Database db = Database.newInstance(config.database.type(), config.database.host(), config.database.name(),
                config.database.username(), config.database.password());
        if (!db.connect()) {
            Log.e(LOG_TAG, "Database connectivity check failed");
            System.exit(1);
        }

        int dbVersion = db.getVersion();
        if (dbVersion < db.getRequiredVersion()) {
            if (dbVersion > 0)
                Log.i(LOG_TAG, "Database must be upgraded from v" + db.getVersion() + " to v"
                        + db.getRequiredVersion());
            else
                Log.i(LOG_TAG, PROJECT_NAME + " is not yet installed in this database");

            boolean upgrade = config.database.autoupgrade();
            if (!upgrade) {
                Log.i(LOG_TAG, (dbVersion > 0 ? "Upgrade" : "Install") + " now? (Y/n)");

                try (Scanner scanner = new Scanner(System.in)) {
                    String line = scanner.nextLine();
                    upgrade = line.isBlank() || line.equalsIgnoreCase("y");
                }
            }

            if (upgrade) {
                Log.i(LOG_TAG, (dbVersion > 0 ? "Upgrading" : "Installing") + " database...");

                for (int i = dbVersion + 1; i <= db.getRequiredVersion(); i++) {
                    if (!db.upgrade(i)) {
                        Log.e(LOG_TAG, "Failed to " + (dbVersion > 0 ? "upgrade" : "install") + " database");
                        System.exit(1);
                    }
                }
            } else {
                Log.e(LOG_TAG, "Aborting!");
                System.exit(1);
            }
        } else if (db.getVersion() > db.getRequiredVersion()) {
            Log.e(LOG_TAG, "A newer version of zms-server is required to use this database!");
            Log.e(LOG_TAG, "zms-server requires database version: " + db.getRequiredVersion());
            Log.e(LOG_TAG, "Actually installed database version: " + db.getVersion());
        }

        ConcurrentHashMap<ServerConfig.Listener, SslContext> sslContexts = new ConcurrentHashMap<>();
        ConcurrentHashMap<ServerConfig.Listener, QuicSslContext> quicSslContexts = new ConcurrentHashMap<>();
        Signal.handle(new Signal("HUP"), signal -> {
            Log.i(LOG_TAG, "Reloading SSL configuration");

            for (ServerConfig.Listener listener : sslContexts.keySet()) {
                SslContext sslCtx = loadSslCertificate(listener);
                if (sslCtx != null)
                    sslContexts.replace(listener, sslCtx);
            }
            for (ServerConfig.Listener listener : quicSslContexts.keySet()) {
                QuicSslContext quicSslCtx = loadQuicSslCertificate(listener);
                if (quicSslCtx != null)
                    quicSslContexts.replace(listener, quicSslCtx);
            }
        });

        boolean epoll = Epoll.isAvailable();
        boolean kqueue = KQueue.isAvailable();
        EventLoopGroup group = new MultiThreadIoEventLoopGroup(epoll ? EpollIoHandler.newFactory()
                : kqueue ? KQueueIoHandler.newFactory() : NioIoHandler.newFactory());

        if (epoll)
            Log.i(LOG_TAG, "Using epoll");
        else if (kqueue)
            Log.i(LOG_TAG, "Using kqueue");

        try {
            List<ChannelFuture> closeFutures = new ArrayList<>();
            for (ServerConfig.Listener listener : config.listeners) {
                Class<? extends ServerChannel> channelClass;
                Class<? extends AbstractChannel> datagramChannelClass = null;
                SocketAddress sockAddr;
                if (listener.type() == ServerConfig.ListenerType.TCP) {
                    if (listener.host() == null || listener.host().isEmpty() || listener.port() == 0)
                        throw new IllegalArgumentException("TCP listener must have host and port configured");

                    channelClass = epoll ? EpollServerSocketChannel.class : kqueue ? KQueueServerSocketChannel.class
                            : NioServerSocketChannel.class;
                    datagramChannelClass = epoll ? EpollDatagramChannel.class : kqueue ? KQueueDatagramChannel.class
                            : NioDatagramChannel.class;
                    sockAddr = new InetSocketAddress(listener.host(), listener.port());
                } else if (listener.type() == ServerConfig.ListenerType.UNIX) {
                    if (listener.path() == null || listener.path().isEmpty())
                        throw new IllegalArgumentException("Unix listener must have a path configured");

                    channelClass = epoll ? EpollServerDomainSocketChannel.class
                            : kqueue ? KQueueServerDomainSocketChannel.class : NioServerDomainSocketChannel.class;
                    sockAddr = epoll || kqueue ? new DomainSocketAddress(listener.path())
                            : UnixDomainSocketAddress.of(listener.path());
                } else {
                    throw new IllegalArgumentException("Unknown listener type: " + listener.type());
                }

                Log.i(LOG_TAG, "Listening on " + sockAddr);

                if (listener.ssl() != null && listener.ssl().enabled()) {
                    SslContext sslCtx = loadSslCertificate(listener);
                    if (sslCtx != null) {
                        sslContexts.put(listener, sslCtx);

                        Log.i(LOG_TAG, "Using SSL");
                    }

                    if (datagramChannelClass != null) {
                        QuicSslContext quicSslCtx = loadQuicSslCertificate(listener);
                        if (quicSslCtx != null) {
                            quicSslContexts.put(listener, quicSslCtx);

                            Log.i(LOG_TAG, "Using QUIC");
                        }
                    } else {
                        Log.w(LOG_TAG, "Skipping QUIC due to unsupported listener type");
                    }
                }

                ChannelGroup allChannels = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE);
                ServerBootstrap bootstrap = new ServerBootstrap();
                bootstrap.group(group)
                        .channel(channelClass)
                        .childHandler(new ChannelInitializer<>() {

                            @Override
                            protected void initChannel(Channel ch) {
                                ChannelPipeline p = ch.pipeline();
                                SslContext sslCtx = sslContexts.get(listener);
                                if (sslCtx != null) {
                                    p.addLast(sslCtx.newHandler(ch.alloc()));
                                    registerNegotiationHandler(p, db, config, allChannels,
                                            quicSslContexts.containsKey(listener));
                                } else {
                                    registerHttp1Handlers(p, db, config, allChannels,
                                            quicSslContexts.containsKey(listener));
                                }
                            }
                        });

                ChannelFuture f = bootstrap.bind(sockAddr).sync();
                closeFutures.add(f.channel().closeFuture());

                QuicSslContext quicSslCtx = quicSslContexts.get(listener);
                if (quicSslCtx != null && datagramChannelClass != null) {
                    Bootstrap quicBootstrap = new Bootstrap();
                    quicBootstrap.group(group)
                            .channel(datagramChannelClass)
                            .handler(createHttp3Handler(quicSslCtx, db, config, allChannels));

                    f = quicBootstrap.bind(sockAddr).sync();
                    closeFutures.add(f.channel().closeFuture());
                }
            }

            for (ChannelFuture f : closeFutures) {
                f.sync();
            }
        } finally {
            group.shutdownGracefully();
        }
    }

}
