import { Session, StytchClient, User } from '@stytch/stytch-js';
import React, { ComponentType, createContext, ReactNode, useContext, useEffect, useState } from 'react';

import { useAsyncState } from 'utils/async';
import { lazyError, noProviderError } from 'utils/errors';
import { invariant } from 'utils/invariant';

type StytchMountedContextType = {
  isMounted: boolean;
  isLazy: boolean;
};

const StytchMountedContext = createContext<StytchMountedContextType>({
  isLazy: false,
  isMounted: false,
});
const StytchUserContext = createContext<User | null>(null);
const StytchSessionContext = createContext<Session | null>(null);
const StytchContext = createContext<StytchClient | null>(null);

export const useIsMounted__INTERNAL = (): boolean => useContext(StytchMountedContext).isMounted;
export const useIsLazy__INTERNAL = (): boolean => useContext(StytchMountedContext).isLazy;
export const useStytchRaw__INTERNAL = (): StytchClient | null => useContext(StytchContext);

export const useStytchUser = (): User | null => {
  invariant(useIsMounted__INTERNAL(), noProviderError('useStytchUser'));
  return useContext(StytchUserContext);
};
export const useStytchSession = (): Session | null => {
  invariant(useIsMounted__INTERNAL(), noProviderError('useStytchSession'));
  return useContext(StytchSessionContext);
};
export const useStytch = (): StytchClient => {
  invariant(useIsMounted__INTERNAL(), noProviderError('useStytch'));
  invariant(!useIsLazy__INTERNAL(), lazyError('useStytch'));
  return useContext(StytchContext);
};
export const useStytchLazy = (): StytchClient => {
  invariant(useIsMounted__INTERNAL(), noProviderError('useStytchLazy'));
  return useContext(StytchContext);
};

export const withStytch = <T extends object>(
  Component: ComponentType<T & { stytch: StytchClient }>,
): ComponentType<T> => {
  const WithStytch: ComponentType<T> = (props) => {
    invariant(useIsMounted__INTERNAL(), noProviderError('withStytch'));
    invariant(!useIsLazy__INTERNAL(), lazyError('withStytch'));
    return <Component {...props} stytch={useStytch()} />;
  };
  WithStytch.displayName = `withStytch(${Component.displayName || Component.name || 'Component'})`;
  return WithStytch;
};
export const withStytchLazy = <T extends object>(
  Component: ComponentType<T & { stytch: StytchClient | null }>,
): ComponentType<T> => {
  const WithStytchLazy: ComponentType<T> = (props) => {
    invariant(useIsMounted__INTERNAL(), noProviderError('withStytchLazy'));
    return <Component {...props} stytch={useStytchLazy()} />;
  };
  WithStytchLazy.displayName = `withStytchLazy(${Component.displayName || Component.name || 'Component'})`;
  return WithStytchLazy;
};
export const withStytchUser = <T extends object>(
  Component: ComponentType<T & { stytchUser: User | null }>,
): ComponentType<T> => {
  const WithStytchUser: ComponentType<T> = (props) => {
    invariant(useIsMounted__INTERNAL(), noProviderError('withStytchUser'));
    return <Component {...props} stytchUser={useStytchUser()} />;
  };
  WithStytchUser.displayName = `withStytchUser(${Component.displayName || Component.name || 'Component'})`;
  return WithStytchUser;
};
export const withStytchSession = <T extends object>(
  Component: ComponentType<T & { stytchSession: Session | null }>,
): ComponentType<T> => {
  const WithStytchSession: ComponentType<T> = (props) => {
    invariant(useIsMounted__INTERNAL(), noProviderError('withStytchSession'));
    return <Component {...props} stytchSession={useStytchSession()} />;
  };
  WithStytchSession.displayName = `withStytchSession(${Component.displayName || Component.name || 'Component'})`;
  return WithStytchSession;
};

export type StytchProviderProps = {
  stytch: StytchClient | null | Promise<StytchClient | null>;
  children?: ReactNode;
};

export const StytchProvider = ({ stytch, children }: StytchProviderProps): JSX.Element => {
  invariant(!useIsMounted__INTERNAL(), 'You cannot render a <StytchProvider> inside another <StytchProvider>.');
  const [stytchClient, setStytchClient] = useState<StytchClient | null>(
    stytch instanceof Promise || !stytch ? null : stytch,
  );
  const [mountedContext] = useState<StytchMountedContextType>({
    isLazy: !stytchClient,
    isMounted: true,
  });
  const [user, setUser] = useAsyncState<User | null>(stytchClient && stytchClient.user.getSync());
  const [session, setSession] = useAsyncState<Session | null>(stytchClient && stytchClient.session.getSync());

  useEffect(() => {
    if (stytch instanceof Promise) {
      stytch.then(setStytchClient);
    }
  }, [stytch]);

  useEffect(() => {
    if (stytchClient) {
      if (mountedContext.isLazy) {
        setUser(stytchClient.user.getSync());
        setSession(stytchClient.session.getSync());
      }
      const unsubscribeUser = stytchClient.user.onChange(setUser);
      const unsubscribeSession = stytchClient.session.onChange(setSession);
      return () => {
        unsubscribeUser();
        unsubscribeSession();
      };
    }
  }, [mountedContext, stytchClient, setUser, setSession]);

  return (
    <StytchMountedContext.Provider value={mountedContext}>
      <StytchContext.Provider value={stytchClient}>
        <StytchUserContext.Provider value={session && user}>
          <StytchSessionContext.Provider value={user && session}>{children}</StytchSessionContext.Provider>
        </StytchUserContext.Provider>
      </StytchContext.Provider>
    </StytchMountedContext.Provider>
  );
};