// Parses a wasm binary's type / function / import / export sections to recover
// each function export's signature, then renders TypeScript declarations.

import type { FnSig, WasmType } from './wasmExports';

type Sig = { params: number[]; results: number[] };

export interface ExportMeta {
    paramNames: string[];
    fnPtrSigs: Map<number, FnSig>;
}

const VALUE_TYPE_TS: Record<number, string> = {
    0x7f: 'number',  // i32
    0x7e: 'bigint',  // i64
    0x7d: 'number',  // f32
    0x7c: 'number',  // f64
};

const WASM_STR_TO_TS: Record<WasmType, string> = {
    i32: 'number',
    i64: 'bigint',
    f32: 'number',
    f64: 'number',
};

function readULEB128(bytes: Uint8Array, offset: number): [number, number] {
    let result = 0;
    let shift = 0;
    let i = 0;
    while (true) {
        const b = bytes[offset + i];
        result |= (b & 0x7f) << shift;
        i++;
        if ((b & 0x80) === 0) break;
        shift += 7;
    }
    return [result, i];
}

function skipLimits(bytes: Uint8Array, offset: number): number {
    const flags = bytes[offset];
    let o = offset + 1;
    const [, l1] = readULEB128(bytes, o); o += l1;
    if (flags & 1) {
        const [, l2] = readULEB128(bytes, o); o += l2;
    }
    return o;
}

export function parseWasmExports(bytes: Uint8Array): Map<string, Sig> {
    if (bytes[0] !== 0x00 || bytes[1] !== 0x61 || bytes[2] !== 0x73 || bytes[3] !== 0x6d) {
        throw new Error('Invalid wasm magic');
    }
    let offset = 8;

    const types: Sig[] = [];
    const funcTypes: number[] = [];      // defined function index → type index
    const importedFuncTypes: number[] = []; // imported function index → type index
    const exportEntries: { name: string; kind: number; index: number }[] = [];
    const decoder = new TextDecoder();

    while (offset < bytes.length) {
        const sectionId = bytes[offset++];
        const [sectionSize, s1] = readULEB128(bytes, offset);
        offset += s1;
        const sectionEnd = offset + sectionSize;

        if (sectionId === 1) {
            const [count, c1] = readULEB128(bytes, offset); offset += c1;
            for (let i = 0; i < count; i++) {
                if (bytes[offset++] !== 0x60) throw new Error('Expected functype 0x60');
                const [pc, p1] = readULEB128(bytes, offset); offset += p1;
                const params: number[] = [];
                for (let j = 0; j < pc; j++) params.push(bytes[offset++]);
                const [rc, r1] = readULEB128(bytes, offset); offset += r1;
                const results: number[] = [];
                for (let j = 0; j < rc; j++) results.push(bytes[offset++]);
                types.push({ params, results });
            }
        } else if (sectionId === 2) {
            const [count, c1] = readULEB128(bytes, offset); offset += c1;
            for (let i = 0; i < count; i++) {
                const [modLen, m1] = readULEB128(bytes, offset); offset += m1 + modLen;
                const [fldLen, f1] = readULEB128(bytes, offset); offset += f1 + fldLen;
                const kind = bytes[offset++];
                if (kind === 0x00) {
                    const [typeIdx, t1] = readULEB128(bytes, offset); offset += t1;
                    importedFuncTypes.push(typeIdx);
                } else if (kind === 0x01) {
                    offset++; // elem type
                    offset = skipLimits(bytes, offset);
                } else if (kind === 0x02) {
                    offset = skipLimits(bytes, offset);
                } else if (kind === 0x03) {
                    offset += 2; // valtype + mutability
                }
            }
        } else if (sectionId === 3) {
            const [count, c1] = readULEB128(bytes, offset); offset += c1;
            for (let i = 0; i < count; i++) {
                const [typeIdx, t1] = readULEB128(bytes, offset); offset += t1;
                funcTypes.push(typeIdx);
            }
        } else if (sectionId === 7) {
            const [count, c1] = readULEB128(bytes, offset); offset += c1;
            for (let i = 0; i < count; i++) {
                const [nameLen, n1] = readULEB128(bytes, offset); offset += n1;
                const name = decoder.decode(bytes.subarray(offset, offset + nameLen));
                offset += nameLen;
                const kind = bytes[offset++];
                const [index, i1] = readULEB128(bytes, offset); offset += i1;
                exportEntries.push({ name, kind, index });
            }
        }

        offset = sectionEnd;
    }

    const result = new Map<string, Sig>();
    for (const { name, kind, index } of exportEntries) {
        if (kind !== 0x00) continue;
        const typeIdx = index < importedFuncTypes.length
            ? importedFuncTypes[index]
            : funcTypes[index - importedFuncTypes.length];
        result.set(name, types[typeIdx]);
    }
    return result;
}

function tsType(valType: number): string {
    return VALUE_TYPE_TS[valType] ?? 'unknown';
}

function renderFnPtrType(sig: FnSig): string {
    const params = sig.params.map((t, i) => `a${i}: ${WASM_STR_TO_TS[t]}`).join(', ');
    const ret = sig.results.length === 0
        ? 'void'
        : sig.results.length === 1
            ? WASM_STR_TO_TS[sig.results[0]]
            : `[${sig.results.map(t => WASM_STR_TO_TS[t]).join(', ')}]`;
    return `(${params}) => ${ret}`;
}

function tsSignature({ params, results }: Sig, meta?: ExportMeta): string {
    const paramList = params
        .map((t, i) => {
            const name = meta?.paramNames[i] || `a${i}`;
            const fnSig = meta?.fnPtrSigs.get(i);
            if (fnSig) return `${name}: (${renderFnPtrType(fnSig)}) | number`;
            return `${name}: ${tsType(t)}`;
        })
        .join(', ');
    let ret: string;
    if (results.length === 0) ret = 'void';
    else if (results.length === 1) ret = tsType(results[0]);
    else ret = `[${results.map(tsType).join(', ')}]`;
    return `(${paramList}) => ${ret}`;
}

export function renderDts(wasmBytes: Uint8Array, meta?: Map<string, ExportMeta>): string {
    const sigs = parseWasmExports(wasmBytes);
    const fns = Array.from(sigs.entries())
        .filter(([name]) => !name.startsWith('_'))
        .sort(([a], [b]) => a.localeCompare(b));

    return [
        '// Auto-generated by wasmPlugin. Do not edit.',
        'declare const _: {',
        ...fns.map(([name, sig]) => `    ${name}: ${tsSignature(sig, meta?.get(name))};`),
        '    memory: WebAssembly.Memory;',
        '    table: WebAssembly.Table;',
        '    readonly data: DataView;',
        '};',
        'export default _;',
        '',
    ].join('\n');
}