import type { NodePath, types } from "@babel/core"
import type { Node } from "@babel/core"
import type { BabelPlugin } from "./esbuild-babel"

// useWhyDidYouUpdate
export const whyDidYouRender =
  ({
    hookName,
    hookPath,
    ignoredHooks,
  }: {
    hookName: string
    hookPath: string
    ignoredHooks: string[]
  }): BabelPlugin =>
  ({ types: t }) => {
    const ignored = new WeakSet<Node>()

    function ignore(node: Node) {
      ignored.add(node)
      return node
    }

    return {
      name: "why-did-you-render",
      visitor: {
        Program(path, state) {
          const id = path.scope.generateUidIdentifier(hookName)
          path.node.body.unshift(
            t.importDeclaration(
              [t.importSpecifier(id, t.identifier(hookName))], //
              t.stringLiteral(hookPath),
            ),
          )
          state.whyDidYouRenderId = id
        },

        VariableDeclaration(path, state) {
          if (ignored.has(path.node)) return

          const decls = path.node.declarations
          if (decls.length !== 1) return

          const [{ init, id }] = decls
          if (
            !t.isCallExpression(init) ||
            !t.isIdentifier(init.callee) ||
            !init.callee.name.startsWith("use") ||
            init.callee.name.length <= 3
          ) {
            return
          }

          if (ignoredHooks.includes(init.callee.name)) {
            return
          }

          const findParent = <T extends types.Node>(
            predicate: (node: types.Node) => node is T,
          ) => path.findParent(path => predicate(path.node)) as NodePath<T> | undefined

          const parentId =
            findParent(t.isFunctionDeclaration)?.node.id!.name ??
            (<types.Identifier>(
              (<types.VariableDeclarator>findParent(t.isArrowFunctionExpression)?.parent)
                ?.id
            ))?.name

          if (!parentId || parentId.startsWith("use")) {
            return
          }

          const callee = t.cloneNode(state.whyDidYouRenderId as types.Identifier)

          if (t.isIdentifier(id)) {
            path.insertAfter(
              t.callExpression(callee, [
                t.stringLiteral(parentId),
                t.stringLiteral(init.callee.name),
                id,
              ]),
            )
            return
          }

          const temporaryId = path.scope.generateUidIdentifier(init.callee.name)

          path.replaceWithMultiple([
            ignore(
              t.variableDeclaration(path.node.kind, [
                t.variableDeclarator(temporaryId, init),
                t.variableDeclarator(id, temporaryId),
              ]),
            ),
            t.expressionStatement(
              t.callExpression(callee, [
                t.stringLiteral(parentId),
                t.stringLiteral(init.callee.name),
                temporaryId,
              ]),
            ),
          ])
        },
      },
    }
  }
