aurak/web/src/contexts/AuthContext.tsx

import React, { createContext, useContext, useState, useEffect } from 'react';

export type UserRole = 'SUPER_ADMIN' | 'TENANT_ADMIN' | 'USER';

export interface User {
    id: string;
    username: string;
    role: UserRole;
    tenantId?: string;
    // Legacy support
    email?: string;
    tenant_name?: string;
    isNotebookEnabled?: boolean;
    displayName?: string;
}

export interface TenantMembership {
    id: string;
    tenantId: string;
    role: UserRole;
    tenant: {
        id: string;
        name: string;
        domain?: string;
        parentId?: string | null;
    };
    features?: {
        isNotebookEnabled: boolean;
    };
}

interface AuthContextType {
    user: User | null;
    apiKey: string;
    availableTenants: TenantMembership[];
    activeTenant: TenantMembership | null;
    login: (key: string, userData: Partial<User> & { role?: string }) => void;
    logout: () => void;
    switchTenant: (tenantId: string) => void;
    isLoading: boolean;
}

const AuthContext = createContext<AuthContextType | undefined>(undefined);

export function AuthProvider({ children }: { children: React.ReactNode }) {
    const [user, setUser] = useState<User | null>(null);
    const [apiKey, setApiKey] = useState<string>(localStorage.getItem('kb_api_key') || '');
    const [availableTenants, setAvailableTenants] = useState<TenantMembership[]>([]);
    const [activeTenant, setActiveTenant] = useState<TenantMembership | null>(null);
    const [isLoading, setIsLoading] = useState(true);

    const fetchTenants = async (key: string, currentTenantId?: string) => {
        try {
            const res = await fetch('/api/users/tenants', {
                headers: { 'x-api-key': key }
            });
            if (res.ok) {
                const tenants: TenantMembership[] = await res.json();
                console.log('[AuthContext] Fetched tenants:', tenants);
                const filteredTenants = tenants.filter(t => t.tenant?.name !== 'Default');
                setAvailableTenants(filteredTenants);

                const savedTenantId = localStorage.getItem('kb_active_tenant_id') || currentTenantId;
                const active = filteredTenants.find(t => t.tenantId === savedTenantId) || filteredTenants[0] || null;
                setActiveTenant(active);
                if (active) {
                    localStorage.setItem('kb_active_tenant_id', active.tenantId);
                }
            }
        } catch (e) {
            console.error('Failed to fetch tenants', e);
        }
    };

    // On mount, restore session
    useEffect(() => {
        const restoreSession = async () => {
            if (!apiKey) {
                setIsLoading(false);
                return;
            }
            try {
                const res = await fetch('/api/users/me', {
                    headers: { 'x-api-key': apiKey }
                });
                if (res.ok) {
                    const data = await res.json();
                    console.log('[AuthContext] Restored user:', data);
                    setUser({
                        id: data.id,
                        username: data.username,
                        role: (data.role as UserRole) ?? 'USER',
                        tenantId: data.tenantId,
                        tenant_name: data.tenantName,
                        isNotebookEnabled: data.isNotebookEnabled ?? true,
                        displayName: data.displayName,
                    });
                    await fetchTenants(apiKey, data.tenantId);
                } else {
                    localStorage.removeItem('kb_api_key');
                    localStorage.removeItem('kb_active_tenant_id');
                    setApiKey('');
                    setUser(null);
                }
            } catch (e) {
                console.error('Failed to restore session', e);
            } finally {
                setIsLoading(false);
            }
        };

        restoreSession();
    }, [apiKey]);

    const login = (key: string, userData: Partial<User> & { role?: string }) => {
        localStorage.setItem('kb_api_key', key);
        setApiKey(key);
        setUser({
            id: userData.id ?? '',
            username: userData.username ?? '',
            role: (userData.role as UserRole) ?? 'USER',
            tenantId: userData.tenantId,
            tenant_name: (userData as any).tenantName || userData.tenant_name,
            isNotebookEnabled: userData.isNotebookEnabled ?? true,
            displayName: userData.displayName,
        });
        fetchTenants(key, userData.tenantId);
    };

    const logout = () => {
        localStorage.removeItem('kb_api_key');
        localStorage.removeItem('kb_active_tenant_id');
        setApiKey('');
        setUser(null);
        setAvailableTenants([]);
        setActiveTenant(null);
    };

    const switchTenant = (tenantId: string) => {
        const target = availableTenants.find(t => t.tenantId === tenantId);
        if (target) {
            setActiveTenant(target);
            localStorage.setItem('kb_active_tenant_id', tenantId);
            // Optionally reload page to reset all states, or just let components re-render
            window.location.reload();
        }
    };

    return (
        <AuthContext.Provider value={{ user, apiKey, availableTenants, activeTenant, login, logout, switchTenant, isLoading }}>
            {children}
        </AuthContext.Provider>
    );
}

export function useAuth() {
    const context = useContext(AuthContext);
    if (context === undefined) {
        throw new Error('useAuth must be used within an AuthProvider');
    }
    return context;
}