// import { relative, resolve } from "path";
import type { PluginBuild } from "esbuild"
import type * as babel from "@babel/core"
import type { NodePath, types } from "@babel/core"
import { expression } from "@babel/template"
import hash from "@emotion/hash"
import { compile, serialize, stringify } from "stylis"
import { dropRightWhile } from "lodash"

const NS = "extract-css"

const runtimeArgs = ["Tag", "className", "vars"] as const
// const root = resolve(__dirname, "../..");

const shared = /* jsx */ `
  ({ as = Tag, className: cls, style, ...rest }, ref) => {
    return _jsx(as, {
      ...rest,
      ref,
      style: vars != null ? getStyle({ ...style }, className, vars) : style,
      className: cx([className, cls])
    })
  }
`

const runtime = /* jsx */ `
  import { jsx } from "react/jsx-runtime";
  import { forwardRef, memo } from "react";
  import { cx } from "@emotion/css";

  const supportsAs = new WeakSet();
  const _jsx = jsx;

  function getStyle(style, className, vars) {
    for (let i = 0; i < vars.length; i++) {
      const name = "--" + className + "-" + i;
      const variable = vars[i];
      if (variable !== null) {
        style[name] = typeof variable === "function" ? variable(props) : variable;
      }
    }
    return style;
  }

  export default function create(${runtimeArgs.join(", ")}) {
    const Component = memo(
      forwardRef(
        supportsAs.has(Tag)
          ? ${shared.replace("as, {", "Tag, { as,")}
          : ${shared}
      )
    );

    Component.withComponent = function (Tag) {
      return create(${runtimeArgs.join(", ")});
    };

    supportsAs.add(Component);

    return Component;
  }
`

const flattenArgs = (args: { [key in (typeof runtimeArgs)[number]]: string }) =>
  runtimeArgs.map(key => args[key]).join(", ")

const styledComponent = expression({
  // plugins: ["jsx"],
  syntacticPlaceholders: true,
})(
  "%%create%%(" +
    flattenArgs({
      className: "%%className%%",
      Tag: "%%tag%%",
      vars: "%%vars%%",
    }) +
    ")",
)

const validParentTypes = new Set<types.Node["type"]>([
  "ArrayExpression",
  "BinaryExpression",
  "CallExpression",
  "JSXExpressionContainer",
  "LogicalExpression",
  "ObjectProperty",
  "VariableDeclarator",
])

interface State extends babel.PluginPass {
  styles: string[]
  program: types.Program
  runtimeHelper?: types.Identifier
}

