// Copyright 2012-2025 The NATS Authors
// Licensed 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.
//
// Adapted from server/auth.go in the NATS server Go source.

using System.Security.Cryptography.X509Certificates;
using ZB.MOM.NatsNet.Server.Auth;
using ZB.MOM.NatsNet.Server.Internal;

namespace ZB.MOM.NatsNet.Server;

/// <summary>
/// Authentication logic for <see cref="NatsServer"/>.
/// Mirrors Go auth.go Server methods.
/// </summary>
public sealed partial class NatsServer
{
    /// <summary>
    /// Wires up auth lookup tables from options.
    /// Mirrors Go <c>configureAuthorization</c>.
    /// </summary>
    internal void ConfigureAuthorization()
    {
        var opts = GetOpts();

        if (opts.CustomClientAuthentication != null)
        {
            _info.AuthRequired = true;
        }
        else if (_trustedKeys != null)
        {
            _info.AuthRequired = true;
        }
        else if (opts.Nkeys != null || opts.Users != null)
        {
            (_nkeys, _users) = BuildNkeysAndUsersFromOptions(opts.Nkeys, opts.Users);
            _info.AuthRequired = true;
        }
        else if (!string.IsNullOrEmpty(opts.Username) || !string.IsNullOrEmpty(opts.Authorization))
        {
            _info.AuthRequired = true;
        }
        else
        {
            _users = null;
            _nkeys = null;
            _info.AuthRequired = false;
        }

        if (opts.AuthCallout != null && string.IsNullOrEmpty(opts.AuthCallout.Account))
            Errorf("Authorization callout account not set");
    }

    private (Dictionary<string, NkeyUser>? nkeys, Dictionary<string, User>? users) BuildNkeysAndUsersFromOptions(
        List<NkeyUser>? nko, List<User>? uo)
    {
        Dictionary<string, NkeyUser>? nkeys = null;
        Dictionary<string, User>? users = null;

        if (nko != null)
        {
            nkeys = new Dictionary<string, NkeyUser>(nko.Count, StringComparer.Ordinal);
            foreach (var u in nko)
            {
                if (u.Permissions != null)
                    AuthHandler.ValidateResponsePermissions(u.Permissions);
                nkeys[u.Nkey] = u;
            }
        }

        if (uo != null)
        {
            users = new Dictionary<string, User>(uo.Count, StringComparer.Ordinal);
            foreach (var u in uo)
            {
                if (u.Permissions != null)
                    AuthHandler.ValidateResponsePermissions(u.Permissions);
                users[u.Username] = u;
            }
        }

        AssignGlobalAccountToOrphanUsers(nkeys, users);
        return (nkeys, users);
    }

    internal void AssignGlobalAccountToOrphanUsers(
        Dictionary<string, NkeyUser>? nkeys,
        Dictionary<string, User>? users)
    {
        if (nkeys != null)
            foreach (var u in nkeys.Values)
                u.Account ??= _gacc;

        if (users != null)
            foreach (var u in users.Values)
                u.Account ??= _gacc;
    }

    /// <summary>
    /// Entry-point auth check — dispatches by client kind.
    /// Mirrors Go <c>checkAuthentication</c>.
    /// </summary>
    internal bool CheckAuthentication(ClientConnection c)
    {
        return c.Kind switch
        {
            ClientKind.Client  => IsClientAuthorized(c),
            ClientKind.Router  => IsRouterAuthorized(c),
            ClientKind.Gateway => IsGatewayAuthorized(c),
            ClientKind.Leaf    => IsLeafNodeAuthorized(c),
            _                  => false,
        };
    }

    /// <summary>Mirrors Go <c>isClientAuthorized</c>.</summary>
    internal bool IsClientAuthorized(ClientConnection c)
        => ProcessClientOrLeafAuthentication(c, GetOpts());

