QuantEngineByItz/src/dotnet/QuantEngine.Web/Client/Infrastructure/CustomAuthenticationStateProvider.cs

using System.Security.Claims;
using Microsoft.AspNetCore.Components.Authorization;
using Microsoft.JSInterop;
using QuantEngine.Web.Client.Services;

namespace QuantEngine.Web.Client.Infrastructure
{
    public class CustomAuthenticationStateProvider : AuthenticationStateProvider
    {
        private readonly LocalStorageService _localStorage;
        private readonly HttpClient _http;
        private readonly IJSRuntime _jsRuntime;
        private readonly ClaimsPrincipal _anonymous = new ClaimsPrincipal(new ClaimsIdentity());
        private const string TokenKey = "quant_admin_access_token";
        private const string UsernameKey = "quant_admin_username";
        private const string RoleKey = "quant_admin_role";
        private const string RememberUsernameKey = "quant_admin_remember_username";

        private AuthenticationState? _cachedState;

        public CustomAuthenticationStateProvider(LocalStorageService localStorage, HttpClient http, IJSRuntime jsRuntime)
        {
            _localStorage = localStorage;
            _http = http;
            _jsRuntime = jsRuntime;
        }

        public override async Task<AuthenticationState> GetAuthenticationStateAsync()
        {
            if (_cachedState != null && _cachedState.User.Identity?.IsAuthenticated == true)
            {
                Console.WriteLine("[Auth] Returning cached authentication state");
                return _cachedState;
            }

            try
            {
                Console.WriteLine("[Auth] GetAuthenticationStateAsync called");

                // Primary: Try to validate via /api/auth/me
                // This works with both cookies (automatic) and Bearer tokens
                try
                {
                    Console.WriteLine("[Auth] Attempting validation via /api/auth/me (cookie or Bearer)...");
                    // BaseAddress is always set to HostEnvironment.BaseAddress by DI.
                    // Never fall back to a hardcoded port — it breaks in production.
                    var meUrl = "api/auth/me";
                    var requestUri = _http.BaseAddress == null ? new Uri($"http://localhost:5265/{meUrl}") : new Uri(_http.BaseAddress, meUrl);
                    Console.WriteLine($"[Auth] /api/auth/me URL: {requestUri}");

                    var meResponse = await _http.GetAsync(requestUri);
                    Console.WriteLine($"[Auth] /api/auth/me status: {meResponse.StatusCode}");

                    if (meResponse.IsSuccessStatusCode)
                    {
                        var json = await meResponse.Content.ReadAsStringAsync();
                        Console.WriteLine($"[Auth] Response JSON: {json}");

                        var meData = System.Text.Json.JsonDocument.Parse(json).RootElement;
                        var authenticated = meData.TryGetProperty("authenticated", out var authProp) && authProp.GetBoolean();
                        var username = meData.TryGetProperty("username", out var userProp) ? userProp.GetString() : null;
                        var role = meData.TryGetProperty("role", out var roleProp) ? roleProp.GetString() : "Admin";

                        Console.WriteLine($"[Auth] Parsed: authenticated={authenticated}, username={username}, role={role}");

                        if (authenticated && !string.IsNullOrWhiteSpace(username))
                        {
                            Console.WriteLine($"[Auth] ✅ SUCCESS: Authenticated as {username}");
                            var identity = new ClaimsIdentity(new[]
                            {
                                new Claim(ClaimTypes.Name, username),
                                new Claim(ClaimTypes.Role, role ?? "Admin")
                            }, "QuantAdminAuth");

                            var state = new AuthenticationState(new ClaimsPrincipal(identity));
                            _cachedState = state;
                            return state;
                        }
                        else
                        {
                            Console.WriteLine($"[Auth] Parsing failed: authenticated={authenticated}, username={username}");
                        }
                    }
                    else
                    {
                        Console.WriteLine($"[Auth] /api/auth/me returned {meResponse.StatusCode}");
                        
                        if (IsLocalhost())
                        {
                            Console.WriteLine("[Auth] Dev SSR fallback: allowing admin authentication on 401");
                            return GetDevAdminState();
                        }
                    }
                }
                catch (Exception meEx)
                {
                    Console.WriteLine($"[Auth] /api/auth/me failed: {meEx.Message}");
                    Console.WriteLine($"[Auth] Exception: {meEx}");
                    
                    if (IsLocalhost())
                    {
                        Console.WriteLine("[Auth] Dev SSR fallback: allowing admin authentication on exception");
                        return GetDevAdminState();
                    }
                }

                // Fallback: Try to read from localStorage
                Console.WriteLine("[Auth] Fallback: checking localStorage...");
                try
                {
                    string token = await _jsRuntime.InvokeAsync<string>("localStorage.getItem", TokenKey);
                    string username = await _jsRuntime.InvokeAsync<string>("localStorage.getItem", UsernameKey);
                    string role = await _jsRuntime.InvokeAsync<string>("localStorage.getItem", RoleKey);

                    Console.WriteLine($"[Auth] localStorage: token={!string.IsNullOrWhiteSpace(token)}, username={username}");

                    if (!string.IsNullOrWhiteSpace(token) && !string.IsNullOrWhiteSpace(username))
                    {
                        var meUrl = "api/auth/me";
                        var requestUri = _http.BaseAddress == null ? new Uri($"http://localhost:5265/{meUrl}") : new Uri(_http.BaseAddress, meUrl);
                        var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
                        request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", token);
                        var response = await _http.SendAsync(request);

                        if (response.IsSuccessStatusCode)
                        {
                            Console.WriteLine($"[Auth] ✅ localStorage token validated: {username}");
                            var identity = new ClaimsIdentity(new[]
                            {
                                new Claim(ClaimTypes.Name, username),
                                new Claim(ClaimTypes.Role, role ?? "Admin")
                            }, "QuantAdminAuth");

                            var state = new AuthenticationState(new ClaimsPrincipal(identity));
                            _cachedState = state;
                            return state;
                        }
                    }
                }
                catch (Exception jsEx)
                {
                    Console.WriteLine($"[Auth] localStorage fallback failed: {jsEx.Message}");
                }

                Console.WriteLine("[Auth] ❌ Not authenticated");
            }
            catch (Exception ex)
            {
                Console.WriteLine($"[Auth] Unexpected error: {ex.Message}");
            }

            _cachedState = new AuthenticationState(_anonymous);
            return _cachedState;
        }

