using Microsoft.Azure.Functions.Worker;
using Microsoft.Azure.Functions.Worker.Http;
using Microsoft.Azure.Functions.Worker.Middleware;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Logging;
using Microsoft.IdentityModel.Protocols;
using Microsoft.IdentityModel.Protocols.OpenIdConnect;
using Microsoft.IdentityModel.Tokens;
using System.IdentityModel.Tokens.Jwt;
using System.Net;
using System.Security.Claims;

namespace AutopilotMonitor.Functions.Middleware;

/// <summary>
/// Middleware to manually validate JWT tokens and populate ClaimsPrincipal
/// Required for Azure Functions .NET 8 Isolated Worker (Microsoft.Identity.Web doesn't work automatically)
/// </summary>
public class AuthenticationMiddleware : IFunctionsWorkerMiddleware
{
    private readonly ILogger<AuthenticationMiddleware> _logger;
    private readonly IConfiguration _configuration;
    private readonly JwtSecurityTokenHandler _tokenHandler;

    // Cache configuration managers per tenant to avoid repeated OIDC metadata fetches
    private readonly Dictionary<string, IConfigurationManager<OpenIdConnectConfiguration>> _configManagerCache;
    private readonly SemaphoreSlim _cacheLock = new SemaphoreSlim(1, 1);

    // Track if we're currently initializing a tenant to prevent race conditions
    private readonly HashSet<string> _initializingTenants = new HashSet<string>();
    private readonly object _initLock = new object();

    public AuthenticationMiddleware(
        ILogger<AuthenticationMiddleware> logger,
        IConfiguration configuration)
    {
        _logger = logger;
        _configuration = configuration;
        _tokenHandler = new JwtSecurityTokenHandler();
        _configManagerCache = new Dictionary<string, IConfigurationManager<OpenIdConnectConfiguration>>();

        // Disable PII logging in production for security
        Microsoft.IdentityModel.Logging.IdentityModelEventSource.ShowPII = false;
    }

