ChatApp/apps/mobile/lib/authContext.tsx

import type { Session, User } from '@supabase/supabase-js';
import { fetchUserKeyBlob } from '@chat-app/shared/auth';
import {
  type ReactNode,
  createContext,
  useCallback,
  useContext,
  useEffect,
  useState,
} from 'react';

import { supabase } from './supabase';
import { cachedUserKey, ensureLegacyMigrated } from './userIdentity';

export type UserKeyState =
  | { status: 'loading' }
  | { status: 'needs-setup' }
  | { status: 'needs-unlock'; lockedUntil: string | null; hasRecovery: boolean }
  | { status: 'unlocked' };

interface AuthContextValue {
  session: Session | null;
  user: User | null;
  userId: string | null;
  ownPrivateKey: Uint8Array | null;
  userKeyState: UserKeyState;
  ready: boolean;
  refreshUserKeyState: () => Promise<void>;
  signOut: () => Promise<void>;
}

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

export function useAuth(): AuthContextValue {
  const v = useContext(Ctx);
  if (!v) throw new Error('useAuth() called outside <AuthProvider>');
  return v;
}

export function AuthProvider({ children }: { children: ReactNode }) {
  const [session, setSession] = useState<Session | null>(null);
  const [ownPrivateKey, setOwnPrivateKey] = useState<Uint8Array | null>(null);
  const [userKeyState, setUserKeyState] = useState<UserKeyState>({ status: 'loading' });
  const [ready, setReady] = useState(false);

  const refreshUserKeyState = useCallback(async () => {
    const s = session;
    if (!s) {
      setUserKeyState({ status: 'loading' });
      setOwnPrivateKey(null);
      return;
    }
    setUserKeyState({ status: 'loading' });
    const cached = await cachedUserKey(s.user.id);
    if (cached) {
      setOwnPrivateKey(cached);
      setUserKeyState({ status: 'unlocked' });
      void ensureLegacyMigrated(s.user.id).catch((err) => {
        console.warn('legacy migration on auth-resume failed', err);
      });
      return;
    }
    const blob = await fetchUserKeyBlob(supabase, s.user.id);
    if (!blob || !blob.exists) {
      setUserKeyState({ status: 'needs-setup' });
      return;
    }
    if (blob.locked) {
      setUserKeyState({
        status: 'needs-unlock',
        lockedUntil: blob.lockedUntil,
        hasRecovery: false,
      });
      return;
    }
    setUserKeyState({
      status: 'needs-unlock',
      lockedUntil: null,
      hasRecovery: blob.recoverySealedPrivateKey !== null,
    });
  }, [session]);

  useEffect(() => {
    let cancelled = false;
    void (async () => {
      const { data } = await supabase.auth.getSession();
      if (cancelled) return;
      setSession(data.session);
      setReady(true);
    })();
    const { data: sub } = supabase.auth.onAuthStateChange((_event, nextSession) => {
      setSession(nextSession);
      setReady(true);
      if (!nextSession) {
        setOwnPrivateKey(null);
        setUserKeyState({ status: 'loading' });
      }
    });
    return () => {
      cancelled = true;
      sub.subscription.unsubscribe();
    };
  }, []);

  useEffect(() => {
    void refreshUserKeyState().catch((err) => {
      console.warn('refreshUserKeyState failed', err);
      setUserKeyState({ status: 'needs-setup' });
    });
  }, [session, refreshUserKeyState]);

  const signOut = useCallback(async () => {
    await supabase.auth.signOut();
    setOwnPrivateKey(null);
    setUserKeyState({ status: 'loading' });
  }, []);

  const value: AuthContextValue = {
    session,
    user: session?.user ?? null,
    userId: session?.user.id ?? null,
    ownPrivateKey,
    userKeyState,
    ready,
    refreshUserKeyState,
    signOut,
  };
  return <Ctx.Provider value={value}>{children}</Ctx.Provider>;
}