feat(auth): finally better auth working as i wanted it to

frontend/src/lib/authClient.ts

@@ -1,11 +1,139 @@
 import { createAuthClient } from "better-auth/client";
 import { usernameClient } from "better-auth/client/plugins";
-export const authClient = createAuthClient({
-    baseURL: `${window.location.origin}/lst/api/auth`, // 👈 This is fine
-    callbacks: {
-        onUpdate: (session: any) => console.log("Session updated", session),
-        onSignIn: (session: any) => console.log("Signed in!", session),
-        onSignOut: () => console.log("Signed out!"),
+import { create } from "zustand";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
+import { useEffect } from "react";
+import { useRouter } from "@tanstack/react-router";
+
+// ---- TYPES ----
+export type Session = typeof authClient.$Infer.Session | null;
+
+// Zustand store type
+type SessionState = {
+    session: Session;
+    setSession: (session: Session) => void;
+    clearSession: () => void;
+};
+
+export type UserRoles = {
+    userRoleId: string;
+    userId: string;
+    module: string;
+    role: "systemAdmin" | "admin" | "manager" | "user" | "viewer";
+};
+
+type UserRoleState = {
+    userRoles: UserRoles[] | null;
+    fetchRoles: (userId: string) => Promise<void>;
+    clearRoles: () => void;
+};
+
+// ---- ZUSTAND STORE ----
+export const useAuth = create<SessionState>((set) => ({
+    session: null,
+    setSession: (session) => set({ session }),
+    clearSession: () => set({ session: null }),
+}));
+
+export const useUserRoles = create<UserRoleState>((set) => ({
+    userRoles: null,
+    fetchRoles: async (userId: string) => {
+        try {
+            const res = await fetch(`/api/${userId}/roles`, {
+                credentials: "include",
+            });
+            if (!res.ok) throw new Error("Failed to fetch roles");
+            const roles = await res.json();
+            set({ userRoles: roles });
+        } catch (err) {
+            console.error("Error fetching roles:", err);
+            set({ userRoles: null });
+        }
     },
+    clearRoles: () => set({ userRoles: null }),
+}));
+// ---- BETTER AUTH CLIENT ----
+export const authClient = createAuthClient({
+    baseURL: `${window.location.origin}/lst/api/auth`,
     plugins: [usernameClient()],
+    callbacks: {
+        callbacks: {
+            onUpdate: (res: any) => {
+                // res has strong type
+                // res.data is `Session | null`
+                useAuth.getState().setSession(res?.data ?? null);
+            },
+            onSignIn: (res: any) => {
+                console.log("Setting session to ", res?.data);
+                useAuth.getState().setSession(res?.data ?? null);
+            },
+            onSignOut: () => {
+                useAuth.getState().clearSession();
+            },
+        },
+    },
 });
+
+// ---- AUTH API HELPERS ----
+export async function signin(data: { username: string; password: string }) {
+    const res = await authClient.signIn.username(data);
+
+    if (res.error) throw res.error;
+    await authClient.getSession();
+    return res.data;
+}
+
+export const useLogout = () => {
+    const { clearSession } = useAuth();
+    const router = useRouter();
+    const logout = async () => {
+        await authClient.signOut();
+
+        router.invalidate();
+        router.clearCache();
+        clearSession();
+
+        window.location.reload();
+    };
+
+    return logout;
+};
+
+export async function getSession() {
+    const res = await authClient.getSession({
+        query: { disableCookieCache: true },
+    });
+    if (res.error) return null;
+    return res.data;
+}
+
+// ---- REACT QUERY INTEGRATION ----
+export function useSession() {
+    const { setSession, clearSession } = useAuth();
+    const qc = useQueryClient();
+
+    const query = useQuery({
+        queryKey: ["session"],
+        queryFn: getSession,
+        refetchInterval: 60_000,
+        refetchOnWindowFocus: true,
+    });
+
+    //console.log("Auth Check", query.data);
+    // react to data change
+    useEffect(() => {
+        if (query.data !== undefined) {
+            setSession(query.data);
+        }
+    }, [query.data, setSession]);
+
+    // react to error
+    useEffect(() => {
+        if (query.error) {
+            clearSession();
+            qc.removeQueries({ queryKey: ["session"] });
+        }
+    }, [query.error, qc, clearSession]);
+
+    return query;
+}