import React, { createContext, useCallback, useContext, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';

type FocusScopeValue = {
  reportDescendantFocusDelta(delta: number): void;
};

type PreferredFocusValue = {
  registerPreferredTarget(target: unknown): () => void;
};

const FocusScopeContext = createContext<FocusScopeValue | null>(null);
const FocusedStyleContext = createContext(false);
const FocusContext = createContext(true);
const PreferredFocusContext = createContext<PreferredFocusValue | null>(null);

export const FocusScope: React.FC<{
  children: React.ReactNode;
  onFocusWithinChange?: (focused: boolean) => void;
}> = ({ children, onFocusWithinChange }) => {
  const parentScope = useContext(FocusScopeContext);
  const focusedCountRef = useRef(0);

  const reportDescendantFocusDelta = useCallback((delta: number) => {
    const previousCount = focusedCountRef.current;
    const nextCount = Math.max(0, previousCount + delta);
    focusedCountRef.current = nextCount;

    if ((previousCount === 0) === (nextCount === 0)) {
      return;
    }

    const isFocused = nextCount > 0;
    onFocusWithinChange?.(isFocused);
    parentScope?.reportDescendantFocusDelta(isFocused ? 1 : -1);
  }, [onFocusWithinChange, parentScope]);

  useEffect(() => () => {
    if (focusedCountRef.current > 0) {
      parentScope?.reportDescendantFocusDelta(-1);
    }
  }, [parentScope]);

  const value = useMemo<FocusScopeValue>(
    () => ({ reportDescendantFocusDelta }),
    [reportDescendantFocusDelta],
  );

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

export const FocusedStyleProvider: React.FC<{
  children: React.ReactNode;
  focused: boolean;
}> = ({ children, focused }) => {
  const parentFocused = useContext(FocusedStyleContext);
  const value = parentFocused || focused;

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

export const FocusProvider: React.FC<{
  children: React.ReactNode;
  enabled: boolean;
}> = ({ children, enabled }) => (
  <FocusContext.Provider value={enabled}>
    {children}
  </FocusContext.Provider>
);

export const PreferredFocusProvider: React.FC<{
  children: React.ReactNode;
  onPreferredTargetChange?: (target: unknown | null) => void;
}> = ({ children, onPreferredTargetChange }) => {
  const registerPreferredTarget = useCallback((target: unknown) => {
    onPreferredTargetChange?.(target);

    return () => {
      onPreferredTargetChange?.(null);
    };
  }, [onPreferredTargetChange]);

  const value = useMemo<PreferredFocusValue>(
    () => ({ registerPreferredTarget }),
    [registerPreferredTarget],
  );

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

export function useInheritedFocusedStyle(): boolean {
  return useContext(FocusedStyleContext);
}

export function useFocusEnabled(): boolean {
  return useContext(FocusContext);
}

export function useRegisterPreferredFocus(target: unknown, enabled: boolean): void {
  const context = useContext(PreferredFocusContext);

  useEffect(() => {
    if (!enabled || target == null) {
      return;
    }

    return context?.registerPreferredTarget(target);
  }, [context, enabled, target]);
}

export function useFocusableState(initialFocused = false, enabled = true) {
  const scope = useContext(FocusScopeContext);
  const [isFocused, setIsFocused] = useState(initialFocused);
  const isFocusedRef = useRef(initialFocused);
  const reportedToScopeRef = useRef(false);

  const setFocused = useCallback((focused: boolean) => {
    if (isFocusedRef.current === focused) {
      return;
    }

    isFocusedRef.current = focused;
    setIsFocused(focused);
    if (focused && !reportedToScopeRef.current) {
      scope?.reportDescendantFocusDelta(1);
      reportedToScopeRef.current = true;
    } else if (!focused && reportedToScopeRef.current) {
      scope?.reportDescendantFocusDelta(-1);
      reportedToScopeRef.current = false;
    }
  }, [scope]);

  useLayoutEffect(() => {
    if (!enabled) {
      setFocused(false);
      return;
    }

    if (initialFocused && !reportedToScopeRef.current) {
      scope?.reportDescendantFocusDelta(1);
      reportedToScopeRef.current = true;
    }

    return () => {
      if (reportedToScopeRef.current) {
        reportedToScopeRef.current = false;
        scope?.reportDescendantFocusDelta(-1);
      }
    };
  }, [enabled, initialFocused, scope, setFocused]);

  return {
    isFocused,
    onFocus: useCallback(() => setFocused(true), [setFocused]),
    onBlur: useCallback(() => setFocused(false), [setFocused]),
    setFocused,
  };
}