ui/src/contexts/AuthContext.tsx

'use client';

import { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import { loadSettings, persistSettings } from '@/utils/storage/settingsStorage';
import { oauthClientId } from '@/utils/env';

type AuthMethod = 'oauth' | 'manual';

interface StoredAuthState {
  token: string;
  namespace: string;
  method: AuthMethod;
}

export type AuthStatus = 'checking' | 'authenticated' | 'unauthenticated' | 'error';

interface AuthContextValue {
  status: AuthStatus;
  token: string | null;
  namespace: string | null;
  method: AuthMethod | null;
  error: string | null;
  oauthAvailable: boolean;
  loginWithOAuth: () => void;
  exchangeCodeForToken: (code: string, state: string) => Promise<boolean>;
  setManualToken: (token: string) => Promise<void>;
  logout: () => void;
}

const STORAGE_KEY = 'HF_AUTH_STATE';

const defaultValue: AuthContextValue = {
  status: 'checking',
  token: null,
  namespace: null,
  method: null,
  error: null,
  oauthAvailable: Boolean(oauthClientId),
  loginWithOAuth: () => {},
  exchangeCodeForToken: async () => false,
  setManualToken: async () => {},
  logout: () => {},
};

const AuthContext = createContext<AuthContextValue>(defaultValue);

async function validateToken(token: string) {
  const res = await fetch('/api/auth/hf/validate', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({ token }),
  });

  if (!res.ok) {
    const data = await res.json().catch(() => ({}));
    const error: any = new Error(data?.error || 'Failed to validate token');
    error.status = res.status;
    throw error;
  }

  return res.json();
}

async function syncTokenWithSettings(token: string) {
  try {
    const current = await loadSettings();
    if (current.HF_TOKEN === token) {
      return;
    }
    current.HF_TOKEN = token;
    await persistSettings(current);
  } catch (error) {
    console.warn('Failed to persist HF token to settings:', error);
  }
}

