'use client';

import { useStatsigClient } from '@statsig/react-bindings';
import { usePathname } from 'next/navigation';
import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useState,
} from 'react';
import storageAvailable from 'storage-available';
import { twMerge } from 'tailwind-merge';
import { useMediaQuery } from 'usehooks-ts';

import {
  THEME_PATHS_EXEMPT_FROM_LIGHT_MODE_REGEX,
  THEME_STORAGE_KEY,
} from '@/utils/constants';
import {
  PolymorphicComponent,
  PolymorphicComponentProps,
} from '@/utils/polymorphic';

export enum ThemeMode {
  Light = 'light',
  Dark = 'dark',
  System = 'system',
}

export type EffectiveThemeMode = Exclude<ThemeMode, ThemeMode.System>;

export const VALID_THEME_MODES = [
  ThemeMode.Light,
  ThemeMode.Dark,
  ThemeMode.System,
];

export const THEME_CLASS_NAME = {
  [ThemeMode.Light]: 'theme-light',
  [ThemeMode.Dark]: 'theme-dark',
  [ThemeMode.System]: '',
};

interface ThemeContextType {
  /**
   * The theme preference setting of the page root
   */
  rootTheme: ThemeMode;
  /**
   * Set the theme preference setting of the page root
   */
  setRootTheme: (theme: ThemeMode) => void;
  /**
   * The effective root theme that the user sees
   *
   * This may be different than the preference setting, such as when we have
   * resolved "system theme" to dark or light mode.
   */
  effectiveRootTheme: EffectiveThemeMode;
  /**
   * The local theme preference setting
   */
  theme: ThemeMode;
  /**
   * Set the theme preference setting
   */
  setTheme: (theme: ThemeMode) => void;
  /**
   * The effective theme that the user sees
   *
   * This may be different than the preference setting, such as when we have
   * resolved "system theme" to dark or light mode.
   */
  effectiveTheme: EffectiveThemeMode;
}

/**
 * Converts a theme mode to a CSS class name
 */
export function getThemeClassName(theme: ThemeMode) {
  return THEME_CLASS_NAME[theme];
}

/**
 * Sets the top-level theme on the document and body of the page
 */
export function applyThemeToDom(theme: ThemeMode) {
  const html = document.documentElement;
  const body = document.body;
  VALID_THEME_MODES.forEach((mode) => {
    const className = THEME_CLASS_NAME[mode];
    const isCurrentTheme = mode === theme;
    if (className) {
      html.classList.toggle(className, isCurrentTheme);
      body.classList.toggle(className, isCurrentTheme);
    }
  });
}

const ThemeContext = createContext<ThemeContextType | undefined>(undefined);

export type RootThemeProviderProps = React.PropsWithChildren<{
  defaultTheme?: ThemeMode;
  syncLocalStorage?: boolean;
  syncRootElement?: boolean;
}>;

/**
 * Sets the theme and syncs with the DOM and localstorage
 */
