/*
 * 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.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.unix.DomainSocketAddress;
import io.netty.handler.codec.http.DefaultFullHttpResponse;
import io.netty.handler.codec.http.FullHttpResponse;
import io.netty.handler.codec.http.HttpHeaderNames;
import io.netty.handler.codec.http.HttpRequest;
import io.netty.util.AttributeKey;
import io.netty.util.ReferenceCountUtil;
import org.aglais.zms.db.core.Database;
import org.aglais.zms.db.core.model.User;
import org.aglais.zms.util.Log;

import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.UnixDomainSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.List;

import static io.netty.handler.codec.http.HttpResponseStatus.UNAUTHORIZED;
import static org.aglais.zms.api.core.APIConstants.X_AGLAIS_ZMS_API_VERSION;

/**
 * Class providing support for HTTP basic authentication
 */
@ChannelHandler.Sharable
public class AglaisHttpAuthenticator extends AglaisBaseHttpRequestHandler {

    /**
     * {@link AttributeKey} to retrieve current authentication retries
     */
    private static final AttributeKey<Integer> AUTH_RETRIES_KEY = AttributeKey.valueOf("auth_retries");
    /**
     * {@link AttributeKey} to retrieve current {@link User} for the active connection
     */
    public static final AttributeKey<User> USER_KEY = AttributeKey.valueOf("user");
    /**
     * {@link AttributeKey} to retrieve the real IP of the currently active connection
     * <p>
     * This will either be the channel's local address or the IP passed by a trusted proxy using the X-Real-IP header.
     */
    public static final AttributeKey<String> IP_KEY = AttributeKey.valueOf("ip");
    /**
     * {@link AttributeKey} to retrieve the API version for the active connection
     */
    public static final AttributeKey<Integer> API_VERSION_KEY = AttributeKey.valueOf("api_version");

    private static final String LOG_TAG = AglaisHttpAuthenticator.class.getSimpleName();
    private static final String BASIC_AUTH_PREFIX = "Basic ";
    private static final String X_REAL_IP = "X-Real-IP";
    private static final int MAX_RETRIES = 3;

    private final Database db;
    private final List<String> trustedProxies;

    /**
     * Create a new HTTP basic authenticator using the specified {@link Database} for querying users
     *
     * @param advertiseQuic  True if support for QUIC should be advertised using the Alt-Svc header
     * @param db             The {@link Database} to load users from
     * @param trustedProxies A list with IP addresses to trust for setting the X-Real-IP header
     */
    public AglaisHttpAuthenticator(boolean advertiseQuic, Database db, List<String> trustedProxies) {
        super(advertiseQuic);
        this.db = db;
        this.trustedProxies = trustedProxies;
    }

    /**
     * Copy an attribute from one channel to another
     *
     * @param src The source channel
     * @param dst The destination channel
     * @param key The attribute to copy
     * @param <T> The attribute's value type
     */
    private <T> void copyAttribute(Channel src, Channel dst, AttributeKey<T> key) {
        if (src.hasAttr(key)) {
            dst.attr(key).set(src.attr(key).get());
        }
    }

