using System.Collections.Concurrent;
using System.Formats.Asn1;
using System.Security.Cryptography;
using Microsoft.IdentityModel.Tokens;
using Net.Pkcs11Interop.Common;
using Net.Pkcs11Interop.HighLevelAPI;
using Net.Pkcs11Interop.HighLevelAPI.Factories;

namespace StellaOps.Cryptography.Kms;

public interface IPkcs11Facade : IDisposable
{
    Task<Pkcs11KeyDescriptor> GetKeyAsync(CancellationToken cancellationToken);

    Task<Pkcs11PublicKeyMaterial> GetPublicKeyAsync(CancellationToken cancellationToken);

    Task<byte[]> SignDigestAsync(ReadOnlyMemory<byte> digest, CancellationToken cancellationToken);
}

public sealed record Pkcs11KeyDescriptor(
    string KeyId,
    string? Label,
    DateTimeOffset CreatedAt);

public sealed record Pkcs11PublicKeyMaterial(
    string KeyId,
    string Curve,
    byte[] Qx,
    byte[] Qy);

internal sealed class Pkcs11InteropFacade : IPkcs11Facade
{
    private readonly Pkcs11Options _options;
    private readonly Pkcs11InteropFactories _factories;
    private readonly IPkcs11Library _library;
    private readonly ISlot _slot;
    private readonly ConcurrentDictionary<string, IObjectAttribute[]> _attributeCache = new(StringComparer.Ordinal);

    public Pkcs11InteropFacade(Pkcs11Options options)
    {
        _options = options ?? throw new ArgumentNullException(nameof(options));
        if (string.IsNullOrWhiteSpace(_options.LibraryPath))
        {
            throw new ArgumentException("PKCS#11 library path must be provided.", nameof(options));
        }

        _factories = new Pkcs11InteropFactories();
        _library = _factories.Pkcs11LibraryFactory.LoadPkcs11Library(_factories, _options.LibraryPath, AppType.MultiThreaded);
        _slot = ResolveSlot(_library, _options)
            ?? throw new InvalidOperationException("Could not resolve PKCS#11 slot.");
    }

    public async Task<Pkcs11KeyDescriptor> GetKeyAsync(CancellationToken cancellationToken)
    {
        using var context = await OpenSessionAsync(cancellationToken).ConfigureAwait(false);
        var session = context.Session;
        var privateHandle = FindKey(session, CKO.CKO_PRIVATE_KEY, _options.PrivateKeyLabel);
        if (privateHandle is null)
        {
            throw new InvalidOperationException("PKCS#11 private key not found.");
        }

        var labelAttr = GetAttribute(session, privateHandle, CKA.CKA_LABEL);
        var label = labelAttr?.GetValueAsString();

        return new Pkcs11KeyDescriptor(
            KeyId: label ?? privateHandle.ObjectId.ToString(),
            Label: label,
            CreatedAt: DateTimeOffset.UtcNow);
    }

    public async Task<Pkcs11PublicKeyMaterial> GetPublicKeyAsync(CancellationToken cancellationToken)
    {
        using var context = await OpenSessionAsync(cancellationToken).ConfigureAwait(false);
        var session = context.Session;
        var publicHandle = FindKey(session, CKO.CKO_PUBLIC_KEY, _options.PublicKeyLabel ?? _options.PrivateKeyLabel);
        if (publicHandle is null)
        {
            throw new InvalidOperationException("PKCS#11 public key not found.");
        }

        var pointAttr = GetAttribute(session, publicHandle, CKA.CKA_EC_POINT)
            ?? throw new InvalidOperationException("Public key missing EC point.");
        var paramsAttr = GetAttribute(session, publicHandle, CKA.CKA_EC_PARAMS)
            ?? throw new InvalidOperationException("Public key missing EC parameters.");

        var ecPoint = ExtractEcPoint(pointAttr.GetValueAsByteArray());
        var (curve, coordinateSize) = DecodeCurve(paramsAttr.GetValueAsByteArray());

        if (ecPoint.Length != 1 + (coordinateSize * 2) || ecPoint[0] != 0x04)
        {
            throw new InvalidOperationException("Unsupported EC point format.");
        }

        var qx = ecPoint.AsSpan(1, coordinateSize).ToArray();
        var qy = ecPoint.AsSpan(1 + coordinateSize, coordinateSize).ToArray();

        var keyId = GetAttribute(session, publicHandle, CKA.CKA_LABEL)?.GetValueAsString()
            ?? publicHandle.ObjectId.ToString();

        return new Pkcs11PublicKeyMaterial(
            keyId,
            curve,
            qx,
            qy);
    }