export const RootThemeProvider: React.FC<RootThemeProviderProps> = (props) => {
  const {
    children,
    defaultTheme = ThemeMode.Dark,
    syncLocalStorage = true,
    syncRootElement = true,
  } = props;

  const [theme, setThemeState] = useState<ThemeMode>(() => {
    // Load theme from localStorage on mount
    if (syncLocalStorage && storageAvailable('localStorage')) {
      const savedTheme = localStorage.getItem(THEME_STORAGE_KEY) as ThemeMode;
      if (savedTheme && VALID_THEME_MODES.includes(savedTheme)) {
        return savedTheme;
      }
    }
    return defaultTheme;
  });

  const systemPrefersDark = useMediaQuery('(prefers-color-scheme: dark)');
  const isPathExemptFromLightMode = usePathExemptFromLightMode();

  // Computed theme that user actually sees
  const effectiveTheme = useMemo(() => {
    if (isPathExemptFromLightMode) {
      return ThemeMode.Dark;
    }
    return theme === ThemeMode.System
      ? systemPrefersDark
        ? ThemeMode.Dark
        : ThemeMode.Light
      : theme;
  }, [theme, systemPrefersDark, isPathExemptFromLightMode]);

  // Hook for context consumers
  const setTheme = useCallback(
    (newTheme: ThemeMode) => {
      setThemeState(newTheme); // Auto-triggers re-render through useEffect below
      if (syncLocalStorage && storageAvailable('localStorage')) {
        localStorage.setItem(THEME_STORAGE_KEY, newTheme);
      }
    },
    [syncLocalStorage]
  );

  // Re-apply theme when pathname changes
  useEffect(() => {
    if (syncRootElement) {
      applyThemeToDom(effectiveTheme);
    }
  }, [syncRootElement, effectiveTheme]);

  // Revert to default theme if theme selector is explicitly disabled
  const statsigClient = useStatsigClient();
  const enableThemeSelector =
    statsigClient.client.loadingStatus === 'Ready'
      ? statsigClient.checkGate('web-theme-selector')
      : undefined;
  const revertToTheme =
    enableThemeSelector === false && theme !== defaultTheme
      ? defaultTheme
      : undefined;
  useEffect(() => {
    if (revertToTheme) {
      setTheme(revertToTheme);
    }
  }, [revertToTheme, setTheme]);

  return (
    <ThemeContext.Provider
      value={{
        rootTheme: theme,
        setRootTheme: setTheme,
        effectiveRootTheme: effectiveTheme,
        theme,
        setTheme,
        effectiveTheme,
      }}
    >
      {children}
    </ThemeContext.Provider>
  );
};

export type ThemeProviderProps = React.PropsWithChildren<{
  defaultTheme?: ThemeMode;
  theme?: ThemeMode;
  honorLightModePathExemption?: boolean;
}>;

/**
 * Sets the local theme for a branch of the component tree
 */
export const ThemeProvider: React.FC<ThemeProviderProps> = (props) => {
  const {
    children,
    defaultTheme: explicitDefaultTheme,
    theme: explicitTheme,
    honorLightModePathExemption = false,
  } = props;

  const parentThemeContext = useThemeContext();

  const [theme, setTheme] = useState<ThemeMode>(
    () => explicitTheme ?? explicitDefaultTheme ?? parentThemeContext.theme
  );

  const systemPrefersDark = useMediaQuery('(prefers-color-scheme: dark)');
  const isPathExemptFromLightMode = usePathExemptFromLightMode(
    honorLightModePathExemption ? undefined : false
  );

  // Computed theme that user actually sees
  const effectiveTheme = useMemo(() => {
    if (isPathExemptFromLightMode) {
      return ThemeMode.Dark;
    }
    return theme === ThemeMode.System
      ? systemPrefersDark
        ? ThemeMode.Dark
        : ThemeMode.Light
      : theme;
  }, [theme, systemPrefersDark, isPathExemptFromLightMode]);

  return (
    <ThemeContext.Provider
      value={{ ...parentThemeContext, theme, setTheme, effectiveTheme }}
    >
      {children}
    </ThemeContext.Provider>
  );
};

/**
 * Creates a container that uses the current theme classname
 */
export const ThemeContainer: PolymorphicComponent = <
  C extends React.ElementType = 'div',
>(
  props: PolymorphicComponentProps<C>
) => {
  const { as: Component = 'div', className, ...restProps } = props;
  const themeClassName = useThemeClassName();
  return (
    <Component className={twMerge(themeClassName, className)} {...restProps} />
  );
};

/**
 * Gets whether the current pathname should be exempt from light mode
 */
export function usePathExemptFromLightMode(override?: boolean) {
  const pathname = usePathname();
  return override ?? THEME_PATHS_EXEMPT_FROM_LIGHT_MODE_REGEX.test(pathname);
}

/**
 * Consume the ThemeContext
 */
export function useThemeContext() {
  const context = useContext(ThemeContext);
  if (context === undefined) {
    throw new Error('useTheme must be used within a ThemeProvider');
  }
  return context;
}

/**
 * Gets the current theme classname based on the closest ThemeProvider
 */
export function useThemeClassName() {
  const { effectiveTheme } = useThemeContext();
  return getThemeClassName(effectiveTheme);
}