import {
    animate,
    AnimationOptions,
    HTMLMotionProps,
    motion,
    MotionConfigContext,
    useIsPresent,
} from "motion/react"
import {
    CSSProperties,
    forwardRef,
    useContext,
    useEffect,
    useImperativeHandle,
    useLayoutEffect,
    useRef,
    useState,
} from "react"
import { useIsInitialRender } from "../hooks/use-is-initial-render"
import { Em } from "../types"
import { getWidthInEm } from "../utils/get-width-in-ems"
import { targetWidths } from "../utils/target-widths"
import { maskHeight } from "./Mask"

export const NumberDigit = forwardRef<
    HTMLSpanElement,
    Omit<HTMLMotionProps<"span">, "children"> & {
        value: number
        initialValue?: number
    }
>(function NumberDigit(
    { value: _value, initialValue: _initialValue = _value, ...rest },
    _ref
) {
    const { transition } = useContext(MotionConfigContext)
    const initialValue = useRef(_initialValue).current // non-reactive, like React's defaultValue props
    const isInitialRender = useIsInitialRender()

    const scope = useRef<HTMLSpanElement>(null)
    const ref = useRef<HTMLSpanElement>(null)
    useImperativeHandle(_ref, () => ref.current!, [])

    const numberRefs = useRef(new Array<HTMLSpanElement | null>(10))

    // Don't use a normal exit animation for this because we want it to trigger a resize:
    const isPresent = useIsPresent()
    const value = isPresent ? _value : 0

    // Set the width to the width of the initial value immediately, so on the next render we animate from that:
    useLayoutEffect(() => {
        if (!scope.current || !numberRefs.current[initialValue]) return
        scope.current.style.width = getWidthInEm(
            numberRefs.current[initialValue]
        )
    }, [])

    // Animate the y in a layout effect, because it's a FLIP
    const prevValue = useRef(_initialValue)
    useLayoutEffect(() => {
        if (!scope.current || value === prevValue.current) return
        const box = scope.current.getBoundingClientRect()
        const refBox = ref.current?.getBoundingClientRect()

        // Using a number seems to ensure Motion ends with "none", which we want:
        // Add the offset between the top of the inner and outer elements to account for
        // any current animation state:

        const initialY =
            box.height * (value - prevValue.current) +
            (box.top - (refBox ? refBox.top || 0 : box.top))

        animate(
            scope.current,
            { y: [initialY, 0] },
            transition as AnimationOptions
        )

        return () => {
            prevValue.current = value
        }
    }, [value])

    const [width, setWidth] = useState<Em>()
    useEffect(() => {
        // Skip setting the width if this is the first render and it's not going to animate:
        if (isInitialRender && initialValue === value) return
        if (!numberRefs.current[value]) return
        const w = getWidthInEm(numberRefs.current[value])
        // Store the target width immediately, so it can be used for the section resize:
        if (ref.current) targetWidths.set(ref.current, w)
        // Trigger the actual layout animation by causing another render:
        setWidth(w)
    }, [value])

    const renderNumber = (i: number) => (
        <span
            key={i}
            style={{
                display: "inline-block",
                padding: `calc(${maskHeight}/2) 0`,
            }}
            ref={(r) => void (numberRefs.current[i] = r)}
        >
            {i}
        </span>
    )

    return (
        <motion.span
            {...rest}
            ref={ref}
            layout="position"
            data-state={isPresent ? undefined : "exiting"}
            style={{
                display: "inline-flex",
                justifyContent: "center",
                width,
            }}
        >
            <span
                ref={scope}
                style={{
                    display: "inline-flex",
                    justifyContent: "center",
                    flexDirection: "column",
                    alignItems: "center",
                    position: "relative",
                }}
            >
                {value !== 0 && (
                    <span
                        style={{
                            ...digitFillStyle,
                            bottom: `100%`,
                            left: 0,
                        }}
                    >
                        {new Array(value)
                            .fill(null)
                            .map((_, i) => renderNumber(i))}
                    </span>
                )}
                {renderNumber(value)}
                {value !== 9 && (
                    <span
                        style={{
                            ...digitFillStyle,
                            top: `100%`,
                            left: 0,
                        }}
                    >
                        {new Array(9 - value)
                            .fill(null)
                            .map((_, i) => renderNumber(value + i + 1))}
                    </span>
                )}
            </span>
        </motion.span>
    )
})

const digitFillStyle: CSSProperties = {
    display: "flex",
    flexDirection: "column",
    alignItems: "center",
    position: "absolute",
    width: "100%",
}