react.tsx

import React, { createContext, FC, useContext, useEffect, useMemo, useRef, useState } from 'react'
import Emitter from './utils/emitter'
const events = new Emitter()

import { IFlagsmith, IFlagsmithTrait, IFlagsmithFeature, IState } from './types'

export const FlagsmithContext = createContext<IFlagsmith<string, string> | null>(null)
export type FlagsmithContextType = {
    flagsmith: IFlagsmith // The flagsmith instance
    options?: Parameters<IFlagsmith['init']>[0] // Initialisation options, if you do not provide this you will have to call init manually
    serverState?: IState
    children: React.ReactNode
}

export const FlagsmithProvider: FC<FlagsmithContextType> = ({ flagsmith, options, serverState, children }) => {
    const firstRenderRef = useRef(true)
    if (flagsmith && !flagsmith?._trigger) {
        flagsmith._trigger = () => {
            // @ts-expect-error using internal function, consumers would never call this
            flagsmith?.log('React - trigger event received')
            events.emit('event')
        }
    }

    if (flagsmith && !flagsmith?._triggerLoadingState) {
        flagsmith._triggerLoadingState = () => {
            events.emit('loading_event')
        }
    }

    if (serverState && !flagsmith.initialised) {
        flagsmith.setState(serverState)
    }

    if (firstRenderRef.current) {
        firstRenderRef.current = false
        if (options) {
            flagsmith
                .init({
                    ...options,
                    state: options.state || serverState,
                    onChange: (...args) => {
                        if (options.onChange) {
                            options.onChange(...args)
                        }
                    },
                })
                .catch((error) => {
                    // @ts-expect-error using internal function, consumers would never call this
                    flagsmith?.log('React - Failed to initialize flagsmith', error)
                    events.emit('event')
                })
        }
    }
    return <FlagsmithContext.Provider value={flagsmith}>{children}</FlagsmithContext.Provider>
}

const useConstant = function <T>(value: T): T {
    const ref = useRef(value)
    if (!ref.current) {
        ref.current = value
    }
    return ref.current
}

const flagsAsArray = (_flags: any): string[] => {
    if (typeof _flags === 'string') {
        return [_flags]
    } else if (typeof _flags === 'object') {
        // eslint-disable-next-line no-prototype-builtins
        if (_flags.hasOwnProperty('length')) {
            return _flags
        }
    }
    throw new Error('Flagsmith: please supply an array of strings or a single string of flag keys to useFlags')
}

const normalizeFlagKey = (key: string) => key.toLowerCase().replace(/ /g, '_')

const getRenderKey = (flagsmith: IFlagsmith, flags: string[], traits: string[] = []) => {
    return flags
        .map((k) => {
            return `${flagsmith.getValue(k)}${flagsmith.hasFeature(k)}${flagsmith.getAllFlags()?.[normalizeFlagKey(k)]?.variant}`
        })
        .concat(traits.map((t) => `${flagsmith.getTrait(t)}`))
        .join(',')
}

const getExperimentRenderKey = (flagsmith: IFlagsmith | null, key: string): string => {
    const flag = flagsmith?.getAllFlags()?.[key]
    const identifier = flagsmith?.getContext().identity?.identifier ?? null
    // Identity is part of the key so that switching identity re-renders (and
    // re-fires the exposure) even when the resolved value is unchanged.
    return `${identifier}|${flag?.value}|${flag?.enabled}|${flag?.variant}`
}

export function useFlagsmithLoading() {
    const flagsmith = useContext(FlagsmithContext)
    const [loadingState, setLoadingState] = useState(flagsmith?.loadingState)

    useEffect(() => {
        if (!flagsmith) return
        setLoadingState(flagsmith.loadingState)
        const unsubscribe = events.on('loading_event', () => {
            setLoadingState(flagsmith.loadingState)
        })
        return () => {
            unsubscribe()
        }
    }, [flagsmith])

    return loadingState
}

type UseFlagsReturn<F extends string | Record<string, any>, T extends string> = F extends string
    ? {
          [K in F]: IFlagsmithFeature
      } & {
          [K in T]: IFlagsmithTrait
      }
    : {
          [K in keyof F]: IFlagsmithFeature<F[K]>
      } & {
          [K in T]: IFlagsmithTrait
      }

