use-tree.tsx

import * as React from "react"
import { useSnapshot } from "valtio"

import { TreeStateContext, TreeMapContext, createInitialTreeState } from "./contexts"
import { useServerComputedData } from "./server"
import { useIndex, useIndexedChildren } from "./use-indexed-children"
import { isServer, useIsomorphicLayoutEffect } from "./utils"

export function useTreeSnapshot<ComputedData extends any>(
  computeData?: (treeMap: Map<string, any>) => ComputedData,
  treeState?: TreeState | null,
  dependencies: React.DependencyList = []
) {
  const contextTreeState = React.useContext(TreeStateContext)
  const parsedTreeState = treeState || contextTreeState

  if (parsedTreeState === null) {
    throw "treeState must be defined"
  }

  const snapshot = useSnapshot(parsedTreeState)
  const serverComputedData = useServerComputedData(computeData)
  const clientComputedData = React.useMemo(
    () => (computeData ? computeData(snapshot.treeMap) : snapshot.treeMap),
    dependencies.concat(snapshot)
  )
  const initialRender = React.useRef(true)
  const computedData = (
    isServer || initialRender.current
      ? serverComputedData
      : clientComputedData || serverComputedData
  ) as ComputedData

  initialRender.current = false

  return computedData
}

/**
 * Control tree state from outside a component.
 *
 * @example
 * import type { TreeState } from "reforest"
 * import { useTree, useTreeData, useTreeState } from "reforest"
 *
 * function Item({ children, value }) {
 *   useTreeData(value)
 *   return <li>{children}</li>
 * }
 *
 * function ItemList({ children }: { children: React.ReactNode, treeState: TreeState }) {
 *   const tree = useTree(children, treeState)
 *   return <ul>{tree.children}</ul>
 * }
 *
 * function App() {
 *   const treeState = useTreeState()
 *   return (
 *     <ItemList treeState={treeState}>
 *       <Item value="apple">Apple</Item>
 *       <Item value="banana">Banana</Item>
 *       <Item value="cherry">Cherry</Item>
 *     </ItemList>
 *   )
 * }
 */
export function useTreeState() {
  const [treeState] = React.useState(() => createInitialTreeState())

  return treeState
}

export type TreeState = ReturnType<typeof useTreeState>

/**
 * Manage ordered data subscriptions for components.
 *
 * @example create a tree of data subscriptions
 * import { useTree, useTreeData } from "reforest"
 *
 * function Item({ children, value }) {
 *   useTreeData(value)
 *   return <li>{children}</li>
 * }
 *
 * function ItemList({ children }: { children: React.ReactNode }) {
 *   const tree = useTree(children)
 *   return <ul>{tree.children}</ul>
 * }
 */
export function useTree(children: React.ReactNode, parentTreeState?: TreeState) {
  const defaultTreeState = useTreeState()
  const [treeMap] = React.useState(() => new Map<string, Record<string, any>>())
  const treeStateContextValue = React.useContext(TreeStateContext)
  const parsedContextValue = treeStateContextValue || parentTreeState || defaultTreeState
  const isRoot = treeStateContextValue === null
  const indexedChildren = useIndexedChildren(children)
  const childrenToRender = isRoot ? (
    <TreeStateContext.Provider value={parsedContextValue}>
      <TreeMapContext.Provider value={treeMap}>{indexedChildren}</TreeMapContext.Provider>
    </TreeStateContext.Provider>
  ) : (
    indexedChildren
  )

  return {
    children: childrenToRender,
    isRoot,
  }
}

/** Subscribe data to the root useTree hook. */
export function useTreeData<TreeValue extends any>(data: TreeValue) {
  const treeState = React.useContext(TreeStateContext)
  const treeMapContext = React.useContext(TreeMapContext)
  const treeId = React.useId().slice(1, -1)

  if (treeState === null) {
    throw new Error("useTreeData must be used in a descendant component of useTree.")
  }

  const index = useIndex()!
  const indexPathString = index.indexPathString

  /** Subscribe tree data to root map on the server and client. */
  if (isServer) {
    treeMapContext.set(treeId, Object.assign({ indexPathString, treeId }, data))
  }

  useIsomorphicLayoutEffect(() => {
    return treeState.subscribeTreeData(treeId, Object.assign({ indexPathString, treeId }, data))
  }, [treeState, treeId, data])

  return treeId
}