using System.IdentityModel.Tokens.Jwt;
using System.Net;
using System.Security.Claims;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Text.Json.Serialization;
using Microsoft.AspNetCore.Authentication.JwtBearer;
using Microsoft.IdentityModel.Tokens;

namespace NpgsqlRestClient;

// AOT-compatible JSON serialization context
[JsonSerializable(typeof(JwtErrorResponse))]
[JsonSerializable(typeof(JwtTokenResponse))]
internal partial class JwtJsonContext : JsonSerializerContext
{
}

internal class JwtErrorResponse
{
    [JsonPropertyName("error")]
    public string Error { get; set; } = null!;
}

internal class JwtTokenResponse
{
    [JsonPropertyName("accessToken")]
    public string AccessToken { get; set; } = null!;

    [JsonPropertyName("refreshToken")]
    public string RefreshToken { get; set; } = null!;

    [JsonPropertyName("tokenType")]
    public string TokenType { get; set; } = "Bearer";

    [JsonPropertyName("expiresIn")]
    public int ExpiresIn { get; set; }

    [JsonPropertyName("refreshExpiresIn")]
    public int RefreshExpiresIn { get; set; }
}

public class JwtTokenConfig
{
    public string Scheme { get; set; } = JwtBearerDefaults.AuthenticationScheme;
    public string Secret { get; set; } = null!;
    public string? Issuer { get; set; }
    public string? Audience { get; set; }
    public int ExpireMinutes { get; set; } = 60;
    public int RefreshExpireDays { get; set; } = 7;
    public bool ValidateIssuer { get; set; }
    public bool ValidateAudience { get; set; }
    public bool ValidateLifetime { get; set; } = true;
    public bool ValidateIssuerSigningKey { get; set; } = true;
    public TimeSpan ClockSkew { get; set; } = TimeSpan.FromMinutes(5);
    public string? RefreshPath { get; set; }

    public SymmetricSecurityKey GetSecurityKey() => new(Encoding.UTF8.GetBytes(Secret));

    public SigningCredentials GetSigningCredentials() => new(GetSecurityKey(), SecurityAlgorithms.HmacSha256);

    public TokenValidationParameters GetTokenValidationParameters() => new()
    {
        ValidateIssuer = ValidateIssuer,
        ValidateAudience = ValidateAudience,
        ValidateLifetime = ValidateLifetime,
        ValidateIssuerSigningKey = ValidateIssuerSigningKey,
        ValidIssuer = Issuer,
        ValidAudience = Audience,
        IssuerSigningKey = GetSecurityKey(),
        ClockSkew = ClockSkew
    };
}

public class JwtTokenGenerator
{
    private readonly JwtTokenConfig _config;

    public JwtTokenGenerator(JwtTokenConfig config)
    {
        _config = config;
    }

    public (string accessToken, string refreshToken, DateTime accessExpires, DateTime refreshExpires) GenerateTokens(ClaimsPrincipal principal)
    {
        var accessExpires = DateTime.UtcNow.AddMinutes(_config.ExpireMinutes);
        var refreshExpires = DateTime.UtcNow.AddDays(_config.RefreshExpireDays);

        var accessToken = GenerateAccessToken(principal.Claims, accessExpires);
        var refreshToken = GenerateRefreshToken(principal.Claims, refreshExpires);

        return (accessToken, refreshToken, accessExpires, refreshExpires);
    }

    public string GenerateAccessToken(IEnumerable<Claim> claims, DateTime expires)
    {
        var tokenDescriptor = new SecurityTokenDescriptor
        {
            Subject = new ClaimsIdentity(claims),
            Expires = expires,
            Issuer = _config.Issuer,
            Audience = _config.Audience,
            SigningCredentials = _config.GetSigningCredentials()
        };

        var tokenHandler = new JwtSecurityTokenHandler();
        var token = tokenHandler.CreateToken(tokenDescriptor);
        return tokenHandler.WriteToken(token);
    }

    public string GenerateRefreshToken(IEnumerable<Claim> claims, DateTime expires)
    {
        // Add a claim to identify this as a refresh token
        var refreshClaims = claims.ToList();
        refreshClaims.Add(new Claim("token_type", "refresh"));

        var tokenDescriptor = new SecurityTokenDescriptor
        {
            Subject = new ClaimsIdentity(refreshClaims),
            Expires = expires,
            Issuer = _config.Issuer,
            Audience = _config.Audience,
            SigningCredentials = _config.GetSigningCredentials()
        };

        var tokenHandler = new JwtSecurityTokenHandler();
        var token = tokenHandler.CreateToken(tokenDescriptor);
        return tokenHandler.WriteToken(token);
    }

    public ClaimsPrincipal? ValidateRefreshToken(string refreshToken)
    {
        var tokenHandler = new JwtSecurityTokenHandler();
        var validationParameters = _config.GetTokenValidationParameters();

        // For refresh tokens, we want to validate even if the access token would be expired
        // but the refresh token itself should not be expired
        validationParameters.ValidateLifetime = true;

        try
        {
            var principal = tokenHandler.ValidateToken(refreshToken, validationParameters, out var validatedToken);

            // Verify this is a refresh token
            var tokenTypeClaim = principal.FindFirst("token_type");
            if (tokenTypeClaim?.Value != "refresh")
            {
                return null;
            }

            return principal;
        }
        catch
        {
            return null;
        }
    }
}

