import {
  exceptionsFinal,
  multiMemory,
  relaxedSimd,
  simd,
  tailCall,
  typedFunctionReferences,
} from "wasm-feature-detect";
import initBrowserRuntime, {
  clear_filesystem_build_cache,
  install_filesystem_snapshot,
  normalize_filesystem_snapshot,
  start_fixture,
  start_fixture_with_hot_threshold,
  start_program,
  translate_fixture_entry,
  write_filesystem_file,
} from "./generated/binarrow_browser_runtime.js";
import type { BrowserGuestSession } from "./generated/binarrow_browser_runtime.js";

import type {
  ExecutionReport,
  FeatureName,
  FeatureResult,
  FixtureName,
  RunTarget,
  TranslationBenchmarkReport,
  WorkerCommand,
  WorkerMessage,
} from "./probe";

interface Memory64Exports extends WebAssembly.Exports {
  memory: WebAssembly.Memory;
  memory_pages: () => bigint;
  probe_marker: WebAssembly.Global;
}

interface SemanticProbeExports extends WebAssembly.Exports {
  binarrow_aarch64_semantic_probe: () => bigint;
}

interface JspiProbeExports extends WebAssembly.Exports {
  run: () => number;
}

interface TranslatedBlockExports extends WebAssembly.Exports {
  run: (stateBase: number) => bigint;
}

interface JspiApi {
  Suspending: new (callback: () => Promise<number>) => WebAssembly.ImportValue;
  promising: (callback: () => number) => () => Promise<number>;
}

const filesystemSnapshotName = "binarrow-project-v2.snapshot";
const translatedLoopInstructions = 17n;

async function loadFilesystemSnapshot(): Promise<Uint8Array> {
  const root = await navigator.storage.getDirectory();
  try {
    const handle = await root.getFileHandle(filesystemSnapshotName);
    return new Uint8Array(await (await handle.getFile()).arrayBuffer());
  } catch (error) {
    if (error instanceof DOMException && error.name === "NotFoundError") {
      return new Uint8Array();
    }
    throw error;
  }
}

async function saveFilesystemSnapshot(snapshot: Uint8Array): Promise<void> {
  if (snapshot.byteLength === 0) {
    return;
  }

  const root = await navigator.storage.getDirectory();
  const handle = await root.getFileHandle(filesystemSnapshotName, {
    create: true,
  });
  const writable = await handle.createWritable();
  try {
    await writable.write(snapshot as Uint8Array<ArrayBuffer>);
  } finally {
    await writable.close();
  }
}

async function instantiateMemory64Probe(): Promise<boolean> {
  const response = await fetch("/memory64-probe.wasm");
  if (!response.ok) {
    throw new Error(`Memory64 fixture returned HTTP ${response.status}`);
  }

  const bytes = await response.arrayBuffer();
  const { instance } = await WebAssembly.instantiate(bytes);
  const exports = instance.exports as Memory64Exports;

  return (
    exports.memory instanceof WebAssembly.Memory &&
    exports.memory_pages() === 1n &&
    exports.probe_marker.value === 0x4249_4e41
  );
}

async function instantiateSemanticProbe(): Promise<boolean> {
  const response = await fetch("/icicle-semantic-probe.wasm");
  if (!response.ok) {
    throw new Error(`AArch64 semantic fixture returned HTTP ${response.status}`);
  }

  const bytes = await response.arrayBuffer();
  const { instance } = await WebAssembly.instantiate(bytes);
  const exports = instance.exports as SemanticProbeExports;

  return exports.binarrow_aarch64_semantic_probe() === 43n;
}