    /**
     * Copy all attributes from one channel to another
     *
     * @param src The source channel
     * @param dst The destination channel
     */
    private void copyAttributes(Channel src, Channel dst) {
        copyAttribute(src, dst, USER_KEY);
        copyAttribute(src, dst, IP_KEY);
        copyAttribute(src, dst, API_VERSION_KEY);
        copyAttribute(src, dst, AUTH_RETRIES_KEY);
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, HttpRequest msg) {
        boolean isChildConnectionChannel = ChannelUtil.isChildConnectionChannel(ctx.channel());
        if (isChildConnectionChannel)
            copyAttributes(ctx.channel().parent(), ctx.channel());

        // Get client IP
        String ip = ctx.channel().attr(IP_KEY).get();
        if (ip == null) {
            SocketAddress addr = ChannelUtil.getRemoteAddress(ctx.channel());
            ip = addr.toString();
            if (addr instanceof InetSocketAddress inetAddr)
                ip = inetAddr.getAddress().getHostAddress().split("%")[0];
            else if (addr instanceof UnixDomainSocketAddress || addr instanceof DomainSocketAddress)
                ip = "unix";
            if (trustedProxies.contains(ip))
                if (msg.headers().contains(X_REAL_IP))
                    ip = msg.headers().get(X_REAL_IP);
            ctx.channel().attr(IP_KEY).set(ip);
        }

        // Get API version
        if (ctx.channel().attr(API_VERSION_KEY).get() == null) {
            String apiVersion = msg.headers().get(X_AGLAIS_ZMS_API_VERSION, "0");
            try {
                ctx.channel().attr(API_VERSION_KEY).set(Integer.parseInt(apiVersion));
            } catch (NumberFormatException e) {
                Log.e(LOG_TAG, "Invalid API version " + apiVersion + " from " + ip);
                ctx.channel().attr(API_VERSION_KEY).set(0);
            }
        }

        // Check if a username and password are provided
        if (ctx.channel().attr(USER_KEY).get() == null) {
            String authorization = msg.headers().get(HttpHeaderNames.AUTHORIZATION);
            String username = null;
            String password;
            boolean validUser = false;
            if (authorization != null && authorization.startsWith(BASIC_AUTH_PREFIX)) {
                authorization = authorization.substring(BASIC_AUTH_PREFIX.length());
                authorization = new String(Base64.getDecoder().decode(authorization), StandardCharsets.UTF_8);

                if (authorization.contains(":")) {
                    username = authorization.substring(0, authorization.indexOf(":"));
                    password = authorization.substring(authorization.indexOf(":") + 1);

                    User user = db.getUserByName(username);
                    validUser = user != null;
                    if (validUser) {
                        if (user.verifyPassword(password)) {
                            ctx.channel().attr(USER_KEY).set(user);
                        }
                    } else {
                        Log.i(LOG_TAG, "Invalid user " + username + " from " + ip);
                    }
                }

                // Disconnect after too many retries
                if (ctx.channel().attr(USER_KEY).get() == null) {
                    Log.i(LOG_TAG, "Failed password for " + (validUser ? "" : "invalid user ") + username
                            + " from " + ip);

                    Integer authRetries = ctx.channel().attr(AUTH_RETRIES_KEY).get();
                    if (authRetries == null)
                        authRetries = 0;
                    authRetries++;
                    if (authRetries >= MAX_RETRIES) {
                        if (isChildConnectionChannel)
                            ctx.channel().parent().close();
                        else
                            ctx.close();

                        Log.i(LOG_TAG, "Disconnected " + (validUser ? "" : "invalid user ") + username + " " + ip);
                    }
                    ctx.channel().attr(AUTH_RETRIES_KEY).set(authRetries);
                    if (isChildConnectionChannel)
                        copyAttributes(ctx.channel(), ctx.channel().parent());
                } else {
                    Log.i(LOG_TAG, "Accepted password for " + username + " from " + ip);
                    if (isChildConnectionChannel)
                        copyAttributes(ctx.channel(), ctx.channel().parent());
                }
            }
        }
        // Still not authenticated? Request authorization.
        if (ctx.channel().attr(USER_KEY).get() == null) {
            FullHttpResponse response = new DefaultFullHttpResponse(msg.protocolVersion(), UNAUTHORIZED);

            response.headers().set(HttpHeaderNames.WWW_AUTHENTICATE, BASIC_AUTH_PREFIX
                    + "realm=\"Authorization required\", charset=\"UTF-8\"");
            writeResponse(ctx, msg, response);
            return;
        }

        ctx.channel().pipeline().remove(this);
        ctx.fireChannelRead(ReferenceCountUtil.retain(msg));
    }
}
