import type esbuild from "esbuild";
import assert from "node:assert";
import crypto from "node:crypto";
import fs from "node:fs/promises";
import { createRequire } from "node:module";
import path from "pathe";
import type { WorkerBundle } from "../worker-bundle.ts";

export function createWasmPlugin() {
  const modules = new Map<string, WorkerBundle.Module>();
  const plugin: esbuild.Plugin = {
    name: "alchemy-wasm",
    setup(build) {
      build.onStart(() => {
        modules.clear();
      });

      // Handle imports like `import "./foo.wasm"` and `import "./foo.wasm?module"`
      // TODO(john): Figure out why this suddenly became necessary
      build.onResolve({ filter: /\.wasm(\?.*)?$/ }, async (args) => {
        const resolved = modules.get(args.path);
        if (resolved) {
          return { path: resolved.path, external: true };
        }

        // Resolve path to source file, excluding the `?module` suffix
        const normalizedPath = path.normalize(args.path).replace(/\?.*$/, "");

        // Determine the source file: use Node module resolution for bare specifiers (npm packages),
        // otherwise resolve relative to the importing file's directory
        let copyFrom: string;
        if (/^[./]/.test(normalizedPath)) {
          copyFrom = path.resolve(args.resolveDir, normalizedPath);
        } else {
          try {
            // createRequire needs a file path to anchor node_modules resolution from.
            // The filename is a dummy — only the directory matters.
            const esmRequire = createRequire(
              path.join(args.resolveDir, "noop.js"),
            );
            copyFrom = esmRequire.resolve(normalizedPath);
          } catch {
            // Fall back to the original behavior if require.resolve fails
            copyFrom = path.resolve(args.resolveDir, normalizedPath);
          }
        }

        // Resolve path to outdir (required for monorepos if the workdir is not the same as process.cwd())
        assert(
          build.initialOptions.absWorkingDir && build.initialOptions.outdir,
          "Missing absWorkingDir or outdir from esbuild options",
        );
        const outdir = path.resolve(
          build.initialOptions.absWorkingDir,
          build.initialOptions.outdir,
        );

        // Use path hash as module specifier for portability (note: the `?module` suffix is not needed in workerd)
        const hash = crypto.createHash("sha256").update(copyFrom).digest("hex");
        const specifier = `${hash}.wasm`;

        // Copy to outdir so it's included in the upload
        const copyTo = path.join(outdir, specifier);
        await fs.mkdir(path.dirname(copyTo), { recursive: true });
        await fs.copyFile(copyFrom, copyTo);
        modules.set(args.path, {
          type: "wasm",
          path: specifier,
        });

        return { path: specifier, external: true };
      });
    },
  };
  return { plugin, modules };
}