import templateBuilder from '@babel/template';

function isAtom(t, callee) {
  if (t.isIdentifier(callee) && atomFunctionNames.includes(callee.name)) {
    return true;
  }
  if (t.isMemberExpression(callee)) {
    const { property } = callee;
    if (t.isIdentifier(property) && atomFunctionNames.includes(property.name)) {
      return true;
    }
  }
  return false;
}
const atomFunctionNames = [
  "atom",
  "atomFamily",
  "atomWithDefault",
  "atomWithObservable",
  "atomWithReducer",
  "atomWithReset",
  "atomWithStorage",
  "freezeAtom",
  "loadable",
  "selectAtom",
  "splitAtom"
];

function reactRefreshPlugin({
  types: t
}) {
  return {
    pre({ opts }) {
      if (!opts.filename) {
        throw new Error("Filename must be available");
      }
    },
    visitor: {
      Program: {
        exit(path) {
          const jotaiAtomCache = templateBuilder(`
          globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || {
            cache: new Map(),
            get(name, inst) { 
              if (this.cache.has(name)) {
                return this.cache.get(name)
              }
              this.cache.set(name, inst)
              return inst
            },
          }`)();
          path.unshiftContainer("body", jotaiAtomCache);
        }
      },
      ExportDefaultDeclaration(nodePath, state) {
        const { node } = nodePath;
        if (t.isCallExpression(node.declaration) && isAtom(t, node.declaration.callee)) {
          const filename = state.filename || "unknown";
          const atomKey = `${filename}/defaultExport`;
          const buildExport = templateBuilder(`export default globalThis.jotaiAtomCache.get(%%atomKey%%, %%atom%%)`);
          const ast = buildExport({
            atomKey: t.stringLiteral(atomKey),
            atom: node.declaration
          });
          nodePath.replaceWith(ast);
        }
      },
      VariableDeclarator(nodePath, state) {
        var _a, _b;
        if (t.isIdentifier(nodePath.node.id) && t.isCallExpression(nodePath.node.init) && isAtom(t, nodePath.node.init.callee) && (((_a = nodePath.parentPath.parentPath) == null ? void 0 : _a.isProgram()) || ((_b = nodePath.parentPath.parentPath) == null ? void 0 : _b.isExportNamedDeclaration()))) {
          const filename = state.filename || "unknown";
          const atomKey = `${filename}/${nodePath.node.id.name}`;
          const buildAtomDeclaration = templateBuilder(`const %%atomIdentifier%% = globalThis.jotaiAtomCache.get(%%atomKey%%, %%atom%%)`);
          const ast = buildAtomDeclaration({
            atomIdentifier: t.identifier(nodePath.node.id.name),
            atomKey: t.stringLiteral(atomKey),
            atom: nodePath.node.init
          });
          nodePath.parentPath.replaceWith(ast);
        }
      }
    }
  };
}

export { reactRefreshPlugin as default };