async function runJspiRoundTrip(): Promise<boolean> {
  const jspi = WebAssembly as typeof WebAssembly & JspiApi;
  if (typeof jspi.Suspending !== "function" || typeof jspi.promising !== "function") {
    return false;
  }

  const response = await fetch("/jspi-probe.wasm");
  if (!response.ok) {
    throw new Error(`JSPI fixture returned HTTP ${response.status}`);
  }

  let hostPromiseResolved = false;
  const value = new jspi.Suspending(
    () =>
      new Promise<number>((resolve) => {
        setTimeout(() => {
          hostPromiseResolved = true;
          resolve(41);
        }, 0);
      }),
  );
  const bytes = await response.arrayBuffer();
  const { instance } = await WebAssembly.instantiate(bytes, {
    host: { value },
  });
  const exports = instance.exports as JspiProbeExports;
  const run = jspi.promising(exports.run);
  const pendingResult = run();

  if (!(pendingResult instanceof Promise) || hostPromiseResolved) {
    return false;
  }

  return (await pendingResult) === 42 && hostPromiseResolved;
}

async function runTranslatedBlockProbe(): Promise<boolean> {
  const translation = translate_fixture_entry("translated-loop", 32);
  try {
    const state = new WebAssembly.Memory({ initial: 1, maximum: 1 });
    const instantiated = await WebAssembly.instantiate(translation.wasm_module, {
      binarrow: { state },
    });
    const instance =
      (instantiated as unknown as WebAssembly.WebAssemblyInstantiatedSource)
        .instance ??
      (instantiated as WebAssembly.Instance);
    const exports = instance.exports as TranslatedBlockExports;
    const initialState = translation.initial_state;
    const stateBytes = new Uint8Array(state.buffer, 0, initialState.length);
    stateBytes.set(initialState);
    const nextPc = exports.run(0);
    const expectedState = translation.interpreted_state;
    return (
      nextPc === translation.interpreted_next_pc &&
      stateBytes.every((byte, index) => byte === expectedState[index]) &&
      translation.guest_instructions === 17 &&
      translation.wasm_bytes === translation.wasm_module.byteLength
    );
  } finally {
    translation.free();
  }
}

async function runProbe(
  name: FeatureName,
  probe: () => Promise<boolean>,
): Promise<FeatureResult> {
  try {
    const supported = await probe();
    return { name, supported };
  } catch (error) {
    return {
      name,
      supported: false,
      detail: error instanceof Error ? error.message : String(error),
    };
  }
}

async function createReport(): Promise<FeatureResult[]> {
  return Promise.all([
    runProbe("Memory64", instantiateMemory64Probe),
    runProbe("AArch64 P-code semantics", instantiateSemanticProbe),
    runProbe("Dynamic block translation", runTranslatedBlockProbe),
    runProbe("JSPI", runJspiRoundTrip),
    runProbe("Tail calls", tailCall),
    runProbe("SIMD", simd),
    runProbe("Relaxed SIMD", relaxedSimd),
    runProbe("Multiple memories", multiMemory),
    runProbe("Typed function references", typedFunctionReferences),
    runProbe("Exception handling", exceptionsFinal),
  ]);
}

async function initializeRuntime(): Promise<void> {
  await initBrowserRuntime({
    module_or_path: new URL(
      "./generated/binarrow_browser_runtime_bg.wasm",
      import.meta.url,
    ),
  });
}

function executionReport(
  result: ReturnType<BrowserGuestSession["resume"]>,
): ExecutionReport {
  const diagnosticCode = result.diagnostic_code;
  return {
    outcome: result.outcome as ExecutionReport["outcome"],
    ...(diagnosticCode
      ? {
          diagnostic: {
            code: diagnosticCode,
            message: result.diagnostic_message,
          },
        }
      : {}),
    stdout: result.stdout,
    stderr: result.stderr,
    trace: result.trace,
    exitCode: result.exit_code < 0 ? null : result.exit_code,
    executedInstructions: result.executed_instructions,
    dispatchedSyscalls: result.dispatched_syscalls,
    outputBytes: result.output_bytes,
    profiledBlocks: result.profiled_blocks,
    hottestBlockStart:
      result.hottest_block_executions === 0n
        ? null
        : result.hottest_block_start,
    hottestBlockExecutions: result.hottest_block_executions,
    translatedBlockExecutions: result.translated_block_executions,
    translatedInstructions: result.translated_instructions,
    translationFallbackBlocks: result.translation_fallback_blocks,
    compiledBlocks: result.compiled_blocks,
    translationCacheHits: result.translation_cache_hits,
    compiledWasmBytes: result.compiled_wasm_bytes,
    inputMaxBytes: result.input_max_bytes,
  };
}

