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: /\.(c(pp)?|wasm(\.ts)?)$/ }, async (args) => {
                let wasmPath = path.resolve(import.meta.dir, '..', '..', 'dist', 'tmp.wasm');
                let jsContent: string = `
                    async function instantiate(url) {
                        const { instance } = await WebAssembly.instantiateStreaming(fetch(url));
                        return {
                            ...instance.exports,
                            data: new DataView(instance.exports.memory.buffer),
                        };
                    }
                    const module = await instantiate(new URL("tmp.wasm", import.meta.url));
                    export default module;
                `;
                if (args.path.endsWith('.ts')) {
                    if (portable) {
                        const contents = await Bun.file(args.path).text();
                        return {
                            contents: `import "assemblyscript/std/portable/index.js";\n${contents}`,
                            loader: 'tsx',
                        }
                    }
                    const jsPath = wasmPath.replace(/\.wasm$/, '.js');
                    const ascArgs = [
                        args.path,
                        '--outFile', wasmPath,
                        '--bindings', 'esm',
                        '-Ospeed'
                    ];

                    const { error, stderr } = await asc.main(ascArgs);
                    if (error) {
                        console.error(stderr.toString(), error.message);
                        throw error;
                    }
                    jsContent = await Bun.file(jsPath).text();
                } else if (args.path.endsWith('.wasm')) {
                    wasmPath = args.path;
                } else {
                    const result = await $`clang --target=wasm32 -O3 --no-standard-libraries -Wl,--export-all -Wl,--no-entry -fno-builtin -o ${wasmPath} ${args.path}`;
                    if (result.exitCode !== 0) {
                        throw new Error('Compile failed, check output');
                    }
                }

                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;