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 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('.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 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;