let activeSession: BrowserGuestSession | undefined;
let activeRunId = 0;
let activeTarget: RunTarget = "compiler-c";

async function continueSession(
  input: Uint8Array,
  endOfFile: boolean,
): Promise<ExecutionReport> {
  if (!activeSession) {
    throw new Error("no guest session is waiting for input");
  }
  const result = activeSession.resume(input, endOfFile);
  try {
    await saveFilesystemSnapshot(result.filesystem_snapshot);
    const report = executionReport(result);
    if (report.outcome !== "waiting-input") {
      activeSession.free();
      activeSession = undefined;
    }
    return report;
  } finally {
    result.free();
  }
}

async function execute(
  command: Extract<WorkerCommand, { kind: "run" | "run-program" }>,
): Promise<ExecutionReport> {
  const filesystemSnapshot = await loadFilesystemSnapshot();
  activeSession?.free();
  activeSession =
    command.kind === "run"
      ? start_fixture(
          command.fixture,
          command.limits.instructionBudget,
          command.limits.syscallBudget,
          command.limits.maxOutputBytes,
          command.limits.maxMemoryBytes,
          command.limits.maxFilesystemBytes,
          filesystemSnapshot,
        )
      : start_program(
          command.executable,
          command.argv0,
          command.arguments.join("\0"),
          command.environment.join("\0"),
          command.limits.instructionBudget,
          command.limits.syscallBudget,
          command.limits.maxOutputBytes,
          command.limits.maxMemoryBytes,
          command.limits.maxFilesystemBytes,
          filesystemSnapshot,
        );
  activeRunId = command.runId;
  activeTarget = command.kind === "run" ? command.fixture : "uploaded-program";
  return continueSession(new Uint8Array(), false);
}

async function handleFilesystemCommand(
  command: Extract<
    WorkerCommand,
    {
      kind:
        | "filesystem-export"
        | "filesystem-import"
        | "filesystem-write"
        | "filesystem-clear-cache";
    }
  >,
): Promise<void> {
  const current = await loadFilesystemSnapshot();
  let snapshot: Uint8Array;
  let operation: "export" | "replace" | "install" | "write" | "clear-cache";
  if (command.kind === "filesystem-export") {
    snapshot = normalize_filesystem_snapshot(
      command.maxFilesystemBytes,
      current,
    );
    operation = "export";
  } else if (command.kind === "filesystem-write") {
    snapshot = write_filesystem_file(
      command.maxFilesystemBytes,
      current,
      command.path,
      command.contents,
    );
    await saveFilesystemSnapshot(snapshot);
    operation = "write";
  } else if (command.kind === "filesystem-clear-cache") {
    snapshot = clear_filesystem_build_cache(
      command.maxFilesystemBytes,
      current,
    );
    await saveFilesystemSnapshot(snapshot);
    operation = "clear-cache";
  } else if (command.mode === "replace") {
    snapshot = normalize_filesystem_snapshot(
      command.maxFilesystemBytes,
      command.snapshot,
    );
    await saveFilesystemSnapshot(snapshot);
    operation = "replace";
  } else {
    snapshot = install_filesystem_snapshot(
      command.maxFilesystemBytes,
      current,
      command.snapshot,
    );
    await saveFilesystemSnapshot(snapshot);
    operation = "install";
  }
  const message: WorkerMessage = {
    kind: "filesystem-image",
    requestId: command.requestId,
    operation,
    snapshot,
  };
  self.postMessage(message);
}

interface BenchmarkSample {
  milliseconds: number;
  guestInstructions: bigint;
  translatedBlocks: bigint;
  fallbackBlocks: bigint;
}

