import type * as babel from "@babel/core"
import type * as t from "@babel/types"

interface State {
  visited: WeakSet<t.ClassDeclaration>
  getTSLib(): t.ImportSpecifier[]
  __param(): t.Identifier
  __decorate(): t.Identifier
}

export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
  name: "tsc-decorators",
  manipulateOptions(_, { plugins }) {
    if (!plugins.includes("decorators-legacy")) {
      plugins.push("decorators-legacy")
    }
  },
  visitor: {
    Program: {
      enter(path, state) {
        const getTSLib = once(() => {
          const specifiers: t.ImportSpecifier[] = []
          const dec = t.importDeclaration(specifiers, t.stringLiteral("tslib"))
          path.node.body.unshift(dec)
          return specifiers
        })

        state.getTSLib = getTSLib

        const createGet = (name: string) =>
          once(() => {
            const specifiers = getTSLib()
            const localID = path.scope.hasBinding(name)
              ? path.scope.generateUidIdentifier(name)
              : t.identifier(name)
            specifiers.push(t.importSpecifier(localID, t.identifier(name)))
            return localID
          })

        const getParam = createGet("__param")
        const getDecorate = createGet("__decorate")

        state.__param = () => t.cloneNode(getParam())
        state.__decorate = () => t.cloneNode(getDecorate())
        state.visited = new WeakSet()
      },
    },

    ClassDeclaration(path, state) {
      if (state.visited.has(path.node)) return
      state.visited.add(path.node)

      const modifiers: (t.Expression | t.Statement)[] = []
      const classDecorators: t.Expression[] = []
      const prefixes: t.Statement[] = []

      const classID = path.node.id

      const isNamedExport = t.isExportNamedDeclaration(path.parent)
      const isDefaultExport = t.isExportDefaultDeclaration(path.parent)

      if (path.node.decorators) {
        classDecorators.push(...path.node.decorators.map(d => d.expression))
        path.node.decorators = []
      }

      for (const child of path.node.body.body) {
        switch (child.type) {
          case "ClassProperty":
            if (!child.decorators) break
            modifiers.push(
              t.callExpression(state.__decorate(), [
                t.arrayExpression(child.decorators.map(d => d.expression)),
                child.static
                  ? classID
                  : t.memberExpression(classID, t.identifier("prototype")),
                t.isIdentifier(child.key) ? t.stringLiteral(child.key.name) : child.key,
                t.unaryExpression("void", t.numericLiteral(0)),
              ])
            )
            child.decorators = null
            break

          case "ClassMethod":
            {
              const list: t.Expression[] = []
              if (child.decorators) {
                list.push(...child.decorators.map(d => d.expression))
                child.decorators = null
              }

              child.params.forEach((param, i) => {
                pseudoAssert(!t.isTSParameterProperty(param))
                if (!param.decorators?.length) return
                list.push(
                  ...param.decorators
                    .map(d => d.expression)
                    .map(e => t.callExpression(state.__param(), [t.numericLiteral(i), e]))
                )
                param.decorators = null
              })

              if (child.kind === "constructor") {
                classDecorators.push(...list)
              } else {
                let key: t.Expression
                if (child.computed) {
                  const newID = path.scope.generateUidIdentifier()
                  prefixes.push(
                    t.variableDeclaration("let", [t.variableDeclarator(newID)])
                  )
                  child.key = t.assignmentExpression("=", newID, child.key)
                  key = newID
                } else if (t.isIdentifier(child.key)) {
                  key = t.stringLiteral(child.key.name)
                } else {
                  key = child.key
                }

                if (!list.length) break

                modifiers.push(
                  t.callExpression(state.__decorate(), [
                    t.arrayExpression(list),
                    t.memberExpression(classID, t.identifier("prototype")),
                    key,
                    t.nullLiteral(),
                  ])
                )
              }
            }
            break
        }
      }

      if (!classDecorators.length && !modifiers.length && !prefixes.length) return

      const replacees: t.Statement[] = prefixes.slice()
      if (classDecorators.length) {
        replacees.push(
          t.variableDeclaration("let", [
            t.variableDeclarator(classID, { ...path.node, type: "ClassExpression" }),
          ])
        )
        modifiers.push(
          t.assignmentExpression(
            "=",
            classID,
            t.callExpression(state.__decorate(), [
              t.arrayExpression(classDecorators),
              classID,
            ])
          )
        )

        if (isNamedExport) {
          modifiers.push(
            t.exportNamedDeclaration(undefined, [t.exportSpecifier(classID, classID)])
          )
        } else if (isDefaultExport) {
          modifiers.push(t.exportDefaultDeclaration(classID))
        }
      } else {
        let node: t.Statement = path.node
        if (isNamedExport) {
          node = t.exportNamedDeclaration(node, [], null)
        } else if (isDefaultExport) {
          node = t.exportDefaultDeclaration(node)
        }
        replacees.push(node)
      }

      replacees.push(
        ...modifiers.map(e => (t.isStatement(e) ? e : t.expressionStatement(e)))
      )
      const replacementTarget = isNamedExport || isDefaultExport ? path.parentPath : path
      replacementTarget.replaceWithMultiple(replacees)
    },
  },
})

function pseudoAssert(condition: any): asserts condition {}

const once = <T extends Function>(fn: T): T => {
  let called = false
  let cache: any
  return function (this: any) {
    if (called) {
      return cache
    } else {
      cache = fn.apply(this, arguments)
      called = true
      return cache
    }
  } as any
}