using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;

namespace Azaion.Missions.JwksMock.Services;

/// <summary>
/// Hand-rolls JWS-compact ES256 tokens for tests. Honors per-call overrides
/// the test harness uses to exercise NFT-SEC-* (alg confusion, unknown kid,
/// claim mismatch, etc.).
/// </summary>
public sealed class TokenSigner
{
    private readonly KeyStore _keys;
    private readonly TimeProvider _clock;
    private readonly string _defaultIssuer;
    private readonly string _defaultAudience;

    public TokenSigner(KeyStore keys, TimeProvider clock, string defaultIssuer, string defaultAudience)
    {
        _keys = keys;
        _clock = clock;
        _defaultIssuer = defaultIssuer;
        _defaultAudience = defaultAudience;
    }

    public SignResult Sign(SignRequest request)
    {
        var active = _keys.Active;
        var kid = request.KidOverride ?? active.Kid;
        var alg = request.AlgOverride ?? "ES256";

        if (request.Permissions is not null && request.PermissionsArray is not null)
            throw new ArgumentException(
                "permissions and permissions_array are mutually exclusive — set at most one.",
                nameof(request));

        // NFT-SEC-11 AC-5.4: the mock refuses kid_override values that don't
        // correspond to a currently-published kid (active or in-grace retired).
        // Without this guard, a tester could mint a token with any kid string
        // and the SUT would simply 401 on JWKS lookup — defeating the
        // "post-grace mock refuses old kid" assertion.
        if (request.KidOverride is not null)
        {
            var known = _keys.PublishedKeys().Select(k => k.Kid).ToHashSet(StringComparer.Ordinal);
            if (!known.Contains(request.KidOverride))
                throw new ArgumentException(
                    $"kid_override '{request.KidOverride}' is not a currently-published kid.",
                    nameof(request));
        }

        var nowUnix = _clock.GetUtcNow().ToUnixTimeSeconds();
        var expUnix = nowUnix + (request.ExpOffsetSeconds ?? 3600);

        var header = new JsonObject
        {
            ["alg"] = alg,
            ["kid"] = kid,
            ["typ"] = "JWT"
        };

        var payload = new JsonObject
        {
            ["iss"] = request.Issuer ?? _defaultIssuer,
            ["aud"] = request.Audience ?? _defaultAudience,
            ["iat"] = nowUnix,
            ["exp"] = expUnix
        };
        if (request.Permissions is not null)
            payload["permissions"] = request.Permissions;
        if (request.PermissionsArray is not null)
        {
            var arr = new JsonArray();
            foreach (var p in request.PermissionsArray)
                arr.Add(p);
            payload["permissions"] = arr;
        }
        if (request.Subject is not null)
            payload["sub"] = request.Subject;

        var headerBytes  = JsonSerializer.SerializeToUtf8Bytes(header);
        var payloadBytes = JsonSerializer.SerializeToUtf8Bytes(payload);
        var headerSeg  = Base64Url.Encode(headerBytes);
        var payloadSeg = Base64Url.Encode(payloadBytes);

        // Signing input is the literal ASCII string "<header>.<payload>" per RFC 7515 §5.1.
        var signingInput = Encoding.ASCII.GetBytes($"{headerSeg}.{payloadSeg}");

        byte[] signature;
        if (alg == "ES256")
        {
            signature = active.Ec.SignData(signingInput, HashAlgorithmName.SHA256, DSASignatureFormat.IeeeP1363FixedFieldConcatenation);
        }
        else if (alg == "HS256")
        {
            // alg-confusion attack vector for NFT-SEC-10. We sign with a key derived
            // from the active public key so a naive validator that fails to enforce
            // alg pinning would accept the token.
            var pubKey = active.Ec.ExportSubjectPublicKeyInfo();
            using var hmac = new HMACSHA256(pubKey);
            signature = hmac.ComputeHash(signingInput);
        }
        else if (alg == "none")
        {
            signature = [];
        }
        else
        {
            throw new ArgumentException($"Unsupported alg_override '{alg}'", nameof(request));
        }

        var sigSeg = Base64Url.Encode(signature);
        var token = $"{headerSeg}.{payloadSeg}.{sigSeg}";
        return new SignResult(token, kid);
    }
}

public sealed record SignRequest(
    string? Issuer,
    string? Audience,
    int? ExpOffsetSeconds,
    string? Permissions,
    string[]? PermissionsArray,
    string? Subject,
    string? AlgOverride,
    string? KidOverride);

public sealed record SignResult(string Token, string Kid);