    /// <summary>
    /// Full authentication dispatch — handles all auth paths.
    /// Mirrors Go <c>processClientOrLeafAuthentication</c>.
    /// </summary>
    internal bool ProcessClientOrLeafAuthentication(ClientConnection c, ServerOptions opts)
    {
        // Auth callout check
        if (opts.AuthCallout != null)
            return ProcessClientOrLeafCallout(c, opts);

        // Proxy check
        var (trustedProxy, proxyOk) = ProxyCheck(c, opts);
        if (trustedProxy && !proxyOk)
        {
            c.SetAuthError(new InvalidOperationException("proxy not trusted"));
            return false;
        }

        // Trusted operators / JWT bearer
        if (_trustedKeys != null)
        {
            var token = c.GetAuthToken();
            if (string.IsNullOrEmpty(token))
            {
                c.SetAuthError(new InvalidOperationException("missing JWT token for trusted operator"));
                return false;
            }
            // TODO: full JWT validation against trusted operators
            return true;
        }

        // NKey authentication
        if (_nkeys != null && _nkeys.Count > 0)
        {
            var nkeyPub = c.GetNkey();
            if (!string.IsNullOrEmpty(nkeyPub) && _nkeys.TryGetValue(nkeyPub, out var nkeyUser))
            {
                var sig   = c.GetNkeySig();
                var nonce = c.GetNonce();   // byte[]?
                if (!string.IsNullOrEmpty(sig) && nonce != null && nonce.Length > 0)
                {
                    try
                    {
                        var kp = NATS.NKeys.KeyPair.FromPublicKey(nkeyPub.AsSpan());
                        // Sig is raw URL-safe base64; convert to standard base64 with padding.
                        var padded = sig.Replace('-', '+').Replace('_', '/');
                        var rem = padded.Length % 4;
                        if (rem == 2) padded += "==";
                        else if (rem == 3) padded += "=";
                        var sigBytes = Convert.FromBase64String(padded);
                        var verified = kp.Verify(
                            new ReadOnlyMemory<byte>(nonce),
                            new ReadOnlyMemory<byte>(sigBytes));
                        if (!verified)
                        {
                            c.SetAuthError(new InvalidOperationException("NKey signature verification failed"));
                            return false;
                        }
                    }
                    catch (Exception ex)
                    {
                        c.SetAuthError(ex);
                        return false;
                    }
                }
                c.SetAccount(nkeyUser.Account);
                c.SetPermissions(nkeyUser.Permissions);
                return true;
            }
        }

        // Username / password
        if (_users != null && _users.Count > 0)
        {
            var username = c.GetUsername();
            if (_users.TryGetValue(username, out var user))
            {
                if (!AuthHandler.ComparePasswords(user.Password, c.GetPassword()))
                {
                    c.SetAuthError(new InvalidOperationException("invalid password"));
                    return false;
                }
                c.SetAccount(user.Account);
                c.SetPermissions(user.Permissions);
                return true;
            }
        }

        // Global username/password (from opts)
        if (!string.IsNullOrEmpty(opts.Username))
        {
            if (c.GetUsername() != opts.Username ||
                !AuthHandler.ComparePasswords(opts.Password, c.GetPassword()))
            {
                c.SetAuthError(new InvalidOperationException("invalid credentials"));
                return false;
            }
            return true;
        }

        // Token (authorization)
        if (!string.IsNullOrEmpty(opts.Authorization))
        {
            if (!AuthHandler.ComparePasswords(opts.Authorization, c.GetAuthToken()))
            {
                c.SetAuthError(new InvalidOperationException("bad authorization token"));
                return false;
            }
            return true;
        }

        // TLS cert mapping
        if (opts.TlsMap)
        {
            var cert = c.GetTlsCertificate();
            if (!AuthHandler.CheckClientTlsCertSubject(cert, _ => true))
            {
                c.SetAuthError(new InvalidOperationException("TLS cert mapping failed"));
                return false;
            }
            return true;
        }

        // No auth required
        if (!_info.AuthRequired) return true;

        c.SetAuthError(new InvalidOperationException("no credentials provided"));
        return false;
    }