        public async Task MarkUserAsAuthenticatedAsync(string username, string accessToken, string role)
        {
            await MarkUserAsAuthenticatedAsync(username, accessToken, role, rememberUsername: true);
        }

        public async Task MarkUserAsAuthenticatedAsync(string username, string accessToken, string role, bool rememberUsername)
        {
            await _localStorage.SetAsync(TokenKey, accessToken);
            if (rememberUsername)
            {
                await _localStorage.SetAsync(UsernameKey, username);
                await _localStorage.SetAsync(RememberUsernameKey, true);
            }
            else
            {
                await _localStorage.DeleteAsync(UsernameKey);
                await _localStorage.SetAsync(RememberUsernameKey, false);
            }
            await _localStorage.SetAsync(RoleKey, role);

            var identity = new ClaimsIdentity(new[]
            {
                new Claim(ClaimTypes.Name, username),
                new Claim(ClaimTypes.Role, role)
            }, "QuantAdminAuth");

            var user = new ClaimsPrincipal(identity);
            var state = new AuthenticationState(user);
            _cachedState = state;
            NotifyAuthenticationStateChanged(Task.FromResult(state));
        }

        public async Task MarkUserAsLoggedOutAsync()
        {
            await _localStorage.DeleteAsync(TokenKey);
            await _localStorage.DeleteAsync(RoleKey);
            var rememberUsername = await _localStorage.GetAsync<bool>(RememberUsernameKey);
            if (!rememberUsername)
            {
                await _localStorage.DeleteAsync(UsernameKey);
            }
            _cachedState = new AuthenticationState(_anonymous);
            NotifyAuthenticationStateChanged(Task.FromResult(_cachedState));
        }

        public async Task LogoutFromServerAsync()
        {
            var token = await _localStorage.GetAsync<string>(TokenKey);
            if (!string.IsNullOrWhiteSpace(token))
            {
                try
                {
                    var request = new HttpRequestMessage(HttpMethod.Post, "api/auth/logout");
                    request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", token);
                    await _http.SendAsync(request);
                }
                catch
                {
                    // Best-effort server revocation; always clear local state.
                }
            }

            await MarkUserAsLoggedOutAsync();
        }

        public async Task<string?> GetRememberedUsernameAsync()
        {
            var rememberUsername = await _localStorage.GetAsync<bool>(RememberUsernameKey);
            if (!rememberUsername)
            {
                return null;
            }

            return await _localStorage.GetAsync<string>(UsernameKey);
        }

        private bool IsLocalhost()
        {
            return _http.BaseAddress == null || _http.BaseAddress.Host == "localhost" || _http.BaseAddress.Host == "127.0.0.1";
        }

        private AuthenticationState GetDevAdminState()
        {
            var identity = new ClaimsIdentity(new[]
            {
                new Claim(ClaimTypes.Name, "admin"),
                new Claim(ClaimTypes.Role, "Admin")
            }, "QuantAdminAuth");
            var state = new AuthenticationState(new ClaimsPrincipal(identity));
            _cachedState = state;
            return state;
        }
    }
}