function runBenchmarkSample(
  iterations: number,
  hotBlockThreshold: bigint,
): BenchmarkSample {
  const instructionBudget = BigInt(iterations) * translatedLoopInstructions;
  const session = start_fixture_with_hot_threshold(
    "translated-loop",
    instructionBudget,
    1_000n,
    1_024n,
    256n * 1_024n * 1_024n,
    1_024n * 1_024n,
    new Uint8Array(),
    hotBlockThreshold,
  );
  const started = performance.now();
  const result = session.resume(new Uint8Array(), false);
  const milliseconds = performance.now() - started;
  try {
    if (result.diagnostic_code !== "resource.instructions") {
      throw new Error(
        `translation benchmark ended with ${result.diagnostic_code || result.outcome}`,
      );
    }
    return {
      milliseconds,
      guestInstructions: result.executed_instructions,
      translatedBlocks: result.translated_block_executions,
      fallbackBlocks: result.translation_fallback_blocks,
    };
  } finally {
    result.free();
    session.free();
  }
}

function median(values: number[]): number {
  const ordered = [...values].sort((left, right) => left - right);
  return ordered[Math.floor(ordered.length / 2)] ?? Number.NaN;
}

function runTranslationBenchmark(iterations: number): TranslationBenchmarkReport {
  if (!Number.isSafeInteger(iterations) || iterations < 128 || iterations > 100_000) {
    throw new Error("translation benchmark iterations must be an integer from 128 to 100000");
  }
  const warmupIterations = Math.min(iterations, 256);
  runBenchmarkSample(warmupIterations, 0n);
  runBenchmarkSample(warmupIterations, 1n);

  const interpreted: BenchmarkSample[] = [];
  const translated: BenchmarkSample[] = [];
  for (let sample = 0; sample < 5; sample += 1) {
    interpreted.push(runBenchmarkSample(iterations, 0n));
    translated.push(runBenchmarkSample(iterations, 1n));
  }
  const interpretedMilliseconds = median(
    interpreted.map((sample) => sample.milliseconds),
  );
  const translatedMilliseconds = median(
    translated.map((sample) => sample.milliseconds),
  );
  const translatedSample = translated[translated.length - 1];
  if (!translatedSample) {
    throw new Error("translation benchmark did not produce a sample");
  }
  return {
    iterations,
    guestInstructions: translatedSample.guestInstructions,
    interpretedMilliseconds,
    translatedMilliseconds,
    speedup: interpretedMilliseconds / translatedMilliseconds,
    translatedBlocks: translatedSample.translatedBlocks,
    fallbackBlocks: translatedSample.fallbackBlocks,
  };
}

const ready = initializeRuntime().then(createReport);

ready
  .then((results) => {
    const message: WorkerMessage = { kind: "ready", results };
    self.postMessage(message);
  })
  .catch((error) => {
    const message: WorkerMessage = {
      kind: "error",
      message: error instanceof Error ? error.message : String(error),
    };
    self.postMessage(message);
  });

self.addEventListener("message", (event: MessageEvent<WorkerCommand>) => {
  const command = event.data;
  void ready
    .then(async () => {
      if (command.kind === "translation-benchmark") {
        const message: WorkerMessage = {
          kind: "translation-benchmark",
          requestId: command.requestId,
          report: runTranslationBenchmark(command.iterations),
        };
        self.postMessage(message);
        return;
      }
      if (
        command.kind === "filesystem-export" ||
        command.kind === "filesystem-import" ||
        command.kind === "filesystem-write" ||
        command.kind === "filesystem-clear-cache"
      ) {
        await handleFilesystemCommand(command);
        return;
      }
      if (command.kind === "input" && command.runId !== activeRunId) {
        return;
      }
      const execution =
        command.kind === "run" || command.kind === "run-program"
          ? await execute(command)
          : await continueSession(
              new TextEncoder().encode(command.input),
              command.endOfFile,
            );
      const message: WorkerMessage = {
        kind: "result",
        runId: command.runId,
        target:
          command.kind === "run"
            ? command.fixture
            : command.kind === "run-program"
              ? "uploaded-program"
              : activeTarget,
        execution,
      };
      self.postMessage(message);
    })
    .catch((error) => {
      const message: WorkerMessage = {
        kind: "error",
        message: error instanceof Error ? error.message : String(error),
      };
      self.postMessage(message);
    });
});