hooks/useUser.ts

/* eslint-disable @typescript-eslint/no-explicit-any */
"use client";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { useCookie } from "react-use";
import { useRouter } from "next/navigation";

import { User } from "@/types";
import { api } from "@/lib/api";
import { toast } from "sonner";
import { 
  storeAuthDataFallback, 
  getAuthDataFallback, 
  clearAuthDataFallback,
  isInIframe 
} from "@/lib/iframe-storage";


export const useUser = (initialData?: {
  user: User | null;
  errCode: number | null;
}) => {
  const client = useQueryClient();
  const router = useRouter();
  const [currentRoute, setCurrentRoute, removeCurrentRoute] = useCookie("deepsite-currentRoute");

  const { data: { user, errCode } = { user: null, errCode: null }, isLoading } =
    useQuery({
      queryKey: ["user.me"],
      queryFn: async () => {
        // Check for fallback data if no initial data provided and we're in iframe
        if (!initialData && isInIframe()) {
          const fallbackData = getAuthDataFallback();
          if (fallbackData.user && fallbackData.token) {
            return { user: fallbackData.user, errCode: null };
          }
        }
        return { user: initialData?.user || null, errCode: initialData?.errCode || null };
      },
      refetchOnWindowFocus: false,
      refetchOnReconnect: false,
      refetchOnMount: false,
      retry: false,
      initialData: initialData
        ? { user: initialData?.user, errCode: initialData?.errCode }
        : undefined,
      enabled: false,
    });

  const { data: loadingAuth } = useQuery({
    queryKey: ["loadingAuth"],
    queryFn: async () => false,
    refetchOnWindowFocus: false,
    refetchOnReconnect: false,
    refetchOnMount: false,
  });
  const setLoadingAuth = (value: boolean) => {
    client.setQueryData(["setLoadingAuth"], value);
  };

  const openLoginWindow = async () => {
    setCurrentRoute(window.location.pathname);
    return router.push("/auth");
  };

  const loginFromCode = async (code: string) => {
    setLoadingAuth(true);
    if (loadingAuth) return;
    await api
      .post("/auth", { code })
      .then(async (res: any) => {
        if (res.data) {
          // Cookie is now set server-side with proper iframe attributes
          // Also store fallback data for iframe contexts
          if (res.data.useLocalStorageFallback) {
            storeAuthDataFallback(res.data.access_token, res.data.user);
          }
          
          client.setQueryData(["user.me"], {
            user: res.data.user,
            errCode: null,
          });
          if (currentRoute) {
            router.push(currentRoute);
            removeCurrentRoute();
          } else {
            router.push("/projects");
          }
          toast.success("Login successful");
        }
      })
      .catch((err: any) => {
        toast.error(err?.data?.message ?? err.message ?? "An error occurred");
      })
      .finally(() => {
        setLoadingAuth(false);
      });
  };

  const logout = async () => {
    try {
      // Call server endpoint to clear the HTTP-only cookie
      await api.post("/auth/logout");
      // Clear fallback storage
      clearAuthDataFallback();
      removeCurrentRoute();
      client.setQueryData(["user.me"], { user: null, errCode: null });
      router.push("/");
      toast.success("Logout successful");
      client.invalidateQueries({ queryKey: ["user.me"] });
      window.location.reload();
    } catch (error) {
      console.error("Logout error:", error);
      // Even if server call fails, clear client state
      clearAuthDataFallback();
      removeCurrentRoute();
      client.setQueryData(["user.me"], { user: null, errCode: null });
      router.push("/");
      toast.success("Logout successful");
      window.location.reload();
    }
  };

  return {
    user,
    errCode,
    loading: isLoading || loadingAuth,
    openLoginWindow,
    loginFromCode,
    logout,
  };
};