export function extractCSS(build: PluginBuild, { className }: { className: string }) {
  const RUNTIME = /^@extract-css\/runtime$/

  build.onResolve({ filter: RUNTIME }, ({ path }) => ({ namespace: NS, path }))

  build.onLoad({ filter: RUNTIME, namespace: NS }, () => ({
    contents: runtime,
    loader: "jsx",
    resolveDir: __dirname,
  }))

  const plugin = ({ types: t }: typeof babel): babel.PluginObj<State> => {
    const classNameMap = new WeakMap<types.Expression, string>()

    function isImportedSpecifierFrom(path: NodePath, name: string, module: string) {
      const declPath = path.scope.getBinding(name)!.path
      return (
        t.isImportSpecifier(declPath.node) &&
        t.isImportDeclaration(declPath.parent) &&
        declPath.parent.source.value === module
      )
    }

    function isDefaultImportedFrom(path: NodePath, name: string, module: string) {
      const binding = path.scope.getBinding(name)
      if (!binding) return false
      const declPath = binding.path
      return (
        t.isImportDefaultSpecifier(declPath.node) &&
        t.isImportDeclaration(declPath.parent) &&
        declPath.parent.source.value === module
      )
    }

    function buildClassName(hash: string, componentName = "styled") {
      return className.replaceAll("[hash]", hash).replaceAll("[name]", componentName)
    }

    function getComponentName(path: NodePath): string | undefined {
      if (t.isVariableDeclarator(path.parent) && t.isIdentifier(path.parent.id)) {
        return path.parent.id.name
      }
    }

    const visitor: babel.Visitor<State> = {
      TaggedTemplateExpression(path, state) {
        const { styles } = state
        const quasiExps = path
          .get("quasi")
          .get("expressions") as NodePath<types.Expression>[]

        const { tag, quasi } = path.node

        if (!validParentTypes.has(path.parent.type)) {
          return
        }

        function extract(getPrefix: (className: string) => string) {
          let bailout = false
          const raws = quasi.quasis.map(q => q.value.raw)
          const className = buildClassName(
            hash(JSON.stringify(raws)),
            getComponentName(path),
          )
          const skipInterpolations = new Set<number>()

          const cssText = raws
            .map((left, i) => {
              if (bailout) {
                return null!
              }

              if (i === raws.length - 1) {
                return left
              }

              const exp = quasi.expressions[i]
              const evaluated = quasiExps[i].evaluate()
              if (evaluated.confident) {
                if (
                  typeof evaluated.value === "string" ||
                  typeof evaluated.value === "number"
                ) {
                  skipInterpolations.add(i)
                  return left + String(evaluated.value)
                }
              }

              if (t.isIdentifier(exp)) {
                const binding = path.scope.getBinding(exp.name)
                if (binding) {
                  const { node } = binding.path
                  if (t.isVariableDeclarator(node)) {
                    const cls = classNameMap.get(node.init!)
                    if (cls) {
                      skipInterpolations.add(i)
                      return left + cls
                    }
                    if (t.isStringLiteral(node.init)) {
                      skipInterpolations.add(i)
                      return left + node.init.value
                    }
                  }

                  bailout = true
                }
              } else if (t.isStringLiteral(exp)) {
                skipInterpolations.add(i)
                return left + exp.value
              }

              return `${left}var(--${className}-${i})`
            })
            .join("")

          const compiled = compile(`${getPrefix(className)}{${cssText}}}`)

          return {
            className,
            skipInterpolations,
            bailout,
            style: serialize(compiled, stringify),
          }
        }

        function processStyled(tag: types.Expression) {
          const { className, skipInterpolations, style, bailout } = extract(c => `.${c}`)
          if (bailout) {
            return
          }

          styles.push(style)

          if (!state.runtimeHelper) {
            state.runtimeHelper = path.scope.generateUidIdentifier("create")
            state.program.body.unshift(
              t.importDeclaration(
                [t.importDefaultSpecifier(state.runtimeHelper)],
                t.stringLiteral("@extract-css/runtime"),
              ),
            )
          }

          const fn = styledComponent({
            create: state.runtimeHelper,
            className: t.stringLiteral(className),
            tag:
              t.isIdentifier(tag) && /^[a-z]/.test(tag.name)
                ? t.stringLiteral(tag.name)
                : tag,
            vars: quasi.expressions.length
              ? t.arrayExpression(
                  dropRightWhile(
                    quasi.expressions.map((e, i) =>
                      skipInterpolations.has(i)
                        ? t.nullLiteral()
                        : (e as types.Expression),
                    ),
                    n => t.isNullLiteral(n),
                  ),
                )
              : t.nullLiteral(),
          })
          classNameMap.set(fn, className)
          const [newPath] = path.replaceWith(fn)
          newPath.addComment("leading", " @__PURE__ ")
        }

        if (t.isIdentifier(tag) && tag.name === "keyframes") {
          if (!isImportedSpecifierFrom(path, tag.name, "@emotion/css")) {
            return
          }

          const { className, style, bailout } = extract(c => `@keyframes ${c}`)
          if (bailout) {
            return
          }
          styles.push(style)
          path.replaceWith(t.stringLiteral(className))
        } else if (t.isIdentifier(tag) && tag.name === "css") {
          if (!isImportedSpecifierFrom(path, tag.name, "@emotion/css")) {
            return
          }

          const { className, style, bailout } = extract(c => `.${c}`)
          if (bailout) {
            return
          }
          styles.push(style)
          path.replaceWith(t.stringLiteral(className))
        } else if (
          t.isMemberExpression(tag) &&
          t.isIdentifier(tag.object, { name: "styled" }) &&
          !t.isPrivateName(tag.property)
        ) {
          if (!isDefaultImportedFrom(path, tag.object.name, "@emotion/styled")) {
            return
          }
          processStyled(tag.property)
        } else if (
          t.isCallExpression(tag) &&
          t.isIdentifier(tag.callee, { name: "styled" }) &&
          tag.arguments.length === 1 &&
          t.isExpression(tag.arguments[0])
        ) {
          if (!isDefaultImportedFrom(path, tag.callee.name, "@emotion/styled")) {
            return
          }
          processStyled(tag.arguments[0])
        }
      },
    }

    return {
      name: "plugin",
      visitor: {
        Program: {
          enter(program, state) {
            state.styles = []
            state.program = program.node
          },
          exit(program, state) {
            program.traverse(visitor, state)

            if (state.styles.length) {
              const css =
                // `/*./${relative(root, state.filename!)}*/` +
                state.styles.join("\n")
              program.node.body.unshift(
                t.importDeclaration(
                  [],
                  t.stringLiteral(`data:text/css,${encodeURI(css)}`),
                ),
              )
            }
          },
        },
      },
    }
  }

  return plugin
}