UMBRA/frontend/src/hooks/useLock.tsx

import {
  createContext,
  useContext,
  useState,
  useCallback,
  useEffect,
  useRef,
  type ReactNode,
} from 'react';
import { useIsMutating, useQueryClient } from '@tanstack/react-query';
import { useSettings } from '@/hooks/useSettings';
import api from '@/lib/api';
import type { AuthStatus } from '@/types';

interface LockContextValue {
  isLocked: boolean;
  isLockResolved: boolean;
  lock: () => Promise<void>;
  unlock: (password: string) => Promise<void>;
}

const LockContext = createContext<LockContextValue | null>(null);

const ACTIVITY_EVENTS = ['mousemove', 'keydown', 'click', 'scroll', 'touchstart'] as const;
const THROTTLE_MS = 5_000; // only reset timer every 5s of activity

export function LockProvider({ children }: { children: ReactNode }) {
  const queryClient = useQueryClient();

  // Initialize lock state from the auth status cache (server-persisted)
  const [isLocked, setIsLocked] = useState(() => {
    const cached = queryClient.getQueryData<AuthStatus>(['auth']);
    return cached?.is_locked ?? false;
  });

  // Track whether lock state has been definitively resolved from the server
  const [isLockResolved, setIsLockResolved] = useState(() => {
    return queryClient.getQueryData<AuthStatus>(['auth']) !== undefined;
  });

  const { settings } = useSettings();
  const activeMutations = useIsMutating();

  const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const lastActivityRef = useRef<number>(Date.now());
  const activeMutationsRef = useRef(activeMutations);
  activeMutationsRef.current = activeMutations;

  // Subscribe to auth query updates to catch server lock state on refresh
  useEffect(() => {
    const unsubscribe = queryClient.getQueryCache().subscribe((event) => {
      if (
        event.type === 'updated' &&
        (event.action.type === 'success' || event.action.type === 'error') &&
        event.query.queryKey[0] === 'auth'
      ) {
        setIsLockResolved(true);
        if (event.action.type === 'success') {
          const data = event.query.state.data as AuthStatus | undefined;
          if (data?.is_locked) {
            setIsLocked(true);
          }
        }
      }
    });
    return unsubscribe;
  }, [queryClient]);

  // Listen for 423 responses from the API interceptor (server-side lock enforcement)
  useEffect(() => {
    const handler = () => setIsLocked(true);
    window.addEventListener('umbra:session-locked', handler);
    return () => window.removeEventListener('umbra:session-locked', handler);
  }, []);

  const lock = useCallback(async () => {
    setIsLocked(true);
    try {
      await api.post('/auth/lock');
    } catch {
      // Lock locally even if server call fails — defense in depth
    }
  }, []);

  const unlock = useCallback(async (password: string) => {
    const { data } = await api.post<{ verified: boolean }>('/auth/verify-password', { password });
    if (data.verified) {
      setIsLocked(false);
      lastActivityRef.current = Date.now();
      // Update auth cache to reflect unlocked state
      queryClient.setQueryData<AuthStatus>(['auth'], (old) =>
        old ? { ...old, is_locked: false } : old
      );
    }
  }, [queryClient]);

  // Auto-lock idle timer
  useEffect(() => {
    const enabled = settings?.auto_lock_enabled ?? false;
    const minutes = settings?.auto_lock_minutes ?? 5;

    if (!enabled || isLocked) {
      // Clear any existing timer when disabled or already locked
      if (timerRef.current) {
        clearTimeout(timerRef.current);
        timerRef.current = null;
      }
      return;
    }

    const timeoutMs = minutes * 60_000;

    const resetTimer = () => {
      // Don't lock while TanStack mutations are in flight
      if (activeMutationsRef.current > 0) return;

      if (timerRef.current) clearTimeout(timerRef.current);
      timerRef.current = setTimeout(() => {
        lock();
      }, timeoutMs);
    };

    const handleActivity = () => {
      const now = Date.now();
      if (now - lastActivityRef.current < THROTTLE_MS) return;
      lastActivityRef.current = now;
      resetTimer();
    };

    // Start the initial timer
    resetTimer();

    // Attach throttled listeners
    for (const event of ACTIVITY_EVENTS) {
      document.addEventListener(event, handleActivity, { passive: true });
    }

    return () => {
      for (const event of ACTIVITY_EVENTS) {
        document.removeEventListener(event, handleActivity);
      }
      if (timerRef.current) {
        clearTimeout(timerRef.current);
        timerRef.current = null;
      }
    };
  }, [settings?.auto_lock_enabled, settings?.auto_lock_minutes, isLocked, lock]);

  return (
    <LockContext.Provider value={{ isLocked, isLockResolved, lock, unlock }}>
      {children}
    </LockContext.Provider>
  );
}

export function useLock(): LockContextValue {
  const ctx = useContext(LockContext);
  if (!ctx) throw new Error('useLock must be used within a LockProvider');
  return ctx;
}