import { createHash } from "node:crypto";
import { readFileSync, writeFileSync } from "node:fs";
import { resolve } from "node:path";

const [inputArgument, outputArgument] = process.argv.slice(2);
if (!inputArgument || !outputArgument) {
  throw new Error("Usage: node transform-rerun-web-viewer-glue.mjs INPUT OUTPUT");
}

const inputPath = resolve(inputArgument);
const outputPath = resolve(outputArgument);
const expectedGeneratedSha256 =
  "cc196a93c5be972c801d46be4dc9934f7f042eb62941f0aa0678f1c8416c6874";

const sha256 = (value) => createHash("sha256").update(value).digest("hex");
let code = readFileSync(inputPath, "utf8");

if (sha256(code) !== expectedGeneratedSha256) {
  throw new Error(`Refusing to transform unexpected wasm-bindgen output: ${inputPath}`);
}

// This is the same no-modules-base transformation used by Rerun 0.34.1's
// rerun_js/web-viewer/build-wasm.mjs. Each factory call gets isolated closure state.
const wrapperStart = "let wasm_bindgen = (function(exports) {";
const wrapperEnd = `return Object.assign(__wbg_init, { initSync }, exports);
})({ __proto__: null });`;

if (!code.includes(wrapperStart) || !code.includes(wrapperEnd)) {
  throw new Error("Rerun wasm-bindgen wrapper markers no longer match");
}
code = code.replace(wrapperStart, "").replace(wrapperEnd, "");

code = `
export default function() {
const exports = { __proto__: null };
${code}

function deinit() {
  __wbg_init.__wbindgen_wasm_module = null;
  wasmModule = null;
  wasm = null;
  cachedDataViewMemory0 = null;
  cachedFloat32ArrayMemory0 = null;
  cachedInt16ArrayMemory0 = null;
  cachedInt32ArrayMemory0 = null;
  cachedInt8ArrayMemory0 = null;
  cachedUint16ArrayMemory0 = null;
  cachedUint32ArrayMemory0 = null;
  cachedUint8ArrayMemory0 = null;
}

return Object.assign(__wbg_init, { initSync, deinit }, exports);
}
`;

const closureDtorsOriginal = `const CLOSURE_DTORS = (typeof FinalizationRegistry === 'undefined')
        ? { register: () => {}, unregister: () => {} }
        : new FinalizationRegistry(state => wasm.__wbindgen_destroy_closure(state.a, state.b));`;
const closureDtorsPatch = `const CLOSURE_DTORS = (typeof FinalizationRegistry === 'undefined')
        ? { register: () => {}, unregister: () => {} }
        : new FinalizationRegistry(state => {
        if (wasm) wasm.__wbindgen_destroy_closure(state.a, state.b);
    });`;

if (!code.includes(closureDtorsOriginal)) {
  throw new Error("Rerun CLOSURE_DTORS block no longer matches");
}
code = code.replace(closureDtorsOriginal, closureDtorsPatch);

const makeMutClosureOriginal = `function makeMutClosure(arg0, arg1, f) {
        const state = { a: arg0, b: arg1, cnt: 1 };
        const real = (...args) => {

            // First up with a closure we increment the internal reference
            // count. This ensures that the Rust closure environment won't
            // be deallocated while we're invoking it.
            state.cnt++;
            const a = state.a;
            state.a = 0;
            try {
                return f(a, state.b, ...args);
            } finally {
                state.a = a;
                real._wbg_cb_unref();
            }
        };
        real._wbg_cb_unref = () => {
            if (--state.cnt === 0) {
                wasm.__wbindgen_destroy_closure(state.a, state.b);
                state.a = 0;
                CLOSURE_DTORS.unregister(state);
            }
        };
        CLOSURE_DTORS.register(real, state, state);
        return real;
    }`;
const makeMutClosurePatch = `function makeMutClosure(arg0, arg1, f) {
        const state = { a: arg0, b: arg1, cnt: 1 };
        const real = (...args) => {
            state.cnt++;
            const a = state.a;
            state.a = 0;
            try {
                if (!wasm) return;
                return f(a, state.b, ...args);
            } finally {
                state.a = a;
                real._wbg_cb_unref();
            }
        };
        real._wbg_cb_unref = () => {
            if (--state.cnt === 0) {
                if (wasm) wasm.__wbindgen_destroy_closure(state.a, state.b);
                state.a = 0;
                CLOSURE_DTORS.unregister(state);
            }
        };
        CLOSURE_DTORS.register(real, state, state);
        return real;
    }`;

if (!code.includes(makeMutClosureOriginal)) {
  throw new Error("Rerun makeMutClosure block no longer matches");
}
code = code.replace(makeMutClosureOriginal, makeMutClosurePatch);

writeFileSync(outputPath, code);
console.log(`Transformed Rerun wasm glue: ${sha256(code)}  ${outputPath}`);