// https://github.com/babel/babel/blob/c38bf12f010520ea7abe8a286f62922b2d1e1f1b/packages/babel-plugin-transform-react-constant-elements/src/index.ts
import type { Visitor } from "@babel/core"
import type { PluginObj, types as t } from "@babel/core"

interface PluginState {
  isImmutable: boolean
  mutablePropsAllowed?: boolean
}

export default (allowMutablePropsOnTags?: string[]): PluginObj => {
  const HOISTED = new WeakSet()

  const immutabilityVisitor: Visitor<PluginState> = {
    enter(path, state) {
      const stop = () => {
        state.isImmutable = false
        path.stop()
      }

      if (path.isJSXClosingElement()) {
        path.skip()
        return
      }

      // Elements with refs are not safe to hoist.
      if (
        path.isJSXIdentifier({ name: "ref" }) &&
        path.parentPath.isJSXAttribute({ name: path.node })
      ) {
        return stop()
      }

      // Ignore identifiers & JSX expressions.
      if (path.isJSXIdentifier() || path.isIdentifier() || path.isJSXMemberExpression()) {
        return
      }

      if (!path.isImmutable()) {
        // If it's not immutable, it may still be a pure expression, such as string concatenation.
        // It is still safe to hoist that, so long as its result is immutable.
        // If not, it is not safe to replace as mutable values (like objects) could be mutated after render.
        // https://github.com/facebook/react/issues/3226
        if (path.isPure()) {
          const { confident, value } = path.evaluate()
          if (confident) {
            // We know the result; check its mutability.
            const isMutable =
              (!state.mutablePropsAllowed && value && typeof value === "object") ||
              typeof value === "function"
            if (!isMutable) {
              // It evaluated to an immutable value, so we can hoist it.
              path.skip()
              return
            }
          }
        }
        stop()
      }
    },
  }

  return {
    name: "transform-react-constant-elements",

    visitor: {
      Program(program) {
        program.traverse({
          JSXElement(path) {
            if (HOISTED.has(path.node)) return
            HOISTED.add(path.node)

            const state: PluginState = { isImmutable: true }

            // This transform takes the option `allowMutablePropsOnTags`, which is an array
            // of JSX tags to allow mutable props (such as objects, functions) on. Use sparingly
            // and only on tags you know will never modify their own props.
            if (allowMutablePropsOnTags != null) {
              // Get the element's name. If it's a member expression, we use the last part of the path.
              // So the option ["FormattedMessage"] would match "Intl.FormattedMessage".
              let namePath = path.get("openingElement").get("name")
              while (namePath.isJSXMemberExpression()) {
                namePath = namePath.get("property")
              }

              const elementName = (namePath.node as t.JSXIdentifier).name
              state.mutablePropsAllowed = allowMutablePropsOnTags.includes(elementName)
            }

            // Traverse all props passed to this element for immutability.
            path.traverse(immutabilityVisitor, state)

            if (state.isImmutable) {
              path.hoist(path.scope)
            }
          },
        })
      },
    },
  }
}