shadow-gunbuilder-ai-proto/context/AuthContext.tsx

"use client";

import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useState,
  ReactNode,
} from "react";

const API_BASE_URL =
  process.env.NEXT_PUBLIC_API_BASE_URL ?? "";

type AuthUser = {
  uuid: string;
  email: string;
  username?: string | null;
  passwordSetAt?: string | null;
  displayName?: string | null;
  role: string;
} | null;

export type RegisterParams = {
  email: string;
  password: string;
  displayName?: string;
  acceptedTos: boolean;
  tosVersion: string;
};

// 1) update the type
type AuthContextValue = {
  user: AuthUser;
  token: string | null;
  loading: boolean;
  login: (params: { email: string; password: string }) => Promise<void>;
  register: (params: RegisterParams) => Promise<void>;
  logout: () => void;
  getAuthHeaders: () => HeadersInit;
  setSession: (token: string, user: NonNullable<AuthUser>) => Promise<void>;
  refreshMeAndSession: () => Promise<void>;

};

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

const TOKEN_KEY = "ballistic_auth_token";
const USER_KEY = "ballistic_auth_user";

export function AuthProvider({ children }: { children: ReactNode }) {
  const [token, setToken] = useState<string | null>(null);
  const [user, setUser] = useState<AuthUser>(null);
  const [loading, setLoading] = useState<boolean>(true);

  /**
   * Persist auth to localStorage
   */
  const persistAuth = useCallback(
    (nextToken: string | null, nextUser: AuthUser) => {
      if (typeof window === "undefined") return;

      if (nextToken) {
        window.localStorage.setItem(TOKEN_KEY, nextToken);
      } else {
        window.localStorage.removeItem(TOKEN_KEY);
      }

      if (nextUser) {
        window.localStorage.setItem(USER_KEY, JSON.stringify(nextUser));
      } else {
        window.localStorage.removeItem(USER_KEY);
      }
    },
    []
  );

  /**
   * ✅ Hydrate from localStorage ONCE
   * IMPORTANT: loading only flips false AFTER user/token are set (if present)
   */
  useEffect(() => {
    if (typeof window === "undefined") return;

    const storedToken = window.localStorage.getItem(TOKEN_KEY);
    const storedUser = window.localStorage.getItem(USER_KEY);

    if (storedToken && storedUser) {
      try {
        const parsedUser = JSON.parse(storedUser);

        setToken(storedToken);
        setUser(parsedUser);

        // ✅ Re-sync cookie so middleware can allow /admin on hard refresh / direct nav
        fetch("/api/auth/session", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            token: storedToken,
            role: parsedUser?.role,
          }),
        }).catch(() => {});
      } catch {
        window.localStorage.removeItem(TOKEN_KEY);
        window.localStorage.removeItem(USER_KEY);
      }
    }

    setLoading(false);
  }, []);

  /**
   * Used by login/register/magic-link
   */
  const setSession = useCallback(
    async (nextToken: string, nextUser: NonNullable<AuthUser>) => {
      setToken(nextToken);
      setUser(nextUser);
      persistAuth(nextToken, nextUser);

      // ✅ await so middleware sees cookie before navigation
      try {
        await fetch("/api/auth/session", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            token: nextToken,
            role: nextUser.role,
          }),
        });
      } catch {
        // ignore
      }
    },
    [persistAuth]
  );

  const refreshMeAndSession = useCallback(async () => {
    if (!token) return;
  
    const res = await fetch(`${API_BASE_URL}/api/users/me`, {
      headers: { Authorization: `Bearer ${token}` },
      cache: "no-store",
    });
  
    if (!res.ok) {
      const text = await res.text().catch(() => res.statusText);
      throw new Error(text || "Failed to refresh session");
    }
  
    const me = await res.json();
  
    const nextUser = {
      uuid: me.uuid,
      email: me.email,
      username: me.username || null,
      displayName: me.displayName || null,
      role: me.role || "USER",
      passwordSetAt: me.passwordSetAt || null,
    } as NonNullable<AuthUser>;
  
    await setSession(token, nextUser);
  }, [token, setSession]);

  const login = useCallback(
    async ({ email, password }: { email: string; password: string }) => {
      setLoading(true);
      try {
        const res = await fetch(`${API_BASE_URL}/api/auth/login`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ email, password }),
        });

        if (!res.ok) {
          const text = await res.text().catch(() => res.statusText);
          throw new Error(text || "Login failed");
        }

        const data = await res.json();

        const nextToken: string = data.token ?? data.accessToken;
        const nextUser: AuthUser =
          data.user ??
          ({
            uuid: data.uuid,
            email: data.email ?? email,
            displayName: data.displayName ?? null,
            passwordSetAt: data.passwordSetAt ?? null, 
            role: data.role ?? "USER",
          } as AuthUser);

        await setSession(nextToken, nextUser as NonNullable<AuthUser>);
      } finally {
        setLoading(false);
      }
    },
    [setSession]
  );

  const register = useCallback(
    async ({
      email,
      password,
      displayName,
      acceptedTos,
      tosVersion,
    }: RegisterParams) => {
      setLoading(true);
      try {
        const res = await fetch(`${API_BASE_URL}/api/auth/register`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            email,
            password,
            displayName,
            acceptedTos,
            tosVersion,
          }),
        });

        if (!res.ok) {
          const text = await res.text().catch(() => res.statusText);
          throw new Error(text || "Registration failed");
        }

        const data = await res.json();

        const nextToken: string = data.token ?? data.accessToken;
        const nextUser: AuthUser =
          data.user ??
          ({
            uuid: data.uuid,
            email: data.email ?? email,
            displayName: data.displayName ?? displayName ?? null,
            passwordSetAt: data.passwordSetAt ?? null,
            role: data.role ?? "USER",
          } as AuthUser);

        await setSession(nextToken, nextUser as NonNullable<AuthUser>);
      } finally {
        setLoading(false);
      }
    },
    [setSession]
  );

  const logout = useCallback(() => {
    setToken(null);
    setUser(null);
    persistAuth(null, null);

    // Clear server session cookies
    fetch("/api/auth/session", { method: "DELETE" }).catch(() => {});
  }, [persistAuth]);

  const getAuthHeaders = useCallback((): HeadersInit => {
    return token ? { Authorization: `Bearer ${token}` } : {};
  }, [token]);

  const value: AuthContextValue = {
    user,
    token,
    loading,
    login,
    register,
    logout,
    getAuthHeaders,
    setSession,
    refreshMeAndSession
  };

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

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