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 decoder = new TextDecoder();
                        const getString = (ptr) => {
                            const view = new Uint8Array(memory.buffer, ptr);
                            const start = ptr;
                            const end = ptr + view.indexOf(0);

                            return decoder.decode(memory.buffer.slice(start, end));
                        };
                        let buf = '';
                        const { instance } = await WebAssembly.instantiateStreaming(fetch(url), {
                            env: {
                                memory,
                                log(format, argsPtr) {
                                    format = getString(format);
                                    let isFormat = false;
                                    let w = 4;
                                    const align = (a) => argsPtr += (argsPtr % a);
                                    for (const c of format) {
                                        if (!isFormat) {
                                            if (c === '%') {
                                                isFormat = true;
                                            } else if (c === '\\n') {
                                                console.log('[wasm]', buf);
                                                buf = '';
                                            } else {
                                                buf += c;
                                            }
                                        } else switch(c) {
                                            case '%': buf += '%'; isFormat = false; break;
                                            case 's': align(4); buf += getString(data.getInt32(argsPtr, true)); argsPtr += 4; isFormat = false; break;
                                            case 'd': align(w); buf += data[w === 4 ? 'getInt32': 'getBigInt64'](argsPtr, true); argsPtr += w; isFormat = false; break;
                                            case 'x': align(w); buf += data[w === 4 ? 'getUint32': 'getBigUint64'](argsPtr, true).toString(16); argsPtr += w; isFormat = false; break;
                                            case 'o': align(w); buf += data[w === 4 ? 'getUint32': 'getBigUint64'](argsPtr, true).toString(8); argsPtr += w; isFormat = false; break;
                                            case 'u': align(w); buf += data[w === 4 ? 'getUint32': 'getBigUint64'](argsPtr, true); argsPtr += w; isFormat = false; break;
                                            case 'f': align(8); buf += data.getFloat64(argsPtr, true); argsPtr += 8; isFormat = false; break;
                                            case 'c': align(4); buf += String.fromCharCode(data.getInt32(argsPtr, true)); argsPtr += 4; isFormat = false; break;
                                            case 'l': w = 8; break;
                                            default: buf += '%' + c; isFormat = false; break;
                                        }
                                    }
                                },
                                grow(blocks) {
                                    if (blocks > 0) {
                                        memory.grow(blocks);
                                        data = new DataView(memory.buffer);
                                    }
                                },
                                Math_round(x) { return Math.round(x); },
                                Math_floor(x) { return Math.floor(x); },
                                Math_ceil(x) { return Math.ceil(x); },
                                Math_exp(x) { return Math.exp(x); },
                                Math_log(x) { return Math.log(x); },
                                Math_log10(x) { return Math.log10(x); },
                                Math_sqrt(x) { return Math.sqrt(x); },
                                Math_sin(x) { return Math.sin(x); },
                                Math_cos(x) { return Math.cos(x); },
                                Math_tan(x) { return Math.tan(x); },
                                Math_asin(x) { return Math.asin(x); },
                                Math_acos(x) { return Math.acos(x); },
                                Math_atan(x) { return Math.atan(x); },
                                Math_sinh(x) { return Math.sinh(x); },
                                Math_cosh(x) { return Math.cosh(x); },
                                Math_tanh(x) { return Math.tanh(x); },

                                Math_fmod(x, y) { return x % y; },
                                Math_pow(x, y) { return Math.pow(x, y); },
                                Math_atan2(x, y) { return Math.atan2(x, y); },
                            }
                        });

                        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 include = `${buildAssets}/include`;
                    const glob = new Bun.Glob(`${buildAssets}/lib/**/*.c`);
                    const stdlib = await Array.fromAsync(glob.scan());
                    const objPath = wasmPath + '.o';

                    const features = [
                        'bulk-memory',
                        'extended-const',
                        'relaxed-simd',
                        'simd128',
                        'tail-call',
                        'sign-ext',
                        'nontrapping-fptoint',
                        'reference-types',
                        'multivalue',

                    ].map(f => `-m${f}`);

                    const flags = [
                        '--target=wasm32',
                        production ? '-O3' : '-O0',
                        '-flto',
                        '-fno-builtin',
                        '--no-standard-libraries',
                        '-Wall',
                        '-Wextra',
                        '-Wpedantic',
                        '-Werror',
                        '-Wshadow',
                        ...features, 
                    ];
                    const std = args.path.endsWith('.cpp') ? '-std=gnu++23': '-std=gnu23';
                    const compileResult = await $`clang -c ${flags} ${std} -I ${include} -o ${objPath} ${args.path}`;

                    if (compileResult.exitCode !== 0) {
                        throw new Error('Compile failed, check output');
                    }

                    const linkFlags = [
                        '--lto-O3',
                        '--no-entry',
                        '--import-memory',
                    ].map(f => `-Wl,${f}`);

                    const linkResult = await $`clang ${flags} -std=gnu23 -I ${include} ${linkFlags} -o ${wasmPath} ${objPath} ${stdlib}`;

                    if (linkResult.exitCode !== 0) {
                        throw new Error('Link 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;