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

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

interface CompilerWithFlags {
    cc: string;
    flags: string[];
}

const wasiArchiveURL = 'https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/wasi-sdk-25.0-x86_64-linux.tar.gz';
const rtURL = 'https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/libclang_rt.builtins-wasm32-wasi-25.0.tar.gz';

const getCompiler = async (): Promise<CompilerWithFlags> => {
    const wasiDir = path.resolve(import.meta.dir, '..', '..', 'dist', 'wasi');
    const cc: CompilerWithFlags = {
        cc: 'clang',
        flags: [
            '--target=wasm32',
            '--no-standard-libraries',
            '-fno-builtin',
        ],
    };

    await fs.mkdir(wasiDir, { recursive: true });

    if (!await Bun.file(path.resolve(wasiDir, 'VERSION')).exists()) {
        const response = await fetch(wasiArchiveURL);

        if (!response.ok) {
            return cc;
        }

        const bytes = await response.bytes();

        await $`tar -xzv -C ${wasiDir} --strip-components=1 < ${bytes}`;

        const rtResponse = await fetch(rtURL);

        if (!rtResponse.ok) {
            return cc;
        }

        const rtBytes = await rtResponse.bytes();

        await $`tar -xzv -C ${wasiDir} --strip-components=1 < ${rtBytes}`;
    }

    cc.cc = `${path.resolve(wasiDir, 'bin', 'clang')}`;
    cc.flags = [
        '--target=wasm32-wasi',
        `--sysroot=${path.resolve(wasiDir, 'share', 'wasi-sysroot')}`,
    ];

    return cc;
}

async function instantiate(url: string) {
    if (typeof WebAssembly.Function !== "function") {
        const typeCodes = {
            i32: 0x7f,
            i64: 0x7e,
            f32: 0x7d,
            f64: 0x7c
        };

        type tc = keyof typeof typeCodes;

        // @ts-ignore
        WebAssembly.Function = function ({
            parameters = [],
            results = [],
        }: { parameters: tc[], results: tc[] }, func: Function) {
            const
                resultAmount = results.length,
                parameterAmount = parameters.length;

            const bytes = new Uint8Array([
                // Headers (Magic and Version)
                0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,

                // Type section
                0x01, (4 + parameterAmount + resultAmount), 0x01, 0x60, // Section ID, Type section length, Count, Function
                parameterAmount, ...parameters.map(p => typeCodes[p]), // Parameter types
                resultAmount, ...results.map(r => typeCodes[r]), // Result types

                // Import and Exports
                0x02, 0x07, 0x01, 0x01, 0x65, 0x01, 0x66, 0x00, 0x00,
                0x07, 0x05, 0x01, 0x01, 0x66, 0x00, 0x00
            ]);

            // Compile and instantiate
            const module = new WebAssembly.Module(bytes);
            const instance = new WebAssembly.Instance(module, {
                e: {
                    f: func
                }
            });
            return instance.exports.f;
        };
    }
    const memory = new WebAssembly.Memory({
        initial: 32,
    });
    const table = new WebAssembly.Table({ initial: 16, element: 'anyfunc' });
    const decoder = new TextDecoder();
    let buf = '';
    let errBuf = '';
    const { instance } = await WebAssembly.instantiateStreaming(fetch(url), {
        env: { memory, __indirect_function_table: table },
        wasi_snapshot_preview1: new Proxy({
            random_get: (ptr: number, length: number) => {
                const data = new DataView(memory.buffer);
                for (let i = 0; i < length; i++) {
                    data.setUint8(ptr + i, Math.random() * 256);
                }
            },
            fd_write: (fd: number, iovsPtr: number, iovsLength: number, bytesWrittenPtr: number) => {
                const iovs = new Uint32Array(memory.buffer, iovsPtr, iovsLength * 2);
                const data = new DataView(memory.buffer);
                if (fd === 1 || fd === 2) {
                    let text = "";
                    let totalBytesWritten = 0;
                    for (let i = 0; i < iovsLength * 2; i += 2) {
                        const offset = iovs[i];
                        const length = iovs[i + 1];
                        const textChunk = decoder.decode(new Int8Array(memory.buffer, offset, length));
                        text += textChunk;
                        totalBytesWritten += length;
                    }
                    data.setInt32(bytesWrittenPtr, totalBytesWritten, true);
                    if (fd === 1) {
                        const lines = (buf + text).split('\n');
                        buf = lines.pop() ?? '';
                        lines.forEach(l => console.log(`[wasm] ${l}`));
                    } else {
                        const lines = (errBuf + text).split('\n');
                        errBuf = lines.pop() ?? '';
                        lines.forEach(l => console.error(`[wasm] ${l}`));
                    }
                }
                return 0;
            },
            clock_time_get: (_id: number, _precision: bigint, timePtr: number) => {
                const now = Date.now();
                const nanosecondNow = BigInt(now) * 1_000_000n;
                const data = new DataView(memory.buffer);
                data.setBigInt64(timePtr, nanosecondNow, true);
                return 0;
            }
        }, {
            get(target, p) {
                if (p in target) {
                    console.debug(`[${String(p)}] exists`);
                    return target[p as keyof typeof target];
                }
                return (...args: any[]) => { console.warn(`[${String(p)}]`, args); return 0; }
            },
        }),
    });

    return {
        ...instance.exports,
        memory,
        table,
        get data() { return new DataView(memory.buffer); },
    };
}

const wasmPlugin = ({ production }: 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 = `
                    ${instantiate}
                    const module = await instantiate(new URL($WASM$));

                    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 cc = await getCompiler();

                    const features = [
                        'bulk-memory',
                        'extended-const',
                        'relaxed-simd',
                        'simd128',
                        'tail-call',
                        'sign-ext',
                        'nontrapping-fptoint',
                        'reference-types',
                        'multivalue',
                    ].map(f => `-m${f}`);

                    const flags = [
                        ...cc.flags,
                        production ? '-O3' : '-O0',
                        '-flto',
                        '-fno-exceptions',
                        '-Wall',
                        '-Wextra',
                        '-Wpedantic',
                        '-Werror',
                        '-Wshadow',
                        '-Wconversion',
                        '-fvisibility=hidden',
                        ...features,
                        '-I', include,
                    ];
                    const std = args.path.endsWith('.cpp') ? '-std=gnu++26' : '-std=gnu23';
                    const compileResult = await $`${cc.cc} -c ${flags} ${std} -o ${objPath} ${args.path}`;

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

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

                    const linkResult = await $`${cc.cc} ${flags} -std=gnu23 ${linkFlags} -lstdc++ -nostartfiles -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;