alrahma_web_client/src/auth/AuthProvider.tsx

import {
  createContext,
  useCallback,
  useContext,
  useMemo,
  useSyncExternalStore,
} from 'react'
import { ApiHttpError, setStoredToken, getStoredToken } from '../api/http'
import type { AuthUser } from '../api/types'
import { loginRequest } from '../api/session'

export type { AuthUser }

type AuthState = {
  token: string | null
  user: AuthUser | null
  selectedRole: string | null
}

const STORAGE_USER = 'alrahma_api_user'
const STORAGE_ROLE = 'alrahma_selected_role'

function enabledRoles(user: AuthUser | null): string[] {
  if (!user?.roles) return []
  return Object.keys(user.roles).filter((k) => user.roles[k])
}

function readUser(): AuthUser | null {
  try {
    const raw = localStorage.getItem(STORAGE_USER)
    if (!raw) return null
    return JSON.parse(raw) as AuthUser
  } catch {
    return null
  }
}

function writeUser(user: AuthUser | null): void {
  if (user) localStorage.setItem(STORAGE_USER, JSON.stringify(user))
  else localStorage.removeItem(STORAGE_USER)
}

function readSelectedRole(): string | null {
  const role = localStorage.getItem(STORAGE_ROLE)
  return role && role.trim() ? role : null
}

function writeSelectedRole(role: string | null): void {
  if (role) localStorage.setItem(STORAGE_ROLE, role)
  else localStorage.removeItem(STORAGE_ROLE)
}

function validSelectedRole(user: AuthUser | null, role: string | null): string | null {
  if (!role) return null
  const roles = enabledRoles(user)
  return roles.includes(role) ? role : null
}

let memory: AuthState = {
  token: getStoredToken(),
  user: readUser(),
  selectedRole: null,
}

memory.selectedRole = validSelectedRole(memory.user, readSelectedRole())

const listeners = new Set<() => void>()

function emit(): void {
  listeners.forEach((fn) => fn())
}

function setAuth(next: AuthState): void {
  memory = next
  emit()
}

function subscribe(listener: () => void): () => void {
  listeners.add(listener)
  return () => listeners.delete(listener)
}

function getSnapshot(): AuthState {
  return memory
}

export function hydrateAuthFromStorage(): void {
  const user = readUser()
  memory = {
    token: getStoredToken(),
    user,
    selectedRole: validSelectedRole(user, readSelectedRole()),
  }
  emit()
}

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const state = useSyncExternalStore(subscribe, getSnapshot, getSnapshot)

  const login = useCallback(async (email: string, password: string) => {
    let result: Awaited<ReturnType<typeof loginRequest>>
    try {
      result = await loginRequest(email, password)
    } catch (e) {
      const message =
        e instanceof ApiHttpError
          ? e.message
          : e instanceof Error
            ? e.message
            : 'Could not reach the server.'
      return { ok: false as const, message }
    }

    if (result.status === false) {
      return {
        ok: false as const,
        message: result.message ?? 'Login failed.',
      }
    }

    const roles = enabledRoles(result.user)
    const nextRole = roles.length === 1 ? roles[0] : null
    setStoredToken(result.token)
    writeUser(result.user)
    writeSelectedRole(nextRole)
    setAuth({ token: result.token, user: result.user, selectedRole: nextRole })
    return {
      ok: true as const,
      requiresRoleSelection: roles.length > 1,
      roles,
      selectedRole: nextRole,
    }
  }, [])

  const logout = useCallback(() => {
    setStoredToken(null)
    writeUser(null)
    writeSelectedRole(null)
    setAuth({ token: null, user: null, selectedRole: null })
  }, [])

  const setSession = useCallback((token: string, user: AuthUser) => {
    const roles = enabledRoles(user)
    const nextRole = roles.length === 1 ? roles[0] : null
    setStoredToken(token)
    writeUser(user)
    writeSelectedRole(nextRole)
    setAuth({ token, user, selectedRole: nextRole })
  }, [])

  const setSelectedRole = useCallback((role: string | null) => {
    const roles = enabledRoles(memory.user)
    const nextRole = role && roles.includes(role) ? role : null
    writeSelectedRole(nextRole)
    setAuth({ ...memory, selectedRole: nextRole })
  }, [])

  const value = useMemo(
    () => ({
      login,
      logout,
      setSession,
      setSelectedRole,
      token: state.token,
      user: state.user,
      selectedRole: state.selectedRole,
      roles: enabledRoles(state.user),
      isAuthenticated: Boolean(state.token && state.user),
    }),
    [state.token, state.user, state.selectedRole, login, logout, setSession, setSelectedRole],
  )

  return <AuthContext.Provider value={value}>{children}</AuthContext.Provider>
}

type AuthContextValue = {
  login: (
    email: string,
    password: string,
  ) => Promise<
    { ok: true; requiresRoleSelection: boolean; roles: string[]; selectedRole: string | null } | { ok: false; message: string }
  >
  logout: () => void
  setSession: (token: string, user: AuthUser) => void
  setSelectedRole: (role: string | null) => void
  token: string | null
  user: AuthUser | null
  selectedRole: string | null
  roles: string[]
  isAuthenticated: boolean
}

const AuthContext = createContext<AuthContextValue | null>(null)

export function useAuth(): AuthContextValue {
  const ctx = useContext(AuthContext)
  if (!ctx) throw new Error('useAuth must be used within AuthProvider')
  return ctx
}