ChatApp/apps/desktop/src/context/AuthContext.tsx

import {
  fetchUserKeyBlob,
  getOwnProfile,
  signOut as supabaseSignOut,
  type Profile,
  updateOwnProfile,
} from '@chat-app/shared/auth';
import { changeLocale, isSupportedLocale } from '@chat-app/shared/i18n';
import type { Session } from '@supabase/supabase-js';
import {
  createContext,
  type ReactNode,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import { useTranslation } from 'react-i18next';

import { ensureInstallId } from '../lib/installId';
import { setSecretStoreUser } from '../lib/secretStore';
import { supabase } from '../lib/supabase';
import { cachedUserKey, ensureLegacyMigrated } from '../lib/userIdentity';
import { registerWebPush } from '../lib/webPush';

// Discriminated union describing the per-user encrypted key blob lifecycle:
//
//   loading      — initial state, or refresh in flight
//   needs-setup  — no row exists on Supabase; user must pick a PIN
//   needs-unlock — row exists but local cache empty; PIN (or recovery code)
//                  required. `lockedUntil` non-null means the server-side
//                  rate limiter is currently rejecting attempts.
//   unlocked     — private key is in the local secret store and ready to use
export type UserKeyState =
  | { status: 'loading' }
  | { status: 'needs-setup' }
  | { status: 'needs-unlock'; lockedUntil: string | null; hasRecovery: boolean }
  | { status: 'unlocked' };

interface AuthContextValue {
  session: Session | null;
  profile: Profile | null;
  userKeyState: UserKeyState;
  // null while we're still resolving the very first auth state.
  ready: boolean;
  refreshProfile: () => Promise<void>;
  refreshUserKeyState: () => Promise<void>;
  signOut: () => Promise<void>;
}

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

export function AuthProvider({ children }: { children: ReactNode }) {
  const { i18n } = useTranslation();
  const [session, setSession] = useState<Session | null>(null);
  const [ready, setReady] = useState(false);
  const [profile, setProfile] = useState<Profile | null>(null);
  const [userKeyState, setUserKeyState] = useState<UserKeyState>({ status: 'loading' });
  const autoOnlineUserRef = useRef<string | null>(null);

  // Initial session + auth subscription. We verify the cached JWT against the
  // server (via getUser) once on mount. Only purge the session on an
  // unambiguous 401/403 — a network failure (Supabase stack offline) must not
  // log the user out, otherwise every local `supabase stop` wipes their session.
  useEffect(() => {
    let cancelled = false;
    (async () => {
      const { data: sessionRes } = await supabase.auth.getSession();
      if (cancelled) return;
      // Flip `ready` immediately on cached session read so the UI unblocks even
      // if the network is slow/down. Validate the token in the background and
      // only wipe on an unambiguous 401/403 — a stalled getUser (Tauri WebView
      // with no network, server unreachable) must not keep the app on the
      // loading spinner forever.
      setSession(sessionRes.session ?? null);
      setReady(true);

      if (sessionRes.session) {
        supabase.auth
          .getUser()
          .then(({ error }) => {
            if (cancelled || !error) return;
            const status = (error as { status?: number }).status;
            if (status === 401 || status === 403) {
              void supabase.auth.signOut({ scope: 'local' }).catch(() => {
                /* ignore */
              });
              setSession(null);
            } else {
              // Network / server unreachable — keep cached session.
              console.warn('auth.getUser failed, keeping cached session:', error);
            }
          })
          .catch((err: unknown) => {
            console.warn('auth.getUser rejected, keeping cached session:', err);
          });
      }
    })();
    const { data: sub } = supabase.auth.onAuthStateChange((_event, s) => {
      setSession(s);
      setReady(true);
      void setSecretStoreUser(s?.user.id ?? null);
    });
    return () => {
      cancelled = true;
      sub.subscription.unsubscribe();
    };
  }, []);

  const refreshProfile = useCallback(async () => {
    if (!session) {
      setProfile(null);
      return;
    }
    const p = await getOwnProfile(supabase);
    setProfile(p);
    if (p && p.locale !== i18n.resolvedLanguage && isSupportedLocale(p.locale)) {
      void changeLocale(p.locale);
    }
  }, [session, i18n.resolvedLanguage]);

  // Resolves the current state of the per-user encrypted key blob:
  //   1. local cache hit  → 'unlocked'
  //   2. no remote row    → 'needs-setup'
  //   3. server rate-limit → 'needs-unlock' with lockedUntil set
  //   4. otherwise         → 'needs-unlock'; hasRecovery reflects whether a
  //                          recovery-code blob is present so the UI can
  //                          conditionally offer the recovery affordance
  const refreshUserKeyState = useCallback(async () => {
    if (!session) {
      setUserKeyState({ status: 'loading' });
      return;
    }
    setUserKeyState({ status: 'loading' });
    const cached = await cachedUserKey(session.user.id);
    if (cached) {
      setUserKeyState({ status: 'unlocked' });
      // Best-effort: re-wrap any unmigrated legacy bundles. Idempotent (RPC
      // uses ON CONFLICT DO NOTHING). Recovers users who set up under 0.18.0
      // where the migration query had a `.eq(null)` bug that made it a no-op.
      void ensureLegacyMigrated(session.user.id).catch((err) => {
        console.warn('legacy conv-key migration on auth-resume failed', err);
      });
      return;
    }
    const blob = await fetchUserKeyBlob(supabase, session.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]);

  // Re-pull profile + user-key state whenever session flips.
  useEffect(() => {
    if (!session) {
      setProfile(null);
      setUserKeyState({ status: 'loading' });
      return;
    }
    void refreshProfile().catch((err: unknown) => {
      console.error('refreshProfile failed', err);
    });
    void refreshUserKeyState().catch((err: unknown) => {
      console.error('refreshUserKeyState failed', err);
      // Treat an unrecoverable lookup error as "needs-setup" so the UI at
      // least drives the user toward the setup/unlock page rather than
      // hanging forever on the spinner.
      setUserKeyState({ status: 'needs-setup' });
    });
  }, [session, refreshProfile, refreshUserKeyState]);

  // Best-effort web-push registration once we have a session. Keyed by an
  // install-id (localStorage UUID) since there's no longer a per-device
  // crypto record to key by. No-op on Tauri (uses native notifications) or
  // when VITE_VAPID_PUBLIC_KEY is unset.
  useEffect(() => {
    if (!session) return;
    const installId = ensureInstallId();
    void registerWebPush(installId);
  }, [session]);

  // Auto online/offline transition.
  //
  // - On mount with a session whose last persisted state is `offline`, flip
  //   to `online`. We never override an explicit `idle`, `dnd`, or
  //   `invisible` choice — those are user intent.
  // - On `pagehide` / `beforeunload`, fire a best-effort update to
  //   `offline`. Browsers don't guarantee delivery during unload, but the
  //   request usually slips through; the next page load corrects state if it
  //   didn't.
  useEffect(() => {
    if (!session) {
      autoOnlineUserRef.current = null;
      return;
    }
    if (!profile) return;
    const shouldApplyInitialAutoOnline = autoOnlineUserRef.current !== session.user.id;
    autoOnlineUserRef.current = session.user.id;
    if (profile.presenceState === 'offline' && shouldApplyInitialAutoOnline) {
      void updateOwnProfile(supabase, { presenceState: 'online' })
        .then(() => refreshProfile())
        .catch((err: unknown) => {
          console.warn('auto online flip failed', err);
        });
    }
    const onLeave = () => {
      // Skip if user explicitly chose a non-online state — they probably
      // want to look unavailable on next reconnect too.
      if (profile.presenceState !== 'online' && profile.presenceState !== 'offline') {
        return;
      }
      void updateOwnProfile(supabase, { presenceState: 'offline' }).catch(() => {});
    };
    window.addEventListener('beforeunload', onLeave);
    window.addEventListener('pagehide', onLeave);
    return () => {
      window.removeEventListener('beforeunload', onLeave);
      window.removeEventListener('pagehide', onLeave);
    };
  }, [session, profile, refreshProfile]);

  const signOut = useCallback(async () => {
    await updateOwnProfile(supabase, { presenceState: 'offline' }).catch((err: unknown) => {
      console.warn('offline update before sign-out failed', err);
    });
    await supabaseSignOut(supabase);
  }, []);

  const value = useMemo<AuthContextValue>(
    () => ({
      session,
      profile,
      userKeyState,
      ready,
      refreshProfile,
      refreshUserKeyState,
      signOut,
    }),
    [session, profile, userKeyState, ready, refreshProfile, refreshUserKeyState, signOut],
  );

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

export function useAuth(): AuthContextValue {
  const ctx = useContext(AuthContext);
  if (!ctx) throw new Error('useAuth must be used inside <AuthProvider>');
  return ctx;
}