using Net.Pkcs11Interop.Common;
using Net.Pkcs11Interop.HighLevelAPI;
using Net.Pkcs11Interop.HighLevelAPI.Factories;
using StellaOps.Cryptography;
using System;
using System.Collections.Generic;
using System.Linq;

namespace StellaOps.Cryptography.Plugin.Pkcs11Gost;

internal static class Pkcs11SignerUtilities
{
    private static readonly Pkcs11InteropFactories Factories = new();

    public static byte[] SignDigest(Pkcs11GostKeyEntry entry, ReadOnlySpan<byte> digest)
    {
        using var pkcs11 = Factories.Pkcs11LibraryFactory.LoadPkcs11Library(Factories, entry.Session.LibraryPath, AppType.MultiThreaded);
        var slot = ResolveSlot(pkcs11, entry.Session);
        if (slot is null)
        {
            throw new InvalidOperationException("No PKCS#11 slot/token matched the provided configuration.");
        }

        using var session = slot.OpenSession(SessionType.ReadWrite);
        var loggedIn = false;
        try
        {
            var pin = ResolvePin(entry.Session);
            if (!string.IsNullOrWhiteSpace(pin))
            {
                session.Login(CKU.CKU_USER, pin);
                loggedIn = true;
            }

            var privateHandle = FindObject(session, CKO.CKO_PRIVATE_KEY, entry.Session.PrivateKeyLabel);
            if (privateHandle is null)
            {
                throw new InvalidOperationException($"Private key with label '{entry.Session.PrivateKeyLabel}' was not found.");
            }

            using var mechanism = Factories.MechanismFactory.Create(entry.SignMechanismId);
            return session.Sign(mechanism, privateHandle, digest.ToArray());
        }
        finally
        {
            if (loggedIn)
            {
                try { session.Logout(); } catch { /* ignored */ }
            }
        }
    }

    private static ISlot? ResolveSlot(IPkcs11Library pkcs11, Pkcs11SessionOptions options)
    {
        var slots = pkcs11.GetSlotList(SlotsType.WithTokenPresent);
        if (slots.Count == 0)
        {
            return null;
        }

        if (!string.IsNullOrWhiteSpace(options.SlotId))
        {
            return slots.FirstOrDefault(slot =>
                string.Equals(slot.SlotId.ToString(), options.SlotId, StringComparison.OrdinalIgnoreCase));
        }

        if (!string.IsNullOrWhiteSpace(options.TokenLabel))
        {
            return slots.FirstOrDefault(slot =>
            {
                var tokenInfo = slot.GetTokenInfo();
                return string.Equals(tokenInfo.Label?.Trim(), options.TokenLabel?.Trim(), StringComparison.OrdinalIgnoreCase);
            });
        }

        return slots[0];
    }

    private static IObjectHandle? FindObject(ISession session, CKO objectClass, string? label)
    {
        var template = new List<IObjectAttribute>
        {
            Factories.ObjectAttributeFactory.Create(CKA.CKA_CLASS, (uint)objectClass)
        };

        if (!string.IsNullOrWhiteSpace(label))
        {
            template.Add(Factories.ObjectAttributeFactory.Create(CKA.CKA_LABEL, label));
        }

        var handles = session.FindAllObjects(template);
        return handles.FirstOrDefault();
    }

    private static string? ResolvePin(Pkcs11SessionOptions options)
    {
        if (!string.IsNullOrWhiteSpace(options.UserPin))
        {
            return options.UserPin;
        }

        if (!string.IsNullOrWhiteSpace(options.UserPinEnvironmentVariable))
        {
            return Environment.GetEnvironmentVariable(options.UserPinEnvironmentVariable);
        }

        return null;
    }
}