#!/usr/bin/env tsx
import assert from 'node:assert';
import { readFileSync } from 'node:fs';
import { extname, resolve } from 'node:path';

import * as babel from '@babel/core';
import type { types as t, types } from '@babel/core';
import babelMacros, { type MacroHandler } from 'babel-plugin-macros';
import type { Loader, Plugin } from 'esbuild';
import PropTypes from 'prop-types';

import * as polyfill from '../src/polyfill';

const polyfills = Object.keys(polyfill);

class HandlerMap {
  map = new Map<string, MacroHandler>();

  set(names: string | string[], handler: MacroHandler) {
    names = Array.isArray(names) ? names : [names];
    const macro = babelMacros.createMacro(handler);
    for (const name of names) {
      this.map.set(name, macro);
    }
    return this;
  }

  get keys() {
    return [...this.map.keys()];
  }

  resolvePath = (module: string) => module;
  require = (module: string) => this.map.get(module);
  isMacrosName = (module: string) => this.map.has(module);
}

const map = new HandlerMap()
  .set(
    'object.assign',
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('assign'))),
  )
  .set(
    ['object-values', 'object.values'],
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('values'))),
  )
  .set(
    'object.fromentries',
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('fromEntries'))),
  )
  .set(
    'object.entries',
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('entries'))),
  )
  .set(
    'hasown',
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('hasOwn'))),
  )
  .set(
    'has',
    replace(t => t.memberExpression(t.identifier('Object'), t.identifier('hasOwn'))),
  )
  .set(
    'array-includes',
    proto(t => t.identifier('includes')),
  )
  .set(
    'array.prototype.flatmap',
    proto(t => t.identifier('flatMap')),
  )
  .set(
    'array.prototype.flat',
    proto(t => t.identifier('flat')),
  )
  .set(
    'array.prototype.findlastindex',
    proto(t => t.identifier('findLastIndex')),
  )
  .set(
    'array.prototype.tosorted',
    proto(t => t.identifier('toSorted')),
  )
  .set(
    'array.prototype.toreversed',
    proto(t => t.identifier('toReversed')),
  )
  .set(
    'array.prototype.findlast',
    proto(t => t.identifier('findLast')),
  )
  .set(
    'string.prototype.matchall',
    proto(t => t.identifier('matchAll')),
  )
  .set(
    'string.prototype.includes',
    proto(t => t.identifier('includes')),
  )
  .set(
    'object.groupby',
    replace(t =>
      t.memberExpression(
        t.callExpression(t.identifier('require'), [t.stringLiteral('lodash-es')]),
        t.identifier('groupBy'),
      ),
    ),
  );

// es-iterator-helpers/Iterator.prototype.*
const polyfillPath = resolve(import.meta.dirname, '../src/polyfill.ts');
const requirePolyfill = (t: typeof types, name: string) =>
  t.memberExpression(
    t.callExpression(t.identifier('require'), [t.stringLiteral(polyfillPath)]),
    t.identifier(name),
  );
map.set(
  `es-iterator-helpers/Iterator.from`,
  replace(t => requirePolyfill(t, 'from')),
);
for (const name of polyfills) {
  map.set(
    `es-iterator-helpers/Iterator.prototype.${name}`,
    replace(t => requirePolyfill(t, name)),
  );
}

map.set(
  'safe-regex-test',
  replace(t => requirePolyfill(t, 'safeRegexTest')),
);

function replace(getReplacement: (types: typeof t) => t.Expression): MacroHandler {
  return ({ references, babel: { types: t } }) => {
    for (const referencePath of references.default) {
      referencePath.replaceWith(getReplacement(t));
    }
  };
}

function proto(getProperty: (types: typeof t) => t.Expression): MacroHandler {
  return ({ references, babel: { types: t } }) => {
    for (const referencePath of references.default) {
      const { parent, parentPath } = referencePath;
      assert(t.isCallExpression(parent));
      const [callee, ...rest] = parent.arguments;
      parentPath!.replaceWith(
        t.callExpression(
          t.memberExpression(callee as t.Expression, getProperty(t)),
          rest,
        ),
      );
    }
  };
}

export const babelPlugin: Plugin = {
  name: 'babel',
  setup(build) {
    const { keys, ...macroOptions } = map;

    build.onLoad({ filter: /\.[jt]sx?$/ }, args => {
      const { path } = args;
      if (path.includes('node_modules/')) {
        return null;
      }

      const source = readFileSync(path, 'utf8')
        .replaceAll("require('object.hasown/polyfill')()", 'Object.hasOwn')
        .replaceAll("require('object.fromentries/polyfill')()", 'Object.fromEntries')
        .replaceAll(
          "Object.keys(require('prop-types'))",
          JSON.stringify(Object.keys(PropTypes)),
        );

      const isFlow = source.includes('@flow');
      const loader = extname(path).slice(1) as Loader;

      if (!isFlow && !keys.some(key => source.includes(key))) {
        return { contents: source, loader };
      }

      const res = babel.transformSync(source, {
        filename: path,
        babelrc: false,
        configFile: false,
        parserOpts: {
          plugins: [isFlow ? 'flow' : 'typescript'],
        },
        plugins: [
          isFlow && '@babel/plugin-transform-flow-strip-types',
          ['babel-plugin-macros', macroOptions],
        ].filter(Boolean),
      })!;

      return {
        contents: res.code!,
        loader,
      };
    });
  },
};