/**
 * Example usage:
 *
 * // A) Using string flags:
 * useFlags<"featureOne"|"featureTwo">(["featureOne", "featureTwo"]);
 *
 * // B) Using an object for F - this can be generated by our CLI: https://github.com/Flagsmith/flagsmith-cli :
 * interface MyFeatureInterface {
 *   featureOne: string;
 *   featureTwo: number;
 * }
 * useFlags<MyFeatureInterface>(["featureOne", "featureTwo"]);
 */
export function useFlags<F extends string | Record<string, any>, T extends string = string>(
    _flags: readonly (F | keyof F)[],
    _traits: readonly T[] = []
) {
    const flags = useConstant<string[]>(flagsAsArray(_flags))
    const traits = useConstant<string[]>(flagsAsArray(_traits))
    const flagsmith = useContext(FlagsmithContext)
    const [renderRef, setRenderRef] = useState(getRenderKey(flagsmith as IFlagsmith, flags, traits))

    useEffect(() => {
        if (!flagsmith) return
        setRenderRef(getRenderKey(flagsmith, flags, traits))
        const unsubscribe = events.on('event', () => {
            setRenderRef((prev) => {
                const next = getRenderKey(flagsmith, flags, traits)
                if (prev === next) return prev
                // @ts-expect-error using internal function, consumers would never call this
                flagsmith?.log('React - useFlags flags and traits have changed')
                return next
            })
        })
        return () => {
            unsubscribe()
        }
    }, [flagsmith, flags, traits])

    const res = useMemo(() => {
        const res: any = {}
        flags
            .map((k) => {
                const variant = flagsmith!.getAllFlags()?.[normalizeFlagKey(k)]?.variant
                res[k] = {
                    enabled: flagsmith!.hasFeature(k),
                    value: flagsmith!.getValue(k),
                    ...(variant != null ? { variant } : {}),
                }
            })
            .concat(
                traits?.map((v) => {
                    res[v] = flagsmith!.getTrait(v)
                })
            )
        return res
    }, [renderRef])

    return res as UseFlagsReturn<F, T>
}

/**
 * Resolve an experiment flag for the currently identified user and record a
 * single `$flag_exposure` event as a side-effect. Re-renders when the flag's
 * value or enabled state changes. When events are disabled (enableEvents is
 * not set) the flag is still returned but no exposure is recorded.
 *
 * Exposures are gated three ways: the effect only runs when the flag value,
 * variant, identity, feature or source change; a ref guard prevents duplicate
 * fires for the same (feature, identifier, value, variant); and the core
 * EventProcessor dedupes within each flush window. Frequent re-renders
 * therefore never amplify into extra events.
 *
 * @experimental @internal
 */
export function useExperiment(featureName: string): IFlagsmithFeature | null {
    const flagsmith = useContext(FlagsmithContext)
    const key = normalizeFlagKey(featureName)
    const lastExposureKey = useRef<string | null>(null)
    const [, setRenderKey] = useState<string>(() => getExperimentRenderKey(flagsmith, key))

    useEffect(() => {
        const listener = () => {
            const next = getExperimentRenderKey(flagsmith, key)
            setRenderKey((prev) => (prev !== next ? next : prev))
        }
        const off = events.on('event', listener)
        listener() // capture any change between first render and subscription
        return () => {
            off()
        }
    }, [flagsmith, key])

    const flag = (flagsmith?.getAllFlags()?.[key] as IFlagsmithFeature | undefined) ?? null
    const identifier = flagsmith?.getContext().identity?.identifier ?? null

    useEffect(() => {
        if (!flagsmith?.eventsEnabled || !flag) {
            return
        }
        const exposureKey = `${key}:${identifier}:${flag.value}:${flag.variant}`
        if (lastExposureKey.current === exposureKey) {
            return
        }
        lastExposureKey.current = exposureKey
        flagsmith.getExperimentFlag(featureName)
        // eslint-disable-next-line react-hooks/exhaustive-deps
    }, [flagsmith, featureName, key, identifier, flag?.value, flag?.enabled, flag?.variant])

    return flag
}

export function useFlagsmith<F extends string | Record<string, any>, T extends string = string>() {
    const context = useContext(FlagsmithContext)

    if (!context) {
        throw new Error('useFlagsmith must be used with in a FlagsmithProvider')
    }

    return context as unknown as IFlagsmith<F, T>
}