// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// Self-hosting entry point for JS proactor wasm binaries.
//
// This file serves as both the orchestrator (main thread) and the wasm runner
// (Worker thread) for proactor-based wasm binaries. The bundler appends it
// after the inlined companions, so createWasmImports() is available in module
// scope. The bundler also defines __IREE_WASM_BINARY with the .wasm filename.
//
// The bundler resolves relative imports (ProactorEventHost, ProactorRing) and
// inlines them into the bundle, making the output fully self-contained.
//
// When `node bundle.mjs` runs:
//
//   Main thread (isMainThread === true):
//     1. Creates SharedArrayBuffer ring and cancel control region
//     2. Starts the event host (timer management, cancel listener)
//     3. Spawns itself as a Worker via `new Worker(import.meta.url, ...)`
//     4. Relays worker messages to the event host
//     5. Exits with the wasm program's exit code
//
//   Worker thread (isMainThread === false):
//     1. Reads workerData for shared buffers and arguments
//     2. Sets up WASI for stdio, filesystem, and clocks
//     3. Calls createWasmImports(context) and merges companion imports
//     4. Compiles and instantiates the wasm module
//     5. Calls wasi.start() and reports the exit code to the parent

import {readFileSync} from 'node:fs';
import {dirname, resolve} from 'node:path';
import {fileURLToPath} from 'node:url';
import {WASI} from 'node:wasi';
import {isMainThread, parentPort, Worker, workerData} from 'node:worker_threads';

import {ProactorEventHost} from './proactor_event_host.mjs';
import {ProactorRing} from './proactor_ring.mjs';

// Ring capacity matches the default token table capacity in the C proactor.
const RING_CAPACITY = 256;

if (isMainThread) {
  // Main thread: create the event host infrastructure, spawn the wasm worker.
  const testArgs = process.argv.slice(2);

  const ringBuffer = ProactorRing.createSharedBuffer(RING_CAPACITY);
  const cancelBuffer = new SharedArrayBuffer(16);

  const eventHost =
      new ProactorEventHost(ringBuffer, cancelBuffer, RING_CAPACITY);
  eventHost.start();

  const worker = new Worker(fileURLToPath(import.meta.url), {
    workerData: {
      wasmArgs: testArgs,
      ringBuffer,
      cancelBuffer,
      ringCapacity: RING_CAPACITY,
    },
  });

  worker.on('message', (message) => {
    if (message.type === 'exit') {
      eventHost.stop();
      process.exit(message.status);
    } else {
      eventHost.handleMessage(message);
    }
  });

  worker.on('error', (error) => {
    console.error('Proactor worker error:', error);
    eventHost.stop();
    process.exit(1);
  });

  // Safety net: if the worker exits without sending an exit message (crash,
  // unhandled exception), propagate a failure exit code.
  worker.on('exit', (code) => {
    eventHost.stop();
    process.exit(code || 1);
  });
} else {
  // Worker thread: set up WASI, compile and run the wasm binary.
  const {wasmArgs, ringBuffer, cancelBuffer, ringCapacity} = workerData;

  // Resolve the wasm binary from the bundler-defined filename.
  const scriptDirectory = dirname(fileURLToPath(import.meta.url));
  const wasmPath = resolve(scriptDirectory, __IREE_WASM_BINARY);

  // Build preopens for directories the test needs to access. Bazel sets
  // TEST_TMPDIR for scratch space and XML_OUTPUT_FILE for gtest XML output.
  const preopens = {};
  if (process.env.TEST_TMPDIR) {
    preopens[process.env.TEST_TMPDIR] = process.env.TEST_TMPDIR;
  }
  if (process.env.XML_OUTPUT_FILE) {
    const xmlDirectory = dirname(process.env.XML_OUTPUT_FILE);
    preopens[xmlDirectory] = xmlDirectory;
  }

  // WASI provides the C runtime environment: stdio, filesystem, clocks.
  const wasi = new WASI({
    version: 'preview1',
    args: [__IREE_WASM_BINARY, ...(wasmArgs || [])],
    env: process.env,
    preopens,
  });

  // Build the wasm import context. The `memory` field is set after
  // instantiation (we need to create imports before instantiation, but memory
  // is exported by the wasm module). Each companion's createImports(context)
  // closes over this object, so setting context.memory later is visible to
  // all import functions.
  const context = {
    memory: null,
    ringBuffer,
    ringCapacity,
    cancelControl: cancelBuffer ? new Int32Array(cancelBuffer) : null,
    postMessage: (message) => parentPort.postMessage(message),
  };

  // Start with WASI imports, merge companion imports on top.
  // createWasmImports is generated by the bundler from all collected
  // companions.
  const imports = wasi.getImportObject();
  const companionImports = createWasmImports(context);
  for (const [moduleName, moduleImports] of Object.entries(companionImports)) {
    if (imports[moduleName]) {
      Object.assign(imports[moduleName], moduleImports);
    } else {
      imports[moduleName] = moduleImports;
    }
  }

  // Compile and instantiate.
  const wasmBytes = readFileSync(wasmPath);
  const {instance} = await WebAssembly.instantiate(wasmBytes, imports);

  // Wire up exported memory for companions that need it (e.g., reading
  // wasm linear memory for ring drain operations).
  context.memory = instance.exports.memory;

  // WASI.start() calls _start() on the wasm instance. The WASI runtime
  // handles proc_exit internally: it sets process.exitCode and throws to
  // unwind.
  let exitStatus = 0;
  try {
    wasi.start(instance);
  } catch (error) {
    // WASI proc_exit throws to unwind the wasm call stack. process.exitCode
    // is already set by the WASI runtime. If it was not set, this is a
    // genuine error (not a proc_exit).
    if (process.exitCode != null) {
      exitStatus = process.exitCode;
      // Reset so it doesn't affect the parent process.
      process.exitCode = undefined;
    } else {
      console.error(error);
      exitStatus = 1;
    }
  }

  parentPort.postMessage({type: 'exit', status: exitStatus});
}