async function clearTokenFromSettings() {
  try {
    const current = await loadSettings();
    if (current.HF_TOKEN !== '') {
      current.HF_TOKEN = '';
      await persistSettings(current);
    }
  } catch (error) {
    console.warn('Failed to clear HF token from settings:', error);
  }
}

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const [status, setStatus] = useState<AuthStatus>('checking');
  const [token, setToken] = useState<string | null>(null);
  const [namespace, setNamespace] = useState<string | null>(null);
  const [method, setMethod] = useState<AuthMethod | null>(null);
  const [error, setError] = useState<string | null>(null);

  const oauthAvailable = Boolean(oauthClientId);

  const applyAuthState = useCallback(async ({ token: nextToken, namespace: nextNamespace, method: nextMethod }: StoredAuthState) => {
    setToken(nextToken);
    setNamespace(nextNamespace);
    setMethod(nextMethod);
    setStatus('authenticated');
    setError(null);

    if (typeof window !== 'undefined') {
      window.localStorage.setItem(
        STORAGE_KEY,
        JSON.stringify({
          token: nextToken,
          namespace: nextNamespace,
          method: nextMethod,
        }),
      );
    }

    syncTokenWithSettings(nextToken).catch(err => {
      console.warn('Failed to sync HF token with settings:', err);
    });
  }, []);

  const clearAuthState = useCallback(async () => {
    setToken(null);
    setNamespace(null);
    setMethod(null);
    setStatus('unauthenticated');
    setError(null);

    if (typeof window !== 'undefined') {
      window.localStorage.removeItem(STORAGE_KEY);
    }

    clearTokenFromSettings().catch(err => {
      console.warn('Failed to clear HF token from settings:', err);
    });
  }, []);

  // Restore stored token on mount
  useEffect(() => {
    if (typeof window === 'undefined') {
      return;
    }

    const restore = async () => {
      const raw = window.localStorage.getItem(STORAGE_KEY);
      if (!raw) {
        setStatus('unauthenticated');
        return;
      }

      try {
        const stored: StoredAuthState = JSON.parse(raw);
        if (!stored?.token) {
          setStatus('unauthenticated');
          return;
        }
        setStatus('checking');
        const data = await validateToken(stored.token);
        await applyAuthState({
          token: stored.token,
          namespace: data?.name || data?.preferred_username || stored.namespace || 'user',
          method: stored.method || 'manual',
        });
      } catch (err: any) {
        console.warn('Stored HF token validation failed:', err);
        if (err?.status === 401 || err?.status === 403) {
          await clearAuthState();
        } else {
          await applyAuthState({
            token: stored.token,
            namespace: stored.namespace || 'user',
            method: stored.method || 'manual',
          });
        }
      }
    };

    restore();
  }, [applyAuthState, clearAuthState]);

  const setManualToken = useCallback(
    async (manualToken: string) => {
      if (!manualToken) {
        setError('Please provide a token');
        setStatus('error');
        return;
      }
      setStatus('checking');
      setError(null);
      try {
        const data = await validateToken(manualToken);
        await applyAuthState({
          token: manualToken,
          namespace: data?.name || data?.preferred_username || 'user',
          method: 'manual',
        });
      } catch (err: any) {
        setError(err?.message || 'Failed to validate token');
        setStatus('error');
      }
    },
    [applyAuthState],
  );

  const exchangeCodeForToken = useCallback(
    async (code: string, state: string) => {
      if (!code || !state) {
        setError('Invalid authorization response.');
        setStatus('error');
        return false;
      }
      if (typeof window !== 'undefined') {
        const storedState = sessionStorage.getItem('HF_OAUTH_STATE');
        if (!storedState || storedState !== state) {
          setError('Invalid or expired OAuth state. Please try again.');
          setStatus('error');
          return false;
        }
      }
      setStatus('checking');
      setError(null);
      try {
        const res = await fetch('/api/auth/hf/exchange', {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          credentials: 'same-origin',
          body: JSON.stringify({ code, state }),
        });

        if (!res.ok) {
          const data = await res.json().catch(() => ({}));
          throw new Error(data?.error || 'Failed to exchange authorization code');
        }

        const data = await res.json();
        await applyAuthState({
          token: data.token,
          namespace: data.namespace || 'user',
          method: 'oauth',
        });
        if (typeof window !== 'undefined') {
          sessionStorage.removeItem('HF_OAUTH_STATE');
        }
        return true;
      } catch (err: any) {
        setError(err?.message || 'Failed to authenticate with Hugging Face');
        setStatus('error');
        if (typeof window !== 'undefined') {
          sessionStorage.removeItem('HF_OAUTH_STATE');
        }
        return false;
      }
    },
    [applyAuthState],
  );

  const loginWithOAuth = useCallback(() => {
    if (typeof window === 'undefined') {
      return;
    }
    if (!oauthAvailable) {
      setError('OAuth is not available on this deployment.');
      setStatus('error');
      return;
    }
    setStatus('checking');
    setError(null);

    const state = window.crypto.randomUUID();
    sessionStorage.setItem('HF_OAUTH_STATE', state);
    const loginUrl = new URL('/api/auth/hf/login', window.location.origin);
    loginUrl.searchParams.set('state', state);
    window.location.href = loginUrl.toString();
  }, []);

  const logout = useCallback(() => {
    clearAuthState();
  }, [clearAuthState]);

  const value = useMemo<AuthContextValue>(
    () => ({
      status,
      token,
      namespace,
      method,
      error,
      oauthAvailable,
      loginWithOAuth,
      exchangeCodeForToken,
      setManualToken,
      logout,
    }),
    [status, token, namespace, method, error, oauthAvailable, loginWithOAuth, exchangeCodeForToken, setManualToken, logout],
  );

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

export function useAuth() {
  return useContext(AuthContext);
}