// Copyright 2022-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 certstore/certstore.go and certstore/certstore_windows.go in
// the NATS server Go source. The .NET implementation uses System.Security.
// Cryptography.X509Certificates.X509Store in place of Win32 P/Invoke calls.

using System.Runtime.InteropServices;
using System.Security.Cryptography.X509Certificates;

namespace ZB.MOM.NatsNet.Server.Auth.CertificateStore;

/// <summary>
/// Provides access to the Windows certificate store for TLS certificate provisioning.
/// Mirrors certstore/certstore.go and certstore/certstore_windows.go.
///
/// On non-Windows platforms all methods that require the Windows store throw
/// <see cref="CertStoreErrors.ErrOSNotCompatCertStore"/>.
/// </summary>
public static class CertificateStoreService
{
    private static readonly IReadOnlyDictionary<string, StoreType> StoreMap =
        new Dictionary<string, StoreType>(StringComparer.OrdinalIgnoreCase)
        {
            ["windowscurrentuser"] = StoreType.WindowsCurrentUser,
            ["windowslocalmachine"] = StoreType.WindowsLocalMachine,
        };

    private static readonly IReadOnlyDictionary<string, MatchByType> MatchByMap =
        new Dictionary<string, MatchByType>(StringComparer.OrdinalIgnoreCase)
        {
            ["issuer"] = MatchByType.Issuer,
            ["subject"] = MatchByType.Subject,
            ["thumbprint"] = MatchByType.Thumbprint,
        };

    // -------------------------------------------------------------------------
    // Cross-platform parse helpers
    // -------------------------------------------------------------------------

    /// <summary>
    /// Parses a cert_store string to a <see cref="StoreType"/>.
    /// Returns an error if the string is unrecognised or not valid on the current OS.
    /// Mirrors <c>ParseCertStore</c>.
    /// </summary>
    public static (StoreType store, Exception? error) ParseCertStore(string certStore)
    {
        if (!StoreMap.TryGetValue(certStore, out var st))
            return (StoreType.Empty, CertStoreErrors.ErrBadCertStore);

        // All currently supported store types are Windows-only.
        if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            return (StoreType.Empty, CertStoreErrors.ErrOSNotCompatCertStore);

        return (st, null);
    }

    /// <summary>
    /// Parses a cert_match_by string to a <see cref="MatchByType"/>.
    /// Mirrors <c>ParseCertMatchBy</c>.
    /// </summary>
    public static (MatchByType matchBy, Exception? error) ParseCertMatchBy(string certMatchBy)
    {
        if (!MatchByMap.TryGetValue(certMatchBy, out var mb))
            return (MatchByType.Empty, CertStoreErrors.ErrBadMatchByType);
        return (mb, null);
    }

    /// <summary>
    /// Returns the issuer certificate for <paramref name="leaf"/> by building a chain.
    /// Returns null if the chain cannot be built or the leaf is self-signed.
    /// Mirrors <c>GetLeafIssuer</c>.
    /// </summary>
    public static X509Certificate2? GetLeafIssuer(X509Certificate2 leaf)
    {
        using var chain = new X509Chain();
        chain.ChainPolicy.RevocationMode = X509RevocationMode.NoCheck;
        chain.ChainPolicy.VerificationFlags = X509VerificationFlags.AllowUnknownCertificateAuthority;

        if (!chain.Build(leaf) || chain.ChainElements.Count < 2)
            return null;

        // chain.ChainElements[0] is the leaf; [1] is its issuer.
        return new X509Certificate2(chain.ChainElements[1].Certificate);
    }

    // -------------------------------------------------------------------------
    // TLS configuration entry point
    // -------------------------------------------------------------------------

    /// <summary>
    /// Finds a certificate in the Windows certificate store matching the given criteria and
    /// returns a <see cref="CertStoreTlsResult"/> suitable for populating TLS options.
    ///
    /// On non-Windows platforms throws <see cref="CertStoreErrors.ErrOSNotCompatCertStore"/>.
    /// Mirrors <c>TLSConfig</c> (certstore_windows.go).
    /// </summary>
    /// <param name="storeType">Which Windows store to use (CurrentUser or LocalMachine).</param>
    /// <param name="matchBy">How to match the certificate (Subject, Issuer, or Thumbprint).</param>
    /// <param name="certMatch">The match value (subject name, issuer name, or thumbprint hex).</param>
    /// <param name="caCertsMatch">Optional list of subject strings to locate CA certificates.</param>
    /// <param name="skipInvalid">If true, skip expired or not-yet-valid certificates.</param>
    public static CertStoreTlsResult TLSConfig(
        StoreType storeType,
        MatchByType matchBy,
        string certMatch,
        IReadOnlyList<string>? caCertsMatch = null,
        bool skipInvalid = false)
    {
        if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            throw CertStoreErrors.ErrOSNotCompatCertStore;
        if (storeType is not (StoreType.WindowsCurrentUser or StoreType.WindowsLocalMachine))
            throw CertStoreErrors.ErrBadCertStore;

        var location = storeType == StoreType.WindowsCurrentUser
            ? StoreLocation.CurrentUser
            : StoreLocation.LocalMachine;

        // Find the leaf certificate.
        var leaf = matchBy switch
        {
            MatchByType.Subject or MatchByType.Empty => CertBySubject(certMatch, location, skipInvalid),
            MatchByType.Issuer => CertByIssuer(certMatch, location, skipInvalid),
            MatchByType.Thumbprint => CertByThumbprint(certMatch, location, skipInvalid),
            _ => throw CertStoreErrors.ErrBadMatchByType,
        } ?? throw CertStoreErrors.ErrFailedCertSearch;

        // Optionally find CA certificates.
        X509Certificate2Collection? caPool = null;
        if (caCertsMatch is { Count: > 0 })
            caPool = CreateCACertsPool(location, caCertsMatch, skipInvalid);

        return new CertStoreTlsResult(leaf, caPool);
    }