    /// <summary>Mirrors Go <c>isRouterAuthorized</c>.</summary>
    internal bool IsRouterAuthorized(ClientConnection c)
    {
        var opts = GetOpts();
        if (opts.Cluster.Port == 0) return true;
        return true; // TODO: full route auth when ClusterOpts is fully typed
    }

    /// <summary>Mirrors Go <c>isGatewayAuthorized</c>.</summary>
    internal bool IsGatewayAuthorized(ClientConnection c)
    {
        var opts = GetOpts();
        if (string.IsNullOrEmpty(opts.Gateway.Name)) return true;
        return true;
    }

    /// <summary>Mirrors Go <c>registerLeafWithAccount</c>.</summary>
    internal bool RegisterLeafWithAccount(ClientConnection c, string accountName)
    {
        var (acc, _) = LookupAccount(accountName);
        if (acc == null) return false;
        c.SetAccount(acc);
        return true;
    }

    /// <summary>Mirrors Go <c>isLeafNodeAuthorized</c>.</summary>
    internal bool IsLeafNodeAuthorized(ClientConnection c)
        => ProcessClientOrLeafAuthentication(c, GetOpts());

    /// <summary>Mirrors Go <c>checkAuthforWarnings</c>.</summary>
    internal void CheckAuthforWarnings()
    {
        var opts = GetOpts();
        if (opts.Users != null && !string.IsNullOrEmpty(opts.Username))
            Warnf("Having a global password along with users/nkeys is not recommended");
    }

    /// <summary>Mirrors Go <c>proxyCheck</c>.</summary>
    internal (bool TrustedProxy, bool Ok) ProxyCheck(ClientConnection c, ServerOptions opts)
    {
        if (!opts.ProxyProtocol) return (false, false);
        // TODO: check remote IP against configured trusted proxy addresses
        return (true, true);
    }

    /// <summary>Mirrors Go <c>processProxiesTrustedKeys</c>.</summary>
    internal void ProcessProxiesTrustedKeys()
    {
        var opts = GetOpts();
        var keys = new HashSet<string>(StringComparer.Ordinal);

        if (opts.Proxies?.Trusted is { Count: > 0 })
        {
            foreach (var proxy in opts.Proxies.Trusted)
            {
                if (!string.IsNullOrWhiteSpace(proxy.Key))
                    keys.Add(proxy.Key.Trim());
            }
        }

        if (opts.TrustedKeys is { Count: > 0 })
        {
            foreach (var key in opts.TrustedKeys)
            {
                if (!string.IsNullOrWhiteSpace(key))
                    keys.Add(key.Trim());
            }
        }

        _proxiesKeyPairs.Clear();
        foreach (var key in keys)
            _proxiesKeyPairs.Add(key);
    }

    /// <summary>
    /// Forwards to AuthCallout.ProcessClientOrLeafCallout.
    /// Mirrors Go <c>processClientOrLeafCallout</c>.
    /// </summary>
    internal bool ProcessClientOrLeafCallout(ClientConnection c, ServerOptions opts)
        => AuthCallout.ProcessClientOrLeafCallout(this, c, opts);

    /// <summary>
    /// Config reload stub.
    /// Mirrors Go <c>Server.Reload</c>.
    /// </summary>
    internal void Reload()
    {
        _reloadMu.EnterWriteLock();
        try
        {
            _configTime = DateTime.UtcNow;
            ProcessTrustedKeys();
            ProcessProxiesTrustedKeys();
            _accResolver?.Reload();
        }
        finally
        {
            _reloadMu.ExitWriteLock();
        }
    }

    /// <summary>
    /// Returns a Task that shuts the server down asynchronously.
    /// Wraps the synchronous <see cref="Shutdown"/> method.
    /// </summary>
    internal Task ShutdownAsync() => Task.Run(Shutdown);
}