    public async Task Invoke(FunctionContext context, FunctionExecutionDelegate next)
    {
        var httpContext = context.GetHttpContext();
        if (httpContext != null)
        {
            // Extract Authorization header
            var authHeader = httpContext.Request.Headers.Authorization.ToString();
            if (!string.IsNullOrEmpty(authHeader) && authHeader.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase))
            {
                var token = authHeader.Substring("Bearer ".Length).Trim();

                try
                {
                    // First, decode the token to get the tenant ID (without validation)
                    var jwtToken = _tokenHandler.ReadJwtToken(token);
                    var tenantId = jwtToken.Claims.FirstOrDefault(c => c.Type == "tid")?.Value;
                    var issuer = jwtToken.Issuer;

                    // Determine which endpoint to use based on the issuer (v1.0 vs v2.0)
                    var isV1Token = issuer.Contains("sts.windows.net");
                    var tenantSpecificAuthority = tenantId != null
                        ? (isV1Token
                            ? $"https://login.microsoftonline.com/{tenantId}"  // v1.0
                            : $"https://login.microsoftonline.com/{tenantId}/v2.0")  // v2.0
                        : "https://login.microsoftonline.com/common/v2.0";

                    // Get or create cached configuration manager for this tenant
                    IConfigurationManager<OpenIdConnectConfiguration>? tenantConfigManager = null;
                    await _cacheLock.WaitAsync();
                    try
                    {
                        if (!_configManagerCache.TryGetValue(tenantSpecificAuthority, out tenantConfigManager))
                        {
                            var tenantMetadataAddress = $"{tenantSpecificAuthority}/.well-known/openid-configuration";
                            tenantConfigManager = new ConfigurationManager<OpenIdConnectConfiguration>(
                                tenantMetadataAddress,
                                new OpenIdConnectConfigurationRetriever(),
                                new HttpDocumentRetriever() { RequireHttps = true })
                            {
                                // Cache configuration for 24 hours (signing keys rarely change)
                                AutomaticRefreshInterval = TimeSpan.FromHours(24),
                                RefreshInterval = TimeSpan.FromHours(24)
                            };
                            _configManagerCache[tenantSpecificAuthority] = tenantConfigManager;
                            _logger.LogInformation($"[Auth Middleware] Created new config manager for {tenantSpecificAuthority}");
                        }
                    }
                    finally
                    {
                        _cacheLock.Release();
                    }

                    if (tenantConfigManager == null)
                    {
                        throw new InvalidOperationException("Failed to get or create configuration manager");
                    }

                    var openIdConfig = await tenantConfigManager.GetConfigurationAsync(CancellationToken.None);

                    // Set up token validation parameters with OIDC signing keys
                    var validationParameters = new TokenValidationParameters
                    {
                        // For multi-tenant, validate issuer format but accept any tenant
                        ValidateIssuer = true,
                        IssuerValidator = (issuer, token, parameters) =>
                        {
                            // Accept issuers from any Azure AD tenant
                            if (issuer.StartsWith("https://login.microsoftonline.com/") ||
                                issuer.StartsWith("https://sts.windows.net/"))
                            {
                                return issuer;
                            }
                            throw new SecurityTokenInvalidIssuerException($"Invalid issuer: {issuer}");
                        },
                        ValidateAudience = true,
                        ValidAudiences = new[]
                        {
                            _configuration["AzureAd:ClientId"],
                            $"api://{_configuration["AzureAd:ClientId"]}"
                        },
                        ValidateLifetime = true,

                        // Signature validation enabled
                        // Token is requested for backend API (api://<clientId>/access_as_user)
                        ValidateIssuerSigningKey = true,
                        RequireSignedTokens = true,

                        // Use key resolver for dynamic key resolution
                        IssuerSigningKeyResolver = (token, securityToken, kid, validationParameters) =>
                        {
                            var matchedKeys = openIdConfig.SigningKeys.Where(k => k.KeyId == kid).ToList();
                            return matchedKeys.Any() ? matchedKeys : openIdConfig.SigningKeys;
                        },

                        // Set algorithm validators to ensure we only accept RS256
                        ValidAlgorithms = new[] { "RS256" }
                    };

                    // Validate the token using JwtSecurityTokenHandler (synchronous API)
                    try
                    {
                        var principal = _tokenHandler.ValidateToken(token, validationParameters, out var validatedToken);

                        // Get user identifier for logging (same logic as TenantHelper.GetUserIdentifier)
                        var userIdentifier = principal.FindFirst("upn")?.Value ??
                                           principal.FindFirst(ClaimTypes.Upn)?.Value ??
                                           principal.FindFirst(ClaimTypes.Email)?.Value ??
                                           principal.FindFirst("preferred_username")?.Value ??
                                           principal.FindFirst(ClaimTypes.Name)?.Value ??
                                           principal.FindFirst("name")?.Value ??
                                           "Unknown";

                        _logger.LogInformation($"[Auth Middleware] ✅ Authenticated: {userIdentifier}");

                        // Set the principal on both the HTTP context AND the FunctionContext
                        // This is critical for Azure Functions Isolated Worker (.NET 8)
                        httpContext.User = principal;
                        context.Items["ClaimsPrincipal"] = principal;
                    }
                    catch (SecurityTokenException ex)
                    {
                        _logger.LogWarning($"[Auth Middleware] ❌ Token validation failed: {ex.Message}");
                        throw;
                    }
                }
                catch (SecurityTokenValidationException ex)
                {
                    _logger.LogWarning($"[Auth Middleware] ❌ Token validation failed: {ex.Message}");
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "[Auth Middleware] ❌ Error validating token");
                }
            }
            else
            {
                _logger.LogDebug("[Auth Middleware] No Bearer token found in Authorization header");
            }
        }

        await next(context);
    }
}