import { plugin, type BunPlugin } from "bun";
import path from 'path';
import asc from 'assemblyscript/asc';

interface WasmLoaderConfig {
    production?: boolean;
    portable?: boolean;
}

const wasmPlugin = ({ production, portable }: WasmLoaderConfig = {}): BunPlugin => {
    const p: BunPlugin = {
        name: "WASM loader",
        async setup(build) {
            build.onLoad({ filter: /\.wasm\.ts$/ }, async (args) => {
                if (portable) {
                    const contents = await Bun.file(args.path).text();
                    return {
                        contents: `import "assemblyscript/std/portable/index.js";\n${contents}`,
                        loader: 'tsx',
                    }
                }
                const wasmPath = path.resolve(import.meta.dir, '..', '..', 'dist', 'tmp.wasm');
                const jsPath = wasmPath.replace(/\.wasm$/, '.js');
                const ascArgs = [
                    args.path,
                    '--outFile', wasmPath,
                    '--bindings', 'esm',
                ];
                ascArgs.push(production ? '--optimize' : '--debug');

                const { error, stderr } = await asc.main(ascArgs);

                if (error) {
                    console.error(stderr.toString(), error.message);
                    throw error;
                } else {
                    const jsContent = await Bun.file(jsPath).text();
                    const wasmContent = await Bun.file(wasmPath).arrayBuffer();
                    const wasmBuffer = Buffer.from(wasmContent).toString('base64');
                    const wasmURL = `data:application/wasm;base64,${wasmBuffer}`;

                    return {
                        loader: 'js',
                        contents: jsContent
                            .replace(/new URL\([^)]*\)/, `new URL(${JSON.stringify(wasmURL)})`),
                    };
                }
            });
        }
    };
    plugin(p);
    return p;
};

export default wasmPlugin;