    public async Task<byte[]> SignDigestAsync(ReadOnlyMemory<byte> digest, CancellationToken cancellationToken)
    {
        using var context = await OpenSessionAsync(cancellationToken).ConfigureAwait(false);
        var session = context.Session;
        var privateHandle = FindKey(session, CKO.CKO_PRIVATE_KEY, _options.PrivateKeyLabel)
            ?? throw new InvalidOperationException("PKCS#11 private key not found.");

        var mechanism = _factories.MechanismFactory.Create(_options.MechanismId);
        return session.Sign(mechanism, privateHandle, digest.ToArray());
    }

    private async Task<SessionContext> OpenSessionAsync(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        var session = _slot.OpenSession(SessionType.ReadOnly);

        var loggedIn = false;

        try
        {
            if (!string.IsNullOrWhiteSpace(_options.UserPin))
            {
                session.Login(CKU.CKU_USER, _options.UserPin);
                loggedIn = true;
            }

            return new SessionContext(session, loggedIn);
        }
        catch
        {
            if (loggedIn)
            {
                try { session.Logout(); } catch { /* ignore */ }
            }

            session.Dispose();
            throw;
        }
    }

    private IObjectHandle? FindKey(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 IObjectAttribute? GetAttribute(ISession session, IObjectHandle handle, CKA type)
    {
        var cacheKey = $"{handle.ObjectId}:{(uint)type}";
        if (_attributeCache.TryGetValue(cacheKey, out var cached))
        {
            return cached.FirstOrDefault();
        }

        var attributes = session.GetAttributeValue(handle, new List<CKA> { type })
            ?.ToArray() ?? Array.Empty<IObjectAttribute>();

        if (attributes.Length > 0)
        {
            _attributeCache[cacheKey] = attributes;
            return attributes[0];
        }

        return null;
    }

    private static ISlot? ResolveSlot(IPkcs11Library pkcs11, Pkcs11Options 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 info = slot.GetTokenInfo();
                return string.Equals(info.Label?.Trim(), options.TokenLabel.Trim(), StringComparison.Ordinal);
            });
        }

        return slots[0];
    }

    private static byte[] ExtractEcPoint(byte[] derEncoded)
    {
        var reader = new AsnReader(derEncoded, AsnEncodingRules.DER);
        var point = reader.ReadOctetString();
        reader.ThrowIfNotEmpty();
        return point;
    }

    private static (string CurveName, int CoordinateSize) DecodeCurve(byte[] ecParamsDer)
    {
        var reader = new AsnReader(ecParamsDer, AsnEncodingRules.DER);
        var oid = reader.ReadObjectIdentifier();
        reader.ThrowIfNotEmpty();

        var curve = oid switch
        {
            "1.2.840.10045.3.1.7" => JsonWebKeyECTypes.P256,
            "1.3.132.0.34" => JsonWebKeyECTypes.P384,
            "1.3.132.0.35" => JsonWebKeyECTypes.P521,
            _ => throw new InvalidOperationException($"Unsupported EC curve OID '{oid}'."),
        };

        var coordinateSize = curve switch
        {
            JsonWebKeyECTypes.P256 => 32,
            JsonWebKeyECTypes.P384 => 48,
            JsonWebKeyECTypes.P521 => 66,
            _ => throw new InvalidOperationException($"Unsupported EC curve '{curve}'."),
        };

        return (curve, coordinateSize);
    }

    public void Dispose()
    {
        _library.Dispose();
    }

    private sealed class SessionContext : System.IDisposable
    {
        private readonly ISession _session;
        private readonly bool _logoutOnDispose;
        private bool _disposed;

        public SessionContext(ISession session, bool logoutOnDispose)
        {
            _session = session ?? throw new System.ArgumentNullException(nameof(session));
            _logoutOnDispose = logoutOnDispose;
        }

        public ISession Session => _session;

        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            if (_logoutOnDispose)
            {
                try
                {
                    _session.Logout();
                }
                catch
                {
                    // ignore logout failures
                }
            }

            _session.Dispose();
            _disposed = true;
        }
    }
}