public class JwtRefreshAuth
{
    public JwtRefreshAuth(JwtTokenConfig? jwtConfig, WebApplication app, ILogger? logger)
    {
        if (jwtConfig is null ||
            string.IsNullOrEmpty(jwtConfig.RefreshPath) ||
            string.IsNullOrEmpty(jwtConfig.Secret))
        {
            return;
        }

        var refreshPath = jwtConfig.RefreshPath;
        var tokenGenerator = new JwtTokenGenerator(jwtConfig);

        app.Use(async (context, next) =>
        {
            if (!context.Request.Path.Equals(refreshPath, StringComparison.OrdinalIgnoreCase))
            {
                await next(context);
                return;
            }

            if (!string.Equals(context.Request.Method, "POST", StringComparison.OrdinalIgnoreCase))
            {
                await next(context);
                return;
            }

            string refreshToken;
            try
            {
                using var reader = new StreamReader(context.Request.Body);
                var body = await reader.ReadToEndAsync();
                var node = JsonNode.Parse(string.IsNullOrWhiteSpace(body) ? "{}" : body);
                refreshToken = node!["refreshToken"]?.ToString() ?? throw new ArgumentException("refreshToken is null");
            }
            catch (Exception ex)
            {
                logger?.LogError(ex, "Failed to read refresh token from request body.");
                context.Response.StatusCode = (int)HttpStatusCode.BadRequest;
                context.Response.ContentType = "application/json";
                await context.Response.WriteAsync(JsonSerializer.Serialize(
                    new JwtErrorResponse { Error = "Invalid request body" },
                    JwtJsonContext.Default.JwtErrorResponse));
                return;
            }

            var principal = tokenGenerator.ValidateRefreshToken(refreshToken);
            if (principal is null)
            {
                logger?.LogWarning("Invalid or expired refresh token.");
                context.Response.StatusCode = (int)HttpStatusCode.Unauthorized;
                context.Response.ContentType = "application/json";
                await context.Response.WriteAsync(JsonSerializer.Serialize(
                    new JwtErrorResponse { Error = "Invalid or expired refresh token" },
                    JwtJsonContext.Default.JwtErrorResponse));
                return;
            }

            // Filter out the token_type claim when creating new tokens
            var claims = principal.Claims.Where(c => c.Type != "token_type").ToList();
            var newPrincipal = new ClaimsPrincipal(new ClaimsIdentity(claims, jwtConfig.Scheme));

            var (accessToken, newRefreshToken, accessExpires, refreshExpires) = tokenGenerator.GenerateTokens(newPrincipal);

            context.Response.StatusCode = (int)HttpStatusCode.OK;
            context.Response.ContentType = "application/json";

            var response = new JwtTokenResponse
            {
                AccessToken = accessToken,
                RefreshToken = newRefreshToken,
                TokenType = "Bearer",
                ExpiresIn = (int)(accessExpires - DateTime.UtcNow).TotalSeconds,
                RefreshExpiresIn = (int)(refreshExpires - DateTime.UtcNow).TotalSeconds
            };

            await context.Response.WriteAsync(JsonSerializer.Serialize(response, JwtJsonContext.Default.JwtTokenResponse));
        });

        logger?.LogDebug("JWT refresh endpoint registered at {RefreshPath}", refreshPath);
    }
}

/// <summary>
/// Static class that provides JWT token generation for login endpoints.
/// This is used by the NpgsqlRest login handler when JWT authentication scheme is specified.
/// </summary>
public static class JwtLoginHandler
{
    private static JwtTokenConfig? _config;
    private static JwtTokenGenerator? _tokenGenerator;

    /// <summary>
    /// Initializes the JWT login handler with the given configuration.
    /// Must be called before any login attempts with JWT scheme.
    /// </summary>
    public static void Initialize(JwtTokenConfig config)
    {
        _config = config;
        _tokenGenerator = new JwtTokenGenerator(config);
    }

    /// <summary>
    /// Gets whether JWT login handling is configured.
    /// </summary>
    public static bool IsConfigured => _config is not null && _tokenGenerator is not null;

    /// <summary>
    /// Gets the configured JWT scheme name, or null if not configured.
    /// </summary>
    public static string? Scheme => _config?.Scheme;

    /// <summary>
    /// Generates JWT tokens for the given claims principal and writes them to the response.
    /// Returns true if handled, false if JWT is not configured.
    /// </summary>
    public static async Task<bool> HandleLoginAsync(HttpContext context, ClaimsPrincipal principal)
    {
        if (_config is null || _tokenGenerator is null)
        {
            return false;
        }

        var (accessToken, refreshToken, accessExpires, refreshExpires) = _tokenGenerator.GenerateTokens(principal);

        context.Response.StatusCode = (int)HttpStatusCode.OK;
        context.Response.ContentType = "application/json";

        var response = new JwtTokenResponse
        {
            AccessToken = accessToken,
            RefreshToken = refreshToken,
            TokenType = "Bearer",
            ExpiresIn = (int)(accessExpires - DateTime.UtcNow).TotalSeconds,
            RefreshExpiresIn = (int)(refreshExpires - DateTime.UtcNow).TotalSeconds
        };

        await context.Response.WriteAsync(JsonSerializer.Serialize(response, JwtJsonContext.Default.JwtTokenResponse));
        return true;
    }
}