using System.Security.Cryptography;
using Azaion.Common.Configs;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Microsoft.IdentityModel.Tokens;

namespace Azaion.Services;

/// <summary>
/// AZ-532 — loads ES256 signing keys from <see cref="JwtConfig.KeysFolder"/>.
/// One key is "active" (used to sign new tokens); the rest stay in JWKS so
/// in-flight tokens minted with older keys still verify during the rotation
/// overlap window. The kid of each key is its filename without <c>.pem</c>.
/// </summary>
public interface IJwtSigningKeyProvider
{
    JwtSigningKey                      Active { get; }
    IReadOnlyList<JwtSigningKey>       All    { get; }
}

public sealed class JwtSigningKey
{
    public string           Kid { get; }
    public ECDsa            Ecdsa { get; }
    public ECDsaSecurityKey SecurityKey { get; }

    public JwtSigningKey(string kid, ECDsa ecdsa)
    {
        Kid         = kid;
        Ecdsa       = ecdsa;
        SecurityKey = new ECDsaSecurityKey(ecdsa) { KeyId = kid };
    }
}

public class JwtSigningKeyProvider : IJwtSigningKeyProvider, IDisposable
{
    private readonly Dictionary<string, JwtSigningKey> _byKid;
    private readonly JwtSigningKey                     _active;

    public JwtSigningKeyProvider(IOptions<JwtConfig> jwtConfig, ILogger<JwtSigningKeyProvider> logger)
    {
        var folder = jwtConfig.Value.KeysFolder;

        if (string.IsNullOrWhiteSpace(folder) || !Directory.Exists(folder))
            throw new InvalidOperationException(
                $"JwtConfig.KeysFolder '{folder}' does not exist. " +
                "Generate a key with scripts/generate-jwt-key.sh and ensure the folder is mounted into the container.");

        var pemFiles = Directory.EnumerateFiles(folder, "*.pem").OrderBy(p => p).ToList();
        if (pemFiles.Count == 0)
            throw new InvalidOperationException(
                $"No *.pem keys found in '{folder}'. Generate a key with scripts/generate-jwt-key.sh.");

        _byKid = new Dictionary<string, JwtSigningKey>(StringComparer.Ordinal);
        foreach (var path in pemFiles)
        {
            var kid   = Path.GetFileNameWithoutExtension(path);
            var ecdsa = ECDsa.Create();
            try
            {
                ecdsa.ImportFromPem(File.ReadAllText(path));
            }
            catch (Exception ex)
            {
                ecdsa.Dispose();
                throw new InvalidOperationException($"Failed to load JWT signing key from '{path}'.", ex);
            }

            EnsureP256(ecdsa, path);
            _byKid[kid] = new JwtSigningKey(kid, ecdsa);
        }

        var requestedActive = jwtConfig.Value.ActiveKid;
        if (!string.IsNullOrEmpty(requestedActive))
        {
            if (!_byKid.TryGetValue(requestedActive, out var resolved))
                throw new InvalidOperationException(
                    $"JwtConfig.ActiveKid '{requestedActive}' is not present in '{folder}'.");
            _active = resolved;
        }
        else
        {
            _active = _byKid[Path.GetFileNameWithoutExtension(pemFiles[0])];
            logger.LogInformation(
                "JwtConfig.ActiveKid not set; falling back to first key by filename: {Kid}", _active.Kid);
        }
    }

    public JwtSigningKey                Active => _active;
    public IReadOnlyList<JwtSigningKey> All    => _byKid.Values.OrderBy(k => k.Kid, StringComparer.Ordinal).ToList();

    private static void EnsureP256(ECDsa ecdsa, string path)
    {
        // ES256 ⇒ P-256 (prime256v1 / secp256r1). Reject anything else so we don't
        // silently sign with the wrong curve and break verifiers expecting ES256.
        var p = ecdsa.ExportParameters(includePrivateParameters: false);
        var oid = p.Curve.Oid?.Value ?? p.Curve.Oid?.FriendlyName;
        if (oid is not ("1.2.840.10045.3.1.7" or "nistP256" or "ECDSA_P256"))
            throw new InvalidOperationException(
                $"Key '{path}' is not on the P-256 curve (got '{oid ?? "unknown"}'). ES256 requires P-256.");
    }

    public void Dispose()
    {
        foreach (var k in _byKid.Values) k.Ecdsa.Dispose();
        _byKid.Clear();
    }
}