// Copyright 2012-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0

using System.Net.Security;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using ZB.MOM.NatsNet.Server.Auth.CertificateIdentityProvider;
using ZB.MOM.NatsNet.Server.Auth.Ocsp;

namespace ZB.MOM.NatsNet.Server;

/// <summary>
/// TLS configuration slot used by OCSP wiring to apply wrapped TLS settings.
/// Mirrors Go <c>tlsConfigKind</c> shape used by <c>configureOCSP</c>.
/// </summary>
internal sealed class OcspTlsConfig
{
    public required string Kind { get; init; }
    public required SslServerAuthenticationOptions TlsConfig { get; init; }
    public required TlsConfigOpts? TlsOptions { get; init; }
    public required Action<SslServerAuthenticationOptions> Apply { get; init; }
    public bool IsLeafSpoke { get; init; }
}

public sealed partial class NatsServer
{
    private const string ClientKindName = "client";
    private const string RouterKindName = "router";
    private const string GatewayKindName = "gateway";
    private const string LeafKindName = "leaf";
    private const string DefaultOcspStoreDir = "ocsp";

    internal OcspMonitor[] GetOcspMonitors()
    {
        _mu.EnterReadLock();
        try
        {
            return _ocsps is null ? [] : [.. _ocsps];
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal Exception? SetupOCSPStapleStoreDir()
    {
        var storeDir = GetOpts().StoreDir;
        if (string.IsNullOrEmpty(storeDir))
        {
            return null;
        }

        var ocspDir = Path.Combine(storeDir, DefaultOcspStoreDir);
        try
        {
            if (!Directory.Exists(ocspDir))
            {
                Directory.CreateDirectory(ocspDir);
            }
            else
            {
                var attributes = File.GetAttributes(ocspDir);
                if ((attributes & FileAttributes.Directory) != FileAttributes.Directory)
                {
                    return new InvalidOperationException("OCSP storage directory is not a directory");
                }
            }
        }
        catch (Exception ex)
        {
            return new InvalidOperationException($"could not create OCSP storage directory - {ex.Message}", ex);
        }

        return null;
    }

    internal List<OcspTlsConfig> ConfigureOCSP()
    {
        var opts = GetOpts();
        var configs = new List<OcspTlsConfig>();

        if (opts.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = ClientKindName,
                TlsConfig = opts.TlsConfig,
                TlsOptions = opts.TlsConfigOpts,
                Apply = tls => opts.TlsConfig = tls,
            });
        }

        if (opts.Websocket.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = ClientKindName,
                TlsConfig = opts.Websocket.TlsConfig,
                TlsOptions = opts.Websocket.TlsConfigOpts,
                Apply = tls => opts.Websocket.TlsConfig = tls,
            });
        }

        if (opts.Mqtt.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = ClientKindName,
                TlsConfig = opts.Mqtt.TlsConfig,
                TlsOptions = opts.Mqtt.TlsConfigOpts,
                Apply = tls => opts.Mqtt.TlsConfig = tls,
            });
        }

        if (opts.Cluster.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = RouterKindName,
                TlsConfig = opts.Cluster.TlsConfig,
                TlsOptions = opts.Cluster.TlsConfigOpts,
                Apply = tls => opts.Cluster.TlsConfig = tls,
            });
        }

        if (opts.LeafNode.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = LeafKindName,
                TlsConfig = opts.LeafNode.TlsConfig,
                TlsOptions = opts.LeafNode.TlsConfigOpts,
                Apply = tls => opts.LeafNode.TlsConfig = tls,
            });
        }

        foreach (var remote in opts.LeafNode.Remotes)
        {
            if (remote.TlsConfig == null)
            {
                continue;
            }

            var capturedRemote = remote;
            configs.Add(new OcspTlsConfig
            {
                Kind = LeafKindName,
                TlsConfig = remote.TlsConfig,
                TlsOptions = remote.TlsConfigOpts,
                IsLeafSpoke = true,
                Apply = tls => capturedRemote.TlsConfig = tls,
            });
        }

        if (opts.Gateway.TlsConfig != null)
        {
            configs.Add(new OcspTlsConfig
            {
                Kind = GatewayKindName,
                TlsConfig = opts.Gateway.TlsConfig,
                TlsOptions = opts.Gateway.TlsConfigOpts,
                Apply = tls => opts.Gateway.TlsConfig = tls,
            });
        }

        foreach (var gateway in opts.Gateway.Gateways)
        {
            if (gateway.TlsConfig == null)
            {
                continue;
            }

            var capturedGateway = gateway;
            configs.Add(new OcspTlsConfig
            {
                Kind = GatewayKindName,
                TlsConfig = gateway.TlsConfig,
                TlsOptions = gateway.TlsConfigOpts,
                Apply = tls => capturedGateway.TlsConfig = tls,
            });
        }

        return configs;
    }

    internal (SslServerAuthenticationOptions? tlsConfig, OcspMonitor? monitor, Exception? error) NewOCSPMonitor(
        OcspTlsConfig config)
    {
        var opts = GetOpts();
        var ocspConfig = opts.OcspConfig;

        var certFile = config.TlsOptions?.CertFile ?? opts.TlsCert;
        var caFile = config.TlsOptions?.CaFile ?? opts.TlsCaCert;

        if (config.TlsConfig.ServerCertificate is not X509Certificate2 leaf)
        {
            return (null, null, new InvalidOperationException("no certificate found"));
        }

        var shutdownOnRevoke = false;
        var mustStaple = OcspHandler.HasOCSPStatusRequest(leaf);
        if (ocspConfig != null)
        {
            switch (ocspConfig.Mode)
            {
                case OcspMode.Never:
                    if (mustStaple)
                    {
                        Warnf("Certificate at '{0}' has MustStaple but OCSP is disabled", certFile);
                    }
                    return (config.TlsConfig, null, null);
                case OcspMode.Always:
                    mustStaple = true;
                    shutdownOnRevoke = true;
                    break;
                case OcspMode.Must when mustStaple:
                    shutdownOnRevoke = true;
                    break;
                case OcspMode.Auto when !mustStaple:
                    return (config.TlsConfig, null, null);
            }
        }

        if (!mustStaple)
        {
            return (config.TlsConfig, null, null);
        }

        var setupError = SetupOCSPStapleStoreDir();
        if (setupError != null)
        {
            return (null, null, setupError);
        }

        var chain = new List<byte[]> { leaf.RawData };
        if (config.TlsConfig.ServerCertificateContext != null)
        {
            foreach (var intermediate in config.TlsConfig.ServerCertificateContext.IntermediateCertificates)
            {
                chain.Add(intermediate.RawData);
            }
        }

        var (issuer, issuerError) = OcspHandler.GetOCSPIssuer(caFile, chain);
        if (issuerError != null || issuer == null)
        {
            return (null, null, issuerError ?? new InvalidOperationException("no issuers found"));
        }

        var monitor = new OcspMonitor
        {
            Kind = config.Kind,
            Server = this,
            CertFile = certFile,
            CaFile = caFile,
            Leaf = leaf,
            Issuer = issuer,
            ShutdownOnRevoke = shutdownOnRevoke,
        };

        var (_, response, statusError) = monitor.GetStatus();
        if (statusError != null)
        {
            return (null, null,
                new InvalidOperationException($"bad OCSP status update for certificate at '{certFile}': {statusError.Message}", statusError));
        }

        if (response != null && response.Status != OcspStatusAssertion.Good && shutdownOnRevoke)
        {
            return (null, null,
                new InvalidOperationException(
                    $"found existing OCSP status for certificate at '{certFile}': {OcspHandler.OcspStatusString((int)response.Status)}"));
        }

        return (config.TlsConfig, monitor, null);
    }

    internal Exception? EnableOCSP()
    {
        var configs = ConfigureOCSP();
        var monitors = new List<OcspMonitor>(configs.Count);

        foreach (var config in configs)
        {
            if (config.Kind != LeafKindName)
            {
                var (tlsConfig, monitor, error) = NewOCSPMonitor(config);
                if (error != null)
                {
                    return error;
                }

                if (monitor != null && tlsConfig != null)
                {
                    monitors.Add(monitor);
                    config.Apply(tlsConfig);
                }
            }

            // OCSP peer verification hook is implemented in batch 9 F4.
        }

        _mu.EnterWriteLock();
        try
        {
            _ocsps = monitors.Count == 0 ? null : [.. monitors];
        }
        finally
        {
            _mu.ExitWriteLock();
        }

        return null;
    }

    internal void StartOCSPMonitoring()
    {
        OcspMonitor[]? monitors;
        _mu.EnterReadLock();
        try
        {
            monitors = _ocsps;
        }
        finally
        {
            _mu.ExitReadLock();
        }

        if (monitors == null || monitors.Length == 0)
        {
            return;
        }

        foreach (var monitor in monitors)
        {
            Noticef("OCSP Stapling enabled for {0} connections", monitor.Kind);
            monitor.Start();
            StartGoRoutine(() => monitor.Run(_quitCts.Token).GetAwaiter().GetResult());
        }
    }

    internal Exception? ReloadOCSP()
    {
        var setupError = SetupOCSPStapleStoreDir();
        if (setupError != null)
        {
            return setupError;
        }

        var existingMonitors = GetOcspMonitors();
        foreach (var monitor in existingMonitors)
        {
            monitor.Stop();
        }

        var configs = ConfigureOCSP();
        var replacement = new List<OcspMonitor>(configs.Count);

        _mu.EnterWriteLock();
        try
        {
            _ocspPeerVerify = false;
        }
        finally
        {
            _mu.ExitWriteLock();
        }

        foreach (var config in configs)
        {
            if (config.Kind != LeafKindName)
            {
                var (tlsConfig, monitor, error) = NewOCSPMonitor(config);
                if (error != null)
                {
                    return error;
                }

                if (monitor != null && tlsConfig != null)
                {
                    replacement.Add(monitor);
                    config.Apply(tlsConfig);
                }
            }
        }

        _mu.EnterWriteLock();
        try
        {
            _ocsps = replacement.Count == 0 ? null : [.. replacement];
        }
        finally
        {
            _mu.ExitWriteLock();
        }

        StartOCSPMonitoring();
        return null;
    }

    internal (SslServerAuthenticationOptions? tlsConfig, bool plugged, Exception? error) PlugTLSOCSPPeer(OcspTlsConfig? config)
    {
        if (config == null || config.TlsConfig == null)
        {
            return (null, false, new InvalidOperationException(OcspMessages.ErrUnableToPlugTLSEmptyConfig));
        }

        var kind = config.Kind;
        var isSpoke = config.IsLeafSpoke;
        var tlsOptions = config.TlsOptions;
        if (tlsOptions?.OcspPeerConfig == null || !tlsOptions.OcspPeerConfig.Verify)
        {
            return (config.TlsConfig, false, null);
        }

        Debugf(OcspMessages.DbgPlugTLSForKind, config.Kind);

        if (kind == ClientKindName || (kind == LeafKindName && !isSpoke))
        {
            if (!tlsOptions.Verify)
            {
                return (null, false, new InvalidOperationException(OcspMessages.ErrMTLSRequired));
            }

            return PlugClientTLSOCSPPeer(config);
        }

        if (kind == LeafKindName && isSpoke)
        {
            return PlugServerTLSOCSPPeer(config);
        }

        return (config.TlsConfig, false, null);
    }

    internal (SslServerAuthenticationOptions? tlsConfig, bool plugged, Exception? error) PlugClientTLSOCSPPeer(OcspTlsConfig? config)
    {
        if (config?.TlsConfig == null || config.TlsOptions == null)
        {
            return (null, false, new InvalidOperationException(OcspMessages.ErrUnableToPlugTLSClient));
        }

        var tlsConfig = config.TlsConfig;
        var tlsOptions = config.TlsOptions;
        if (tlsOptions.OcspPeerConfig == null || !tlsOptions.OcspPeerConfig.Verify)
        {
            return (tlsConfig, false, null);
        }

        tlsConfig.RemoteCertificateValidationCallback = (_, _, chain, _) =>
        {
            if (chain?.ChainElements == null || chain.ChainElements.Count == 0)
            {
                return false;
            }

            var converted = chain.ChainElements
                .Select(e => e.Certificate)
                .OfType<X509Certificate2>()
                .ToArray();
            return TlsClientOCSPValid([converted], tlsOptions.OcspPeerConfig);
        };

        return (tlsConfig, true, null);
    }

    internal (SslServerAuthenticationOptions? tlsConfig, bool plugged, Exception? error) PlugServerTLSOCSPPeer(OcspTlsConfig? config)
    {
        if (config?.TlsConfig == null || config.TlsOptions == null)
        {
            return (null, false, new InvalidOperationException(OcspMessages.ErrUnableToPlugTLSServer));
        }

        var tlsConfig = config.TlsConfig;
        var tlsOptions = config.TlsOptions;
        if (tlsOptions.OcspPeerConfig == null || !tlsOptions.OcspPeerConfig.Verify)
        {
            return (tlsConfig, false, null);
        }

        tlsConfig.RemoteCertificateValidationCallback = (_, _, chain, _) =>
        {
            if (chain?.ChainElements == null || chain.ChainElements.Count == 0)
            {
                return false;
            }

            var converted = chain.ChainElements
                .Select(e => e.Certificate)
                .OfType<X509Certificate2>()
                .ToArray();
            return TlsServerOCSPValid([converted], tlsOptions.OcspPeerConfig);
        };

        return (tlsConfig, true, null);
    }

    internal bool TlsServerOCSPValid(X509Certificate2[][] chains, OcspPeerConfig options)
    {
        Debugf(OcspMessages.DbgNumServerChains, chains.Length);
        return PeerOCSPValid(chains, options);
    }

    internal bool TlsClientOCSPValid(X509Certificate2[][] chains, OcspPeerConfig options)
    {
        Debugf(OcspMessages.DbgNumClientChains, chains.Length);
        return PeerOCSPValid(chains, options);
    }

    internal bool PeerOCSPValid(X509Certificate2[][] chains, OcspPeerConfig options)
    {
        var peer = OcspHandler.PeerFromVerifiedChains(chains);
        if (peer == null)
        {
            Errorf(OcspMessages.ErrPeerEmptyAutoReject);
            return false;
        }

        for (var chainIndex = 0; chainIndex < chains.Length; chainIndex++)
        {
            var chain = chains[chainIndex];
            Debugf(OcspMessages.DbgLinksInChain, chainIndex, chain.Length);

            if (chain.Length == 1)
            {
                Debugf(OcspMessages.DbgSelfSignedValid, chainIndex);
                return true;
            }

            var chainEligible = false;
            var eligibleLinks = new List<ChainLink>();
            for (var linkPos = 0; linkPos < chain.Length - 1; linkPos++)
            {
                var cert = chain[linkPos];
                var link = new ChainLink { Leaf = cert };
                if (OcspUtilities.CertOCSPEligible(link))
                {
                    chainEligible = true;
                    link.Issuer = chain[linkPos + 1];
                    eligibleLinks.Add(link);
                }
            }

            if (!chainEligible)
            {
                Debugf(OcspMessages.DbgValidNonOCSPChain, chainIndex);
                return true;
            }

            Debugf(OcspMessages.DbgChainIsOCSPEligible, chainIndex, eligibleLinks.Count);
            var chainValid = true;
            foreach (var link in eligibleLinks)
            {
                var (reason, good) = CertOCSPGood(link, options);
                if (good)
                {
                    continue;
                }

                Debugf(reason);
                chainValid = false;
                break;
            }

            if (chainValid)
            {
                Debugf(OcspMessages.DbgChainIsOCSPValid, chainIndex);
                return true;
            }
        }

        Debugf(OcspMessages.DbgNoOCSPValidChains);
        return false;
    }

    internal (string reason, bool good) CertOCSPGood(ChainLink? link, OcspPeerConfig options)
    {
        if (link?.Leaf == null || link.Issuer == null || link.OcspWebEndpoints == null || link.OcspWebEndpoints.Count == 0)
        {
            return ("Empty chainlink found", false);
        }

        var log = new OcspLog
        {
            Debugf = (fmt, args) => Debugf(fmt, args),
            Noticef = (fmt, args) => Noticef(fmt, args),
            Warnf = (fmt, args) => Warnf(fmt, args),
            Errorf = (fmt, args) => Errorf(fmt, args),
        };

        var fingerprint = Convert.ToHexString(SHA256.HashData(link.Leaf.RawData)).ToLowerInvariant();
        var cached = _ocsprc?.Get(fingerprint);
        if (cached is { Length: > 0 })
        {
            var (parsed, parseError) = OcspHandler.ParseOcspResponse(cached);
            if (parseError == null && parsed != null && OcspUtilities.OCSPResponseCurrent(parsed, options, log))
            {
                if (parsed.Status == OcspStatusAssertion.Revoked ||
                    (parsed.Status == OcspStatusAssertion.Unknown && !options.UnknownIsGood))
                {
                    if (options.WarnOnly)
                    {
                        Warnf("allowing OCSP peer due warn_only for [{0}]", link.Leaf.Subject);
                        return (string.Empty, true);
                    }

                    return ($"OCSP response invalid status: {OcspStatusAssertionExtensions.GetStatusAssertionStr((int)parsed.Status)}", false);
                }

                Debugf(OcspMessages.DbgOCSPValidPeerLink, link.Leaf.Subject);
                return (string.Empty, true);
            }
        }

        if (options.AllowWhenCAUnreachable || options.WarnOnly)
        {
            if (options.WarnOnly)
            {
                Warnf("allowing OCSP peer due warn_only for [{0}]", link.Leaf.Subject);
            }

            if (options.AllowWhenCAUnreachable)
            {
                Warnf("allowing OCSP peer due unreachable CA for [{0}]", link.Leaf.Subject);
            }

            return (string.Empty, true);
        }

        return ("failed to fetch OCSP response", false);
    }
}