babel.ts

import { declare } from '@babel/helper-plugin-utils'
import { types as t } from '@babel/core'
import type {
  ImportDeclaration,
  TaggedTemplateExpression,
  Program,
  CallExpression,
} from '@babel/types'
import type { NodePath } from '@babel/core'

import { generateExtractableModule } from './module-generator'
import { executeModule } from './module-executor'
import { compileCss } from './css-compiler'

interface State {
  imports: string[]
  nodes: NodePath<TaggedTemplateExpression>[]
  extractable: string[]
  xxx: any[]
}

const processed = new Set<string>()

export default declare((api, opts) => {
  api.assertVersion(7)

  const isVirtualModuleCaller = api.caller((caller) => caller?.name === 'virtual-module-evaluator')

  const options = Object.assign(opts, {
    allowedModules: ['@steely/core', '@steely/react'],
    allowedMethods: ['css', 'keyframes', 'createGlobalStyle', 'createTheme'],
  })

  const mapper = new Map<string, string>([
    ['@steely/core', '@steely/core/lib/node'],
    ['@steely/react', '@steely/react/lib/node'],
  ])

  function collectImports(path: NodePath<ImportDeclaration>, state: State) {
    if (!options.allowedModules.includes(path.node.source.value)) {
      return
    }

    for (const specifier of path.node.specifiers) {
      if (options.allowedMethods.includes(specifier.local.name)) {
        state.imports.push(specifier.local.name)
      }
    }
  }

  // TODO: rename
  function collectExtractable(path: NodePath<TaggedTemplateExpression>, state: State) {
    if (t.isIdentifier(path.node.tag) && state.imports.includes(path.node.tag.name)) {
      state.nodes.push(path)
      // @ts-expect-error (TODO: Fix this case)
      state.extractable.push(path.parentPath.node.id.name)
    }
  }

  function collectExtractableTheme(path: NodePath<CallExpression>, state: State) {
    if (t.isIdentifier(path.node.callee) && state.imports.includes(path.node.callee.name)) {
      const xxx = { kind: 'theme', node: path, vars: [] as string[] }

      if (t.isVariableDeclarator(path.parentPath.node)) {
        if (t.isObjectPattern(path.parentPath.node.id)) {
          for (const prop of path.parentPath.node.id.properties) {
            if (t.isObjectProperty(prop) && t.isIdentifier(prop.value)) {
              xxx.vars.push(prop.value.name)
            }
          }
        }
      }

      state.xxx.push(xxx)
    }
  }

  function extractStyles(path: NodePath<Program>, state: State) {
    const code = generateExtractableModule(path, state.extractable)

    try {
      // @ts-expect-error
      const extractable = executeModule(code, state.file.opts.filename, mapper)

      for (let i = 0; i < extractable.length; i++) {
        const chunk = extractable[i]
        const className = chunk.className
        const css = compileCss(chunk.css)

        state.nodes[i].replaceWith(
          t.objectExpression([
            t.objectProperty(t.identifier('css'), t.stringLiteral(css)),
            t.objectProperty(t.identifier('className'), t.stringLiteral(className)),
          ]),
        )
      }
    } catch (error) {
      console.log(error)
    }
  }

  function extract(path: NodePath<Program>, state: State) {
    const exports = state.xxx.flatMap((xxx) => xxx.vars)

    // const f =

    const code = generateExtractableModule(path, exports)
    // console.log('>>> code', code)

    try {
      // @ts-expect-error
      const extractable = executeModule(code, state.file.opts.filename, mapper)

      console.log('>>> extractable', extractable)
      console.log('>>> state.xxx', state.xxx)

      for (let i = 0; i < state.xxx.length; i++) {
        const f = state.xxx[i]

        if (f.kind === 'theme') {
          // const chunk = extractable[i]
          // const className = chunk.className
          // const css = compileCss(chunk.css)

          f.node.replaceWith(
            t.objectExpression([
              t.objectProperty(
                t.identifier('theme'),
                t.objectExpression([
                  t.objectProperty(t.identifier('css'), t.stringLiteral('css')),
                  t.objectProperty(t.identifier('className'), t.stringLiteral('className')),
                ]),
              ),
            ]),
          )
        }
      }

      // const extractable = executeModule(code, state.file.opts.filename, mapper)
      // for (let i = 0; i < extractable.length; i++) {
      //   const chunk = extractable[i]
      //   const className = generateClassName(chunk.css, state)
      //   const css = compileCss(chunk.css.replace(chunk.className, className))
      //   state.nodes[i].replaceWith(
      //     t.objectExpression([
      //       t.objectProperty(t.identifier('css'), t.stringLiteral(css)),
      //       t.objectProperty(t.identifier('className'), t.stringLiteral(className)),
      //     ]),
      //   )
      // }
    } catch (error) {
      console.log(error)
    }
  }

  return {
    name: '@steely/babel-plugin',
    visitor: {
      Program: {
        enter: (path, state) => {
          if (isVirtualModuleCaller || processed.has(state.filename)) {
            return
          }

          processed.add(state.filename)

          state.imports = []
          state.nodes = []
          state.extractable = []
          state.xxx = []

          path.traverse({
            ImportDeclaration: (p) => {
              // @ts-expect-error (TODO: Fix ts issue)
              collectImports(p, state)
            },

            TaggedTemplateExpression: (p) => {
              // @ts-expect-error (TODO: Fix ts issue)
              collectExtractable(p, state)
            },

            CallExpression: (p) => {
              // @ts-expect-error (TODO: Fix ts issue)
              collectExtractableTheme(p, state)
            },
          })

          // @ts-expect-error (TODO: Fix ts issue)
          if (state.extractable.length > 0) {
            // @ts-expect-error (TODO: Fix ts issue)
            extractStyles(path, state)
          }

          // @ts-expect-error
          if (state.xxx.length > 0) {
            // @ts-expect-error
            extract(path, state)
          }
          // for (const f of state.xxx) {
          // if (f.kind === 'theme') {
          // }
          // }
        },
        exit: () => {
          processed.clear()
        },
      },
    },
  }
})