FPSMS-frontend/src/app/(main)/axios/AxiosProvider.tsx

// src/app/(main)/axios/AxiosProvider.tsx

"use client";

import React, { createContext, useContext, useEffect, useState, useCallback } from "react";
import { getSession } from "next-auth/react";
import { SessionWithTokens } from "@/config/authConfig";
import {
  isBackendJwtExpired,
  LOGIN_SESSION_EXPIRED_HREF,
} from "@/app/utils/authToken";
import axiosInstance, { SetupAxiosInterceptors } from "./axiosInstance";

const AxiosContext = createContext(axiosInstance);

const TokenContext = createContext<{
  setAccessToken: (token: string | null) => void;
}>({
  setAccessToken: () => {},
});

export const AxiosProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
  const [accessToken, setAccessToken] = useState<string | null>(null);
  const [isHydrated, setIsHydrated] = useState(false);

  // Hydrate token only on client
  useEffect(() => {
    try {
      const token = localStorage.getItem("accessToken");
      if (token) setAccessToken(token);
    } catch (e) {
      console.warn("localStorage unavailable", e);
    } finally {
      setIsHydrated(true);
    }
  }, []);

  /**
   * Detect expired/missing backend JWT before user actions (e.g. /report search).
   * Sync accessToken from next-auth session into localStorage if missing, then
   * redirect to login when the Bearer token is absent or past `exp`.
   */
  useEffect(() => {
    if (!isHydrated || typeof window === "undefined") return;

    let cancelled = false;
    (async () => {
      try {
        let token = localStorage.getItem("accessToken")?.trim() ?? "";

        if (!token) {
          const session = (await getSession()) as SessionWithTokens | null;
          if (cancelled) return;
          if (session?.accessToken) {
            token = session.accessToken;
            localStorage.setItem("accessToken", token);
            setAccessToken(token);
          }
        }

        if (!token) {
          window.location.href = LOGIN_SESSION_EXPIRED_HREF;
          return;
        }

        if (isBackendJwtExpired(token)) {
          window.location.href = LOGIN_SESSION_EXPIRED_HREF;
        }
      } catch (e) {
        console.warn("Auth token check failed", e);
      }
    })();

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

  // Apply token + interceptors
  useEffect(() => {
    if (accessToken) {
      axiosInstance.defaults.headers.Authorization = `Bearer ${accessToken}`;
      SetupAxiosInterceptors(accessToken);
    } else {
      delete axiosInstance.defaults.headers.Authorization;
    }
  }, [accessToken]);

  const handleSetAccessToken = useCallback((token: string | null) => {
    setAccessToken(token);
    try {
      if (token) {
        localStorage.setItem("accessToken", token);
      } else {
        localStorage.removeItem("accessToken");
      }
    } catch (e) {
      // ignore (e.g. private mode)
    }
  }, []);

  // Critical fix: never return null → always return children wrapped in fragment
  return (
    <AxiosContext.Provider value={axiosInstance}>
      <TokenContext.Provider value={{ setAccessToken: handleSetAccessToken }}>
        {/* Render children immediately – they will just not have the token for 1-2ms */}
        {children}
      </TokenContext.Provider>
    </AxiosContext.Provider>
  );
};

export const useAxios = () => useContext(AxiosContext);
export const useToken = () => useContext(TokenContext);