import { createContext, useContext, useState, useEffect, useCallback, useRef } from 'react';
import { PublicClientApplication, InteractionStatus } from '@azure/msal-browser';
import { MsalProvider, useMsal } from '@azure/msal-react';
import { msalConfig, loginRequest, GATEWAY_URL } from './authConfig';

// ─── MSAL instance (singleton) ───
const msalInstance = new PublicClientApplication(msalConfig);
const msalReady = msalInstance.initialize();

// ─── Context ───
const AuthContext = createContext(null);
export const useAuth = () => useContext(AuthContext);

// ─── Session storage keys ───
const SK_TOKEN   = 'adp_session_token';
const SK_USER    = 'adp_session_user';

// ─── Inner provider (needs useMsal) ───
function AuthInner({ children }) {
  const { instance, inProgress, accounts } = useMsal();

  // States: unauthenticated | authenticating | needsRegistration | active | error
  const [authState, setAuthState]     = useState('authenticating');
  const [sessionToken, setSessionToken] = useState(null);
  const [sessionUser, setSessionUser] = useState(null);
  const [error, setError]             = useState(null);
  const exchangingRef = useRef(false);

  // ─── Exchange JWT for Gateway session ───
  const exchangeForSession = useCallback(async (jwt) => {
    if (exchangingRef.current) return;
    exchangingRef.current = true;

    console.log('[Auth] Exchanging JWT for session at', `${GATEWAY_URL}/api/auth/session`);
    try {
      const res = await fetch(`${GATEWAY_URL}/api/auth/session`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${jwt}`
        }
      });

      const body = await res.json();
      console.log('[Auth] Session response:', res.status, body.ok);

      if (body.ok && body.data) {
        const d = body.data;
        const user = {
          sessionId:    d.sessionId,
          userId:       d.userId,
          email:        d.userEmail,
          displayName:  d.displayName,
          role:         d.role,
          clientId:     d.clientId,
          clientName:   d.clientName,
          expiresAt:    d.expiresAt,
          availableClients: d.availableClients || []
        };
        setSessionToken(d.sessionToken);
        setSessionUser(user);
        sessionStorage.setItem(SK_TOKEN, d.sessionToken);
        sessionStorage.setItem(SK_USER, JSON.stringify(user));
        setAuthState('active');
        console.log('[Auth] Session established for', d.userEmail, '| client:', d.clientName);
      } else {
        const errMsg = body.error || 'Session creation failed';
        console.warn('[Auth] Session exchange error:', errMsg);

        if (/no client access|user not found|not registered/i.test(errMsg)) {
          setAuthState('needsRegistration');
        } else {
          setError(errMsg);
          setAuthState('error');
        }
      }
    } catch (err) {
      console.error('[Auth] Network error during session exchange:', err);
      setError('Unable to connect to the server. Please try again.');
      setAuthState('error');
    } finally {
      exchangingRef.current = false;
    }
  }, []);

  // ─── Validate existing session via /api/auth/me ───
  const validateSession = useCallback(async (token) => {
    try {
      const res = await fetch(`${GATEWAY_URL}/api/auth/me`, {
        headers: { 'Authorization': `Bearer ${token}` }
      });
      const body = await res.json();
      if (body.ok && body.data) {
        const d = body.data;
        const user = {
          sessionId:   d.sessionId,
          userId:      d.userId,
          email:       d.userEmail,
          displayName: d.displayName,
          role:        d.role,
          clientId:    d.clientId,
          clientName:  d.clientName,
          expiresAt:   d.expiresAt,
          availableClients: d.availableClients || []
        };
        setSessionToken(token);
        setSessionUser(user);
        setAuthState('active');
        console.log('[Auth] Session restored for', d.userEmail);
        return true;
      }
    } catch (e) {
      console.warn('[Auth] Session validation failed:', e.message);
    }
    return false;
  }, []);

  // ─── Handle MSAL redirect + session restoration ───
  useEffect(() => {
    let cancelled = false;

    async function init() {
      await msalReady;

      // 1. Check for MSAL redirect response
      try {
        const response = await instance.handleRedirectPromise();
        if (response && response.idToken && !cancelled) {
          console.log('[Auth] MSAL redirect received, exchanging token');
          await exchangeForSession(response.idToken);
          return;
        }
      } catch (err) {
        console.warn('[Auth] MSAL redirect error:', err);
      }

      if (cancelled) return;

      // 2. Try restoring existing session from storage
      const savedToken = sessionStorage.getItem(SK_TOKEN);
      if (savedToken) {
        const valid = await validateSession(savedToken);
        if (valid) return;
        // Stored session invalid — clear it
        sessionStorage.removeItem(SK_TOKEN);
        sessionStorage.removeItem(SK_USER);
      }

      // 3. If MSAL has an active account, try silent token + session exchange
      const account = instance.getActiveAccount() || instance.getAllAccounts()[0];
      if (account) {
        try {
          const silent = await instance.acquireTokenSilent({ ...loginRequest, account });
          if (silent?.idToken && !cancelled) {
            await exchangeForSession(silent.idToken);
            return;
          }
        } catch (e) {
          console.warn('[Auth] Silent token acquisition failed:', e.message);
        }
      }

      if (!cancelled) {
        setAuthState('unauthenticated');
      }
    }

    init();
    return () => { cancelled = true; };
  }, [instance, exchangeForSession, validateSession]);

  // ─── Actions ───
  const signIn = useCallback(() => {
    setAuthState('authenticating');
    instance.loginRedirect(loginRequest).catch(err => {
      console.error('[Auth] Login redirect error:', err);
      setAuthState('unauthenticated');
    });
  }, [instance]);

  const signOut = useCallback(async () => {
    // Tell Gateway to invalidate session
    if (sessionToken) {
      try {
        await fetch(`${GATEWAY_URL}/api/auth/signoff`, {
          method: 'POST',
          headers: { 'Authorization': `Bearer ${sessionToken}` }
        });
      } catch (e) { /* best effort */ }
    }
    sessionStorage.removeItem(SK_TOKEN);
    sessionStorage.removeItem(SK_USER);
    setSessionToken(null);
    setSessionUser(null);
    setAuthState('unauthenticated');
    instance.logoutRedirect({ postLogoutRedirectUri: window.location.origin });
  }, [instance, sessionToken]);

  const retrySignIn = useCallback(() => {
    sessionStorage.clear();
    setError(null);
    setAuthState('unauthenticated');
    instance.logoutRedirect({ postLogoutRedirectUri: window.location.origin }).catch(() => {
      setAuthState('unauthenticated');
    });
  }, [instance]);

  const switchClient = useCallback(async (clientId) => {
    if (!sessionToken) return;
    try {
      const res = await fetch(`${GATEWAY_URL}/api/auth/switch-client`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${sessionToken}`
        },
        body: JSON.stringify({ clientId })
      });
      const body = await res.json();
      if (body.ok && body.data) {
        const d = body.data;
        setSessionUser(prev => ({
          ...prev,
          clientId:   d.clientId   || prev.clientId,
          clientName: d.clientName || prev.clientName,
          role:       d.role       || prev.role
        }));
        sessionStorage.setItem(SK_USER, JSON.stringify({ ...sessionUser, ...d }));
      }
    } catch (e) {
      console.error('[Auth] Switch client error:', e);
    }
  }, [sessionToken, sessionUser]);

  // ─── MSAL account info (available even before session) ───
  const msalAccount = accounts?.[0] || null;

  const value = {
    authState,
    sessionToken,
    sessionUser,
    error,
    msalAccount,
    gatewayUrl: GATEWAY_URL,
    signIn,
    signOut,
    retrySignIn,
    switchClient,
    clearError: () => { setError(null); setAuthState('unauthenticated'); }
  };

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

// ─── Outer wrapper (provides MsalProvider) ───
export default function AuthProvider({ children }) {
  return (
    <MsalProvider instance={msalInstance}>
      <AuthInner>{children}</AuthInner>
    </MsalProvider>
  );
}


// ─── Named export alias (index.js uses { AuthProvider }) ───
export { AuthProvider };
