import { useCallback, useState } from 'react';

export function useSessionProtection() {
  const [activeSessions, setActiveSessions] = useState<Set<string>>(new Set());
  const [processingSessions, setProcessingSessions] = useState<Set<string>>(new Set());
  // sessionId -> epoch ms when processing began (for the live elapsed timer)
  const [processingStartTimes, setProcessingStartTimes] = useState<Map<string, number>>(new Map());
  // sessions waiting on the user (permission / AskUserQuestion pending)
  const [needsInputSessions, setNeedsInputSessions] = useState<Set<string>>(new Set());

  const markSessionAsActive = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setActiveSessions((prev) => new Set([...prev, sessionId]));
  }, []);

  const markSessionAsInactive = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setActiveSessions((prev) => {
      const next = new Set(prev);
      next.delete(sessionId);
      return next;
    });
  }, []);

  const markSessionAsProcessing = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setProcessingSessions((prev) => new Set([...prev, sessionId]));
    setProcessingStartTimes((prev) => {
      // Preserve an existing start time so the elapsed timer doesn't reset on
      // repeated marks for the same run.
      if (prev.has(sessionId)) {
        return prev;
      }
      const next = new Map(prev);
      next.set(sessionId, Date.now());
      return next;
    });
  }, []);

  const markSessionAsNotProcessing = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setProcessingSessions((prev) => {
      const next = new Set(prev);
      next.delete(sessionId);
      return next;
    });
    setProcessingStartTimes((prev) => {
      if (!prev.has(sessionId)) {
        return prev;
      }
      const next = new Map(prev);
      next.delete(sessionId);
      return next;
    });
  }, []);

  const markSessionAsNeedsInput = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setNeedsInputSessions((prev) => new Set([...prev, sessionId]));
  }, []);

  const clearSessionNeedsInput = useCallback((sessionId?: string | null) => {
    if (!sessionId) {
      return;
    }

    setNeedsInputSessions((prev) => {
      if (!prev.has(sessionId)) {
        return prev;
      }
      const next = new Set(prev);
      next.delete(sessionId);
      return next;
    });
  }, []);

  return {
    activeSessions,
    processingSessions,
    processingStartTimes,
    needsInputSessions,
    markSessionAsActive,
    markSessionAsInactive,
    markSessionAsProcessing,
    markSessionAsNotProcessing,
    markSessionAsNeedsInput,
    clearSessionNeedsInput,
  };
}
