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

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

namespace QuantEngine.Web.Client.Infrastructure
{
    public class CustomAuthenticationStateProvider : AuthenticationStateProvider
    {
        private readonly LocalStorageService _localStorage;
        private readonly HttpClient _http;
        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";

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

        public override async Task<AuthenticationState> GetAuthenticationStateAsync()
        {
            try
            {
                var token = await _localStorage.GetAsync<string>(TokenKey);
                var username = await _localStorage.GetAsync<string>(UsernameKey);
                var role = await _localStorage.GetAsync<string>(RoleKey) ?? "Admin";

                if (!string.IsNullOrWhiteSpace(token) && !string.IsNullOrWhiteSpace(username))
                {
                    try
                    {
                        var request = new HttpRequestMessage(HttpMethod.Get, "api/auth/me");
                        request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", token);
                        var response = await _http.SendAsync(request);

                        if (!response.IsSuccessStatusCode)
                        {
                            Console.WriteLine($"[Auth] /api/auth/me failed: {response.StatusCode}");
                            await MarkUserAsLoggedOutAsync();
                            return new AuthenticationState(_anonymous);
                        }

                        Console.WriteLine($"[Auth] User authenticated: {username}");
                        var identity = new ClaimsIdentity(new[]
                        {
                            new Claim(ClaimTypes.Name, username),
                            new Claim(ClaimTypes.Role, role)
                        }, "QuantAdminAuth");

                        var user = new ClaimsPrincipal(identity);
                        return new AuthenticationState(user);
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine($"[Auth] Error during /api/auth/me call: {ex.Message}");
                        // Fall through to anonymous
                    }
                }
                else
                {
                    Console.WriteLine("[Auth] No token or username found in localStorage");
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"[Auth] Error accessing localStorage: {ex.Message}");
                // Return anonymous if localStorage isn't ready
            }

            return new AuthenticationState(_anonymous);
        }

        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);
            NotifyAuthenticationStateChanged(Task.FromResult(new AuthenticationState(user)));
        }

        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);
            }
            NotifyAuthenticationStateChanged(Task.FromResult(new AuthenticationState(_anonymous)));
        }

        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);
        }
    }
}