    // -------------------------------------------------------------------------
    // Certificate search helpers (mirror winCertStore.certByXxx / certSearch)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Finds the first certificate in the personal (MY) store by subject name.
    /// Mirrors <c>certBySubject</c>.
    /// </summary>
    public static X509Certificate2? CertBySubject(string subject, StoreLocation location, bool skipInvalid) =>
        CertSearch(StoreName.My, location, X509FindType.FindBySubjectName, subject, skipInvalid);

    /// <summary>
    /// Finds the first certificate in the personal (MY) store by issuer name.
    /// Mirrors <c>certByIssuer</c>.
    /// </summary>
    public static X509Certificate2? CertByIssuer(string issuer, StoreLocation location, bool skipInvalid) =>
        CertSearch(StoreName.My, location, X509FindType.FindByIssuerName, issuer, skipInvalid);

    /// <summary>
    /// Finds the first certificate in the personal (MY) store by SHA-1 thumbprint (hex string).
    /// Mirrors <c>certByThumbprint</c>.
    /// </summary>
    public static X509Certificate2? CertByThumbprint(string thumbprint, StoreLocation location, bool skipInvalid) =>
        CertSearch(StoreName.My, location, X509FindType.FindByThumbprint, thumbprint, skipInvalid);

    /// <summary>
    /// Searches Root, AuthRoot, and CA stores for certificates matching the given subject name.
    /// Returns all matching certificates across all three locations.
    /// Mirrors <c>caCertsBySubjectMatch</c>.
    /// </summary>
    public static IReadOnlyList<X509Certificate2> CaCertsBySubjectMatch(
        string subject,
        StoreLocation location,
        bool skipInvalid)
    {
        if (string.IsNullOrEmpty(subject))
            throw CertStoreErrors.ErrBadCaCertMatchField;

        var results = new List<X509Certificate2>();
        var searchLocations = new[] { StoreName.Root, StoreName.AuthRoot, StoreName.CertificateAuthority };

        foreach (var storeName in searchLocations)
        {
            var cert = CertSearch(storeName, location, X509FindType.FindBySubjectName, subject, skipInvalid);
            if (cert != null)
                results.Add(cert);
        }

        if (results.Count == 0)
            throw CertStoreErrors.ErrFailedCertSearch;

        return results;
    }

    /// <summary>
    /// Core certificate search — opens the specified store and finds a matching certificate.
    /// Returns null if not found.
    /// Mirrors <c>certSearch</c>.
    /// </summary>
    public static X509Certificate2? CertSearch(
        StoreName storeName,
        StoreLocation storeLocation,
        X509FindType findType,
        string findValue,
        bool skipInvalid)
    {
        using var store = new X509Store(storeName, storeLocation, OpenFlags.ReadOnly | OpenFlags.OpenExistingOnly);
        var certs = store.Certificates.Find(findType, findValue, validOnly: skipInvalid);
        if (certs.Count == 0)
            return null;

        // Pick first that has a private key (mirrors certKey requirement in Go).
        foreach (var cert in certs)
        {
            if (cert.HasPrivateKey)
                return cert;
        }

        // Fall back to first even without private key (e.g. CA cert lookup).
        return certs[0];
    }

    // -------------------------------------------------------------------------
    // CA cert pool builder (mirrors createCACertsPool)
    // -------------------------------------------------------------------------

    /// <summary>
    /// Builds a collection of CA certificates from the trusted Root, AuthRoot, and CA stores
    /// for each subject name in <paramref name="caCertsMatch"/>.
    /// Mirrors <c>createCACertsPool</c>.
    /// </summary>
    public static X509Certificate2Collection CreateCACertsPool(
        StoreLocation location,
        IReadOnlyList<string> caCertsMatch,
        bool skipInvalid)
    {
        var pool = new X509Certificate2Collection();
        var failCount = 0;

        foreach (var subject in caCertsMatch)
        {
            try
            {
                var matches = CaCertsBySubjectMatch(subject, location, skipInvalid);
                foreach (var cert in matches)
                    pool.Add(cert);
            }
            catch
            {
                failCount++;
            }
        }

        if (failCount == caCertsMatch.Count)
            throw new InvalidOperationException("unable to match any CA certificate");

        return pool;
    }
}
