import { plugin, $, type BunPlugin } from "bun";
import path from 'path';

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)$/ }, async (args) => {
                let wasmPath = path.resolve(import.meta.dir, '..', 'dist', 'tmp.wasm');
                let jsContent: string = `
                    async function instantiate(url) {
                        const memory = new WebAssembly.Memory({
                            initial: 32,
                        });
                        let data = new DataView(memory.buffer);                        
                        const { instance } = await WebAssembly.instantiateStreaming(fetch(url), {
                            env: {
                                memory,
                                log(...args) {
                                    console.log('[wasm]', ...args);
                                },
                                grow(blocks) {
                                    if (blocks > 0) {
                                        memory.grow(blocks);
                                        data = new DataView(memory.buffer);
                                    }
                                }
                            }
                        });

                        return {
                            ...instance.exports,
                            memory,
                            get data() { return data; },
                        };
                    }
                    const module = await instantiate(new URL($WASM$));
                    if (typeof module.__srand === 'function') module.__srand(BigInt(Date.now()));

                    export default module;
                `;
                if (args.path.endsWith('.wasm')) {
                    wasmPath = args.path;
                } else {
                    const buildAssets = path.resolve(import.meta.dir, 'assets');
                    const stdlib = `${buildAssets}/stdlib.c`;
                    const opt = production ? '-O3': '-O0';
                    const result = await $`clang --target=wasm32 ${opt} -flto -fno-builtin --no-standard-libraries -I ${buildAssets} -Wall -Wextra -Wpedantic -Werror -Wl,--lto-O3 -Wl,--no-entry -Wl,--import-memory -o ${wasmPath} ${args.path} ${stdlib}`;

                    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;