lst/frontend/src/lib/authClient.ts

import { useQuery, useQueryClient } from "@tanstack/react-query";
import { redirect, useNavigate, useRouter } from "@tanstack/react-router";
import { createAuthClient } from "better-auth/client";
import { usernameClient } from "better-auth/client/plugins";
import { useEffect } from "react";
import { create } from "zustand";
import { api } from "./axiosAPI";

// ---- 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:
		| "viewer"
		| "technician"
		| "supervisor"
		| "manager"
		| "admin"
		| "systemAdmin";
};

type UserRoleState = {
	userRoles: UserRoles[] | null;
	fetchRoles: () => 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 () => {
		try {
			const res = await api.get("/api/user/roles");
			const roles = res.data;
			set({ userRoles: roles.data });
		} catch (err) {
			console.error("Error fetching roles:", err);
			set({ userRoles: null });
		}
	},
	clearRoles: () => set({ userRoles: null }),
}));

export function userAccess(
	moduleName: string | null,
	roles: UserRoles["role"] | UserRoles["role"][],
): boolean {
	const { userRoles } = useUserRoles();

	if (!userRoles) return false;

	const roleArray = Array.isArray(roles) ? roles : [roles];

	return userRoles.some(
		(m) =>
			(moduleName ? m.module === moduleName : true) &&
			roleArray.includes(m.role),
	);
}

export async function checkUserAccess({
	allowedRoles,
	moduleName,
}: {
	allowedRoles: UserRoles["role"][];
	moduleName?: string;
	//location: { pathname: string; search: string };
}) {
	try {
		// fetch roles from your API (credentials required)
		const res = await api.get("/api/user/roles", { withCredentials: true });
		const roles = res.data.data as UserRoles[];

		const hasAccess = roles.some(
			(r) =>
				(moduleName ? r.module === moduleName : true) &&
				allowedRoles.includes(r.role),
		);

		if (!hasAccess) {
			throw redirect({
				to: "/",
				search: { from: location.pathname + location.search },
			});
		}

		// return roles so the route component can use them if needed
		return roles;
	} catch {
		throw redirect({
			to: "/login",
			search: { redirect: location.pathname + location.search },
		});
	}
}
// ---- BETTER AUTH CLIENT ----
export const authClient = createAuthClient({
	baseURL: `${window.location.origin}/lst/api/auth`,
	plugins: [usernameClient()],
	options: {
		autoPopup: false,
		requireAuth: false,
	},
	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 { clearRoles } = useUserRoles();
	const navigate = useNavigate();

	const router = useRouter();
	const logout = async () => {
		await authClient.signOut();

		router.invalidate();
		router.clearCache();
		clearSession();
		clearRoles();
		navigate({ to: "/" });
		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;
}