using System.Text;
using System.Text.Json;
using NATS.NKeys;

namespace NATS.Server.Auth.Jwt;

/// <summary>
/// Provides NATS JWT decode, verify, and claim extraction.
/// NATS JWTs are standard JWT format (base64url header.payload.signature) with Ed25519 signing.
/// All NATS JWTs start with "eyJ" (base64url for '{"').
/// </summary>
/// <remarks>
/// Reference: golang/nats-server/server/jwt.go and github.com/nats-io/jwt/v2
/// </remarks>
public static class NatsJwt
{
    private const string JwtPrefix = "eyJ";

    /// <summary>
    /// Returns true if the string appears to be a JWT (starts with "eyJ").
    /// </summary>
    public static bool IsJwt(string token)
    {
        return !string.IsNullOrEmpty(token) && token.StartsWith(JwtPrefix, StringComparison.Ordinal);
    }

    /// <summary>
    /// Decodes a JWT token into its constituent parts without verifying the signature.
    /// Returns null if the token is structurally invalid.
    /// </summary>
    public static JwtToken? Decode(string token)
    {
        if (string.IsNullOrEmpty(token))
            return null;

        var parts = token.Split('.');
        if (parts.Length != 3)
            return null;

        try
        {
            var headerBytes = Base64UrlDecode(parts[0]);
            var payloadBytes = Base64UrlDecode(parts[1]);
            var signatureBytes = Base64UrlDecode(parts[2]);

            var header = JsonSerializer.Deserialize<JwtHeader>(headerBytes);
            if (header is null)
                return null;

            var payloadJson = Encoding.UTF8.GetString(payloadBytes);
            var signingInput = $"{parts[0]}.{parts[1]}";

            return new JwtToken
            {
                Header = header,
                PayloadJson = payloadJson,
                Signature = signatureBytes,
                SigningInput = signingInput,
            };
        }
        catch
        {
            return null;
        }
    }

    /// <summary>
    /// Decodes a JWT token and deserializes the payload as <see cref="UserClaims"/>.
    /// Returns null if the token is structurally invalid or cannot be deserialized.
    /// </summary>
    public static UserClaims? DecodeUserClaims(string token)
    {
        var jwt = Decode(token);
        if (jwt is null)
            return null;

        try
        {
            return JsonSerializer.Deserialize<UserClaims>(jwt.PayloadJson);
        }
        catch
        {
            return null;
        }
    }

    /// <summary>
    /// Decodes a JWT token and deserializes the payload as <see cref="AccountClaims"/>.
    /// Returns null if the token is structurally invalid or cannot be deserialized.
    /// </summary>
    public static AccountClaims? DecodeAccountClaims(string token)
    {
        var jwt = Decode(token);
        if (jwt is null)
            return null;

        try
        {
            return JsonSerializer.Deserialize<AccountClaims>(jwt.PayloadJson);
        }
        catch
        {
            return null;
        }
    }

    /// <summary>
    /// Verifies the Ed25519 signature on a JWT token against the given NKey public key.
    /// </summary>
    public static bool Verify(string token, string publicNkey)
    {
        try
        {
            var jwt = Decode(token);
            if (jwt is null)
                return false;

            var kp = KeyPair.FromPublicKey(publicNkey);
            var signingInputBytes = Encoding.UTF8.GetBytes(jwt.SigningInput);
            return kp.Verify(signingInputBytes, jwt.Signature);
        }
        catch
        {
            return false;
        }
    }

    /// <summary>
    /// Verifies a nonce signature against the given NKey public key.
    /// Tries base64url decoding first, then falls back to standard base64 (Go compatibility).
    /// </summary>
    public static bool VerifyNonce(byte[] nonce, string signature, string publicNkey)
    {
        try
        {
            var sigBytes = TryDecodeSignature(signature);
            if (sigBytes is null)
                return false;

            var kp = KeyPair.FromPublicKey(publicNkey);
            return kp.Verify(nonce, sigBytes);
        }
        catch
        {
            return false;
        }
    }

    /// <summary>
    /// Decodes a base64url-encoded byte array.
    /// Replaces URL-safe characters and adds padding as needed.
    /// </summary>
    internal static byte[] Base64UrlDecode(string input)
    {
        var s = input.Replace('-', '+').Replace('_', '/');
        switch (s.Length % 4)
        {
            case 2: s += "=="; break;
            case 3: s += "="; break;
        }

        return Convert.FromBase64String(s);
    }

    /// <summary>
    /// Attempts to decode a signature string. Tries base64url first, then standard base64.
    /// Returns null if neither encoding works.
    /// </summary>
    private static byte[]? TryDecodeSignature(string signature)
    {
        // Try base64url first
        try
        {
            return Base64UrlDecode(signature);
        }
        catch (FormatException)
        {
            // Fall through to standard base64
        }

        // Try standard base64
        try
        {
            return Convert.FromBase64String(signature);
        }
        catch (FormatException)
        {
            return null;
        }
    }
}

/// <summary>
/// Represents a decoded JWT token with its constituent parts.
/// </summary>
public sealed class JwtToken
{
    /// <summary>The decoded JWT header.</summary>
    public required JwtHeader Header { get; init; }

    /// <summary>The raw JSON string of the payload.</summary>
    public required string PayloadJson { get; init; }

    /// <summary>The raw signature bytes.</summary>
    public required byte[] Signature { get; init; }

    /// <summary>The signing input (header.payload in base64url) used for signature verification.</summary>
    public required string SigningInput { get; init; }
}

/// <summary>
/// NATS JWT header. Algorithm is "ed25519-nkey" for NATS JWTs.
/// </summary>
public sealed class JwtHeader
{
    [System.Text.Json.Serialization.JsonPropertyName("alg")]
    public string? Algorithm { get; set; }

    [System.Text.Json.Serialization.JsonPropertyName("typ")]
    public string? Type { get; set; }
}
