import { plugin, $, type BunPlugin } from "bun";
import path from 'path';
import fs from 'fs/promises';
import { renderDts } from './wasmTypes';
import { extractExports } from './wasmExports';

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

interface CompilerWithFlags {
    cc: string;
    wasmOpt: string | null;
    reactorCrt: string | null;
    flags: string[];
}

const STD_C = ['-std=gnu23'];
const STD_CPP = ['-std=gnu++26', '-fno-rtti'];

const WASI_VERSION = '32.0';
const WASI_MAJOR = WASI_VERSION.split('.')[0];
const wasiArchiveURL = `https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-${WASI_MAJOR}/wasi-sdk-${WASI_VERSION}-x86_64-linux.tar.gz`;
const rtURL = `https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-${WASI_MAJOR}/libclang_rt-${WASI_VERSION}.tar.gz`;

let compilerPromise: Promise<CompilerWithFlags> | null = null;

const getCompiler = (): Promise<CompilerWithFlags> => {
    if (compilerPromise) return compilerPromise;
    compilerPromise = (async (): Promise<CompilerWithFlags> => {
        const wasiDir = path.resolve(import.meta.dir, '..', '..', 'dist', 'wasi');
        const cc: CompilerWithFlags = {
            cc: 'clang',
            wasmOpt: Bun.which('wasm-opt') ?? null,
            reactorCrt: null,
            flags: [
                '--target=wasm32',
                '--no-standard-libraries',
                '-fno-builtin',
            ],
        };

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

        const installedVersion = (await Bun.file(path.resolve(wasiDir, 'VERSION')).text().catch(() => '')).slice(0, WASI_VERSION.length);

        if (installedVersion !== WASI_VERSION) {
            console.log(`WASI version mismatch. Downloading WASI SDK ${wasiArchiveURL}...`)
            const response = await fetch(wasiArchiveURL);

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

            const bytes = await response.bytes();

            console.log(`Extracting WASI SDK...`);
            await $`tar -xzv -C ${wasiDir} --strip-components=1 < ${bytes}`;

            console.log(`Downloading libclang_rt.builtins-wasm32-wasi-${WASI_VERSION}...`);
            const rtResponse = await fetch(rtURL);

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

            const rtBytes = await rtResponse.bytes();

            console.log(`Extracting libclang_rt.builtins-wasm32-wasi-${WASI_VERSION}...`);
            await $`tar -xzv -C ${wasiDir} --strip-components=1 < ${rtBytes}`;
        }

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

        const wasmOptInSdk = path.resolve(wasiDir, 'bin', 'wasm-opt');
        cc.wasmOpt = await fs.access(wasmOptInSdk).then(() => wasmOptInSdk).catch(() => Bun.which('wasm-opt') ?? null);

        const reactorCrt = path.resolve(wasiDir, 'share', 'wasi-sysroot', 'lib', 'wasm32-wasip1', 'crt1-reactor.o');
        cc.reactorCrt = await fs.access(reactorCrt).then(() => reactorCrt).catch(() => null);

        return cc;
    })();
    return compilerPromise;
};

async function mtimeMs(p: string): Promise<number> {
    return fs.stat(p).then(s => s.mtimeMs).catch(() => 0);
}

async function needsRebuild(output: string, inputs: string[], flags?: string[]): Promise<boolean> {
    const [outTime, storedFlags] = await Promise.all([
        mtimeMs(output),
        flags ? Bun.file(output + '.flags').text().catch(() => null) : Promise.resolve(null),
    ]);
    if (outTime === 0) return true;
    if (flags && storedFlags !== flags.join('\0')) return true;
    const inTimes = await Promise.all(inputs.map(mtimeMs));
    return inTimes.some(t => t > outTime);
}

async function saveFlags(output: string, flags: string[]): Promise<void> {
    await Bun.write(output + '.flags', flags.join('\0'));
}

