AuthContext.tsx

"use client";

import {
  createContext,
  useContext,
  useState,
  useCallback,
  useEffect,
} from "react";
import { APIKey } from "@/app/lib/types/credentials";
import {
  User,
  GoogleProfile,
  Session,
  AuthContextValue,
} from "@/app/lib/types/auth";
import { apiFetch } from "@/app/lib/apiClient";
import {
  AUTH_EXPIRED_EVENT,
  COOKIE_KEYS,
  FEATURES_UPDATED_EVENT,
  STORAGE_KEYS,
} from "@/app/lib/constants";
import { useChatStore } from "@/app/lib/store/chat";
import { clearAllStorage, readClientCookie } from "@/app/lib/utils";
export type { User, GoogleProfile, Session } from "@/app/lib/types/auth";

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

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const [apiKeys, setApiKeys] = useState<APIKey[]>([]);
  const [isHydrated, setIsHydrated] = useState(false);
  const [currentUser, setCurrentUser] = useState<User | null>(null);
  const [session, setSession] = useState<Session | null>(null);

  useEffect(() => {
    try {
      const rawMeta = readClientCookie(COOKIE_KEYS.API_KEY_META);
      if (rawMeta) setApiKeys([JSON.parse(rawMeta) as APIKey]);

      const storedSession = localStorage.getItem(STORAGE_KEYS.SESSION);
      if (storedSession) {
        const parsed = JSON.parse(storedSession) as Session;
        setSession(parsed);
        if (parsed.user) setCurrentUser(parsed.user);
      }
    } catch {
      /* ignore malformed data */
    }
    setIsHydrated(true);
  }, []);

  // Always fetch the latest user profile on hydration.
  useEffect(() => {
    if (!isHydrated) return;
    const hasApiKey = apiKeys.length > 0;
    const hasSession = !!session;
    if (!hasApiKey && !hasSession) return;

    let cancelled = false;

    (async () => {
      try {
        const res = await apiFetch<{ success: boolean; data: User } | User>(
          "/api/users/me",
          apiKeys[0]?.key ?? "",
        );
        const userData =
          "success" in res && res.data ? res.data : (res as User);
        if (!cancelled) {
          setCurrentUser(userData);
          const storedRaw = localStorage.getItem(STORAGE_KEYS.SESSION);
          if (storedRaw) {
            try {
              const stored = JSON.parse(storedRaw);
              stored.user = userData;
              localStorage.setItem(
                STORAGE_KEYS.SESSION,
                JSON.stringify(stored),
              );
            } catch {
              /* ignore */
            }
          }
        }
      } catch {
        // silently ignore
      }
    })();

    return () => {
      cancelled = true;
    };
  }, [apiKeys, session, isHydrated]);

  const addKey = useCallback(async (input: { key: string; label: string }) => {
    const res = await fetch("/api/apikeys", {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      credentials: "include",
      body: JSON.stringify(input),
    });
    const body = (await res.json().catch(() => ({}))) as {
      data?: APIKey;
      error?: string;
    };
    if (!res.ok || !body.data) {
      throw new Error(body.error || "Failed to add API key");
    }
    setApiKeys([body.data]);
    window.dispatchEvent(new Event("kaapi-auth-changed"));
  }, []);

  const removeKey = useCallback(async () => {
    await fetch("/api/apikeys", { method: "DELETE", credentials: "include" });
    setApiKeys([]);
    window.dispatchEvent(new Event("kaapi-auth-changed"));
  }, []);

  const loginWithToken = useCallback(
    (accessToken: string, user?: User, googleProfile?: GoogleProfile) => {
      const newSession: Session = {
        accessToken,
        user: user ?? null,
        googleProfile: googleProfile ?? null,
      };
      setSession(newSession);
      localStorage.setItem(STORAGE_KEYS.SESSION, JSON.stringify(newSession));

      if (user) setCurrentUser(user);
    },
    [],
  );

  const logout = useCallback(async () => {
    try {
      await fetch("/api/auth/logout", {
        method: "POST",
        credentials: "include",
      });
    } catch {
      // Clear local state even if the backend call fails
    }
    setSession(null);
    setCurrentUser(null);
    clearAllStorage();
    useChatStore.getState().reset();
    setApiKeys([]);
    window.dispatchEvent(new Event("kaapi-auth-changed"));
    window.location.replace("/chat");
  }, []);

  // logout when both access + refresh tokens are expired
  useEffect(() => {
    const handleExpired = () => logout();
    window.addEventListener(AUTH_EXPIRED_EVENT, handleExpired);
    return () => window.removeEventListener(AUTH_EXPIRED_EVENT, handleExpired);
  }, [logout]);

  // keep auth state in sync when feature flags are updated client-side
  useEffect(() => {
    const handleFeaturesUpdated = (event: Event) => {
      const customEvent = event as CustomEvent<{ features?: string[] }>;
      const features = customEvent.detail?.features;
      if (!Array.isArray(features)) return;

      setCurrentUser((prev) =>
        prev ? { ...prev, features: [...features] } : prev,
      );
      setSession((prev) =>
        prev?.user
          ? {
              ...prev,
              user: { ...prev.user, features: [...features] },
            }
          : prev,
      );
    };

    window.addEventListener(
      FEATURES_UPDATED_EVENT,
      handleFeaturesUpdated as EventListener,
    );
    return () =>
      window.removeEventListener(
        FEATURES_UPDATED_EVENT,
        handleFeaturesUpdated as EventListener,
      );
  }, []);

  const activeKey = apiKeys[0] ?? null;
  const isAuthenticated = !!activeKey || !!session;
  const features = currentUser?.features ?? [];
  const hasFeature = useCallback(
    (flag: string) => features.includes(flag),
    [features],
  );

  return (
    <AuthContext.Provider
      value={{
        apiKeys,
        activeKey,
        isHydrated,
        currentUser,
        googleProfile: session?.googleProfile ?? null,
        session,
        isAuthenticated,
        features,
        hasFeature,
        addKey,
        removeKey,
        loginWithToken,
        logout,
      }}
    >
      {children}
    </AuthContext.Provider>
  );
}

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