import {
  createContext,
  FC,
  ReactNode,
  useCallback,
  useContext,
  useRef,
  useState,
} from 'react';

import { useAuth } from 'src/contexts/JWTContext';
import useLocalStorage from 'src/hooks/useLocalStorage';

type OpenMenuCallback = () => Promise<HTMLElement | null>;
type StepChangeCallback = (step: number) => void;

type ShiftsCoachmarkContextValue = {
  currentStep: number;
  setCurrentStep: (step: number) => void;
  startTour: () => void;
  handleCoachmarkClose: () => void;
  hasCompletedCoachmark: boolean;
  registerOpenStep3Menu: (callback: OpenMenuCallback) => void;
  openStep3Menu: () => Promise<HTMLElement | null>;
  registerOnStepChange: (callback: StepChangeCallback) => void;
};

const ShiftsCoachmarkContext =
  createContext<ShiftsCoachmarkContextValue | null>(null);

export const ShiftsCoachmarkProvider: FC<{ children: ReactNode }> = ({
  children,
}) => {
  const { user } = useAuth();
  const storageKey = `shifts-coachmark-${user?.id}`;
  const [hasCompletedCoachmark, setHasCompletedCoachmark] =
    useLocalStorage<boolean>(storageKey, false);
  const [currentStepState, setCurrentStepState] = useState(0);

  // Ref to store the callback for opening step 3 menu
  const openStep3MenuRef = useRef<OpenMenuCallback | null>(null);
  // Ref to store callbacks for step changes
  const stepChangeCallbacksRef = useRef<Set<StepChangeCallback>>(new Set());

  /** Set current step and notify all registered callbacks */
  const setCurrentStep = useCallback((step: number) => {
    setCurrentStepState(step);
    stepChangeCallbacksRef.current.forEach(callback => callback(step));
  }, []);

  /** Close intro modal and start the coachmark tour */
  const startTour = useCallback(() => {
    setCurrentStep(1);
  }, [setCurrentStep]);

  /** Called when coachmark closes (skip, close X, or finish) */
  const handleCoachmarkClose = useCallback(() => {
    setHasCompletedCoachmark(true);
  }, [setHasCompletedCoachmark]);

  /** Register callback to open step 3 menu */
  const registerOpenStep3Menu = useCallback((callback: OpenMenuCallback) => {
    openStep3MenuRef.current = callback;
  }, []);

  /** Call the registered callback to open step 3 menu */
  const openStep3Menu = useCallback(async () => {
    if (openStep3MenuRef.current) {
      return openStep3MenuRef.current();
    }
    return null;
  }, []);

  /** Register callback to be called when step changes */
  const registerOnStepChange = useCallback((callback: StepChangeCallback) => {
    stepChangeCallbacksRef.current.add(callback);
    // Return cleanup function to unregister
    return () => {
      stepChangeCallbacksRef.current.delete(callback);
    };
  }, []);

  const value: ShiftsCoachmarkContextValue = {
    currentStep: currentStepState,
    setCurrentStep,
    startTour,
    handleCoachmarkClose,
    hasCompletedCoachmark,
    registerOpenStep3Menu,
    openStep3Menu,
    registerOnStepChange,
  };

  return (
    <ShiftsCoachmarkContext.Provider value={value}>
      {children}
    </ShiftsCoachmarkContext.Provider>
  );
};

/**
 * Hook to access the Shifts coachmark context.
 * Must be used within a ShiftsCoachmarkProvider.
 */
export const useShiftsCoachmark = (): ShiftsCoachmarkContextValue => {
  const context = useContext(ShiftsCoachmarkContext);
  if (!context) {
    throw new Error(
      'useShiftsCoachmark must be used within a ShiftsCoachmarkProvider',
    );
  }
  return context;
};