async function compileLibGroup(
    cc: CompilerWithFlags,
    flags: string[],
    sources: string[],
    std: string[],
    objDir: string,
): Promise<string[]> {
    if (sources.length === 0) return [];

    await fs.mkdir(objDir, { recursive: true });
    const objPaths = sources.map(src => path.resolve(objDir, path.basename(src) + '.o'));

    const compilationFlags = [...flags, ...std];
    await Promise.all(sources.map(async (src, i) => {
        if (await needsRebuild(objPaths[i], [src], compilationFlags)) {
            const result = await $`${cc.cc} -c ${flags} ${std} -o ${objPaths[i]} ${src}`;
            if (result.exitCode !== 0) throw new Error(`Compile failed: ${src}`);
            await saveFlags(objPaths[i], compilationFlags);
        }
    }));

    return objPaths;
}

type FnSig = { params: ('i32' | 'i64' | 'f32' | 'f64')[]; results: ('i32' | 'i64' | 'f32' | 'f64')[] };
type FnPtrMeta = Record<string, [number, FnSig][]>;

async function instantiate(url: string, fnPtrs: FnPtrMeta = {}) {
    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 decoder = new TextDecoder();
    let buf = '';
    let errBuf = '';
    const { instance } = await WebAssembly.instantiateStreaming(fetch(url), {
        env: { memory },
        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}`));
                    }
                }
            },
            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);
            },
        }, {
            get(target, p) {
                if (p in target) {
                    return (...args: any[]) => {
                        const fn = target[p as keyof typeof target] as Function;
                        const result = fn(...args) ?? 0;
                        console.debug(`${String(p)}(${args.join(', ')}) = ${result}`);

                        return result;
                    }
                }
                return (...args: any[]) => {
                    console.warn(`${String(p)}(${args.join(', ')}) is undefined`);
                    return 0;
                }
            },
        }),
    });

    (instance.exports._initialize as (() => void) | undefined)?.();

    const table = instance.exports.__indirect_function_table as WebAssembly.Table | undefined;
    const exported: Record<string, unknown> = Object.fromEntries(
        Object.entries(instance.exports).filter(([k]) => !k.startsWith('_'))
    );

    // Wrap exports that take function-pointer args so callers can pass JS functions
    // directly. Same JS function reuses the same table slot across calls.
    const fnPtrCache = new WeakMap<Function, number>();
    for (const [name, ptrs] of Object.entries(fnPtrs)) {
        const orig = exported[name];
        if (typeof orig !== 'function' || !table) continue;
        exported[name] = (...args: unknown[]) => {
            for (const [i, sig] of ptrs) {
                const v = args[i];
                if (typeof v !== 'function') continue;
                let idx = fnPtrCache.get(v);
                if (idx === undefined) {
                    const wasmFn = new WebAssembly.Function({ parameters: sig.params, results: sig.results }, v as (...a: number[]) => void);
                    idx = table.grow(1, wasmFn);
                    fnPtrCache.set(v, idx);
                }
                args[i] = idx;
            }
            return (orig as (...a: unknown[]) => unknown)(...args);
        };
    }

    return {
        ...exported,
        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) => {
                const distDir = path.resolve(import.meta.dir, '..', '..', 'dist');
                const inputBasename = path.basename(args.path);
                let wasmPath = path.resolve(distDir, inputBasename.replace(/\.(c|cpp)$/, '.wasm'));
                const fnPtrMeta: Record<string, [number, FnSig][]> = {};
                if (args.path.endsWith('.wasm')) {
                    wasmPath = args.path;
                } else {
                    await fs.mkdir(distDir, { recursive: true });

                    const buildAssets = path.resolve(import.meta.dir, '..', 'assets');
                    const include = `${buildAssets}/include`;
                    const cc = await getCompiler();

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

                    const flags = [
                        ...cc.flags,
                        production ? '-Os' : '-O0',
                        '-flto',
                        '-fno-exceptions',
                        '-fno-stack-protector',
                        '-ffunction-sections',
                        '-fdata-sections',
                        '-Wall',
                        '-Wextra',
                        '-Wpedantic',
                        '-Werror',
                        '-Wshadow',
                        '-Wconversion',
                        '-fvisibility=hidden',
                        ...features,
                        '-I', include,
                    ];

                    const std = args.path.endsWith('.cpp') ? STD_CPP : STD_C;

                    const inputObjPath = path.resolve(distDir, inputBasename + '.o');
                    const libCObjDir = path.resolve(distDir, 'lib.c');
                    const libCppObjDir = path.resolve(distDir, 'lib.cpp');

                    const libCGlob = new Bun.Glob(`${buildAssets}/lib/**/*.c`);
                    const libCppGlob = new Bun.Glob(`${buildAssets}/lib/**/*.cpp`);
                    const [libCFiles, libCppFiles] = await Promise.all([
                        Array.fromAsync(libCGlob.scan()),
                        Array.fromAsync(libCppGlob.scan()),
                    ]);

                    const [, libCObjs, libCppObjs] = await Promise.all([
                        (async () => {
                            const compilationFlags = [...flags, ...std];
                            if (!await needsRebuild(inputObjPath, [args.path], compilationFlags)) return;
                            const result = await $`${cc.cc} -c ${flags} ${std} -o ${inputObjPath} ${args.path}`;
                            if (result.exitCode !== 0) throw new Error('Compile failed, check output');
                            await saveFlags(inputObjPath, compilationFlags);
                        })(),
                        compileLibGroup(cc, flags, libCFiles, STD_C, libCObjDir),
                        compileLibGroup(cc, flags, libCppFiles, STD_CPP, libCppObjDir),
                    ]);

                    const allObjs = [inputObjPath, ...libCObjs, ...libCppObjs];

                    const exportInfos = extractExports(await Bun.file(args.path).text());
                    for (const e of exportInfos) {
                        if (e.fnPtrSigs.size > 0) {
                            fnPtrMeta[e.name] = [...e.fnPtrSigs.entries()];
                        }
                    }

                    if (await needsRebuild(wasmPath, allObjs)) {
                        const meta = new Map(exportInfos.map(e => [e.name, {
                            paramNames: e.paramNames,
                            fnPtrSigs: e.fnPtrSigs,
                        }]));

                        const linkFlags = [
                            '--strip-all',
                            '--lto-O3',
                            '--gc-sections',
                            '--no-entry',
                            '--import-memory',
                            '--export-table',
                            '--growable-table',
                            ...exportInfos.flatMap(e => ['--export-if-defined', e.name]),
                        ].map(f => `-Wl,${f}`);

                        const reactorInputs = cc.reactorCrt ? [cc.reactorCrt] : [];
                        const linkResult = await $`${cc.cc} ${flags} ${linkFlags} -o ${wasmPath} -lstdc++ -nostartfiles ${reactorInputs} ${allObjs}`;

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

                        if (production && cc.wasmOpt) {
                            const optResult = await $`${cc.wasmOpt} -Os --converge -o ${wasmPath} ${wasmPath}`;
                            if (optResult.exitCode !== 0) throw new Error('wasm-opt failed, check output');
                        }

                        const wasmBytes = new Uint8Array(await Bun.file(wasmPath).arrayBuffer());
                        await Bun.write(`${args.path}.d.ts`, renderDts(wasmBytes, meta));
                    }
                }

                const wasmContent = await Bun.file(wasmPath).arrayBuffer();
                const wasmBuffer = Buffer.from(wasmContent).toString('base64');
                const wasmURL = `data:application/wasm;base64,${wasmBuffer}`;

                const jsContent = `
                    ${instantiate}
                    const module = await instantiate(new URL(${JSON.stringify(wasmURL)}), ${JSON.stringify(fnPtrMeta)});

                    export default module;
                `;

                return {
                    loader: 'js',
                    contents: jsContent,
                };
            });
        }
    };
    plugin(p);
    return p;
};

export default wasmPlugin;