import {
  ArrayBufferPrototypeGetResizable as resizable,
  Map,
  MapPrototypeDelete as mapDelete,
  MapPrototypeGet as mapGet,
  MapPrototypeSet as mapSet,
  ObjectDefineProperty as define,
  ReflectApply as apply,
  TypedArrayPrototypeGetBuffer as buffer,
  TypedArrayPrototypeGetTag as tag,
  TypeError,
  Uint8Array,
} from "ext:core/primordials";
import {
  op_wasm_streaming_abort as abort,
  op_wasm_streaming_finish as finish,
  op_wasm_streaming_register as register,
  op_wasm_streaming_url as setUrl,
  op_wasm_streaming_write as write,
} from "ext:core/ops";
import { markPromiseHandled, safePromise } from "ext:core/promises";
import { onRuntimeStart } from "ext:core/bootstrap";
import { messageState, usable } from "ext:http_types/body.js";
import { getHeader } from "ext:http_types/headers.js";
import { extractMime } from "ext:http_types/mime.js";
import {
  read,
  readerCancel,
  release,
  streamReader,
} from "ext:webstreams/streams.js";
import type {
  ReadableStreamDefaultReader,
} from "ext:webstreams/streams.js";
import type { ResponseState } from "ext:http_types/body.js";
import type {
  CompileOptions,
  CompileStreaming,
  Imports,
  InstantiateStreaming,
  InstantiatedSource,
  WasmModule,
  WasmNamespace,
} from "./types.js";

declare const WebAssembly: WasmNamespace;

let wasm: WasmNamespace;
let nativeCompileStreaming: CompileStreaming;
let nativeInstantiateStreaming: InstantiateStreaming;
const pending = new Map<unknown, unknown>();
const maxToken = 9_007_199_254_740_991;
let nextToken = 0;

function setLength<T extends Function>(fn: T, name: string, length: number): T {
  define(fn, "name", { value: name });
  define(fn, "length", { value: length });
  return fn;
}

async function consume(id: number, source: unknown): Promise<void> {
  let reader: ReadableStreamDefaultReader<unknown> | undefined;
  let body: ResponseState["body"];
  try {
    const state: ResponseState = messageState(source, "response");
    body = state.body;
    setUrl([id, state.url]);
    usable(body);
    if (body.stream === null) {
      finish(id);
      return;
    }
    body.used = true;
    reader = streamReader(body.stream);
    while (true) {
      const result = await safePromise(read(reader));
      if (result.done) break;
      if (tag(result.value) !== "Uint8Array") {
        throw new TypeError(
          "Response body stream must contain Uint8Array chunks",
        );
      }
      let bytes = result.value as Uint8Array<ArrayBufferLike>, copy: boolean;
      try {
        copy = resizable(buffer(bytes));
      } catch {
        copy = true;
      }
      if (copy) bytes = new Uint8Array(bytes);
      write(id, bytes as Uint8Array<ArrayBuffer>);
    }
    finish(id);
  } catch (error) {
    if (reader !== undefined) {
      try {
        markPromiseHandled(readerCancel(reader, error));
      } catch (error) {
        void error;
      }
    }
    try {
      abort(id, error);
    } catch (error) {
      void error;
    }
  } finally {
    if (reader !== undefined) {
      try {
        release(reader);
      } catch (error) {
        void error;
      }
    }
  }
}

function onNativeStart(token: unknown, id: number): void {
  const response = mapGet(pending, token);
  mapDelete(pending, token);
  if (response === undefined) {
    abort(id, new TypeError("Missing WebAssembly streaming response"));
    return;
  }
  consume(id, response);
}

interface StreamingState {
  active: boolean;
  token: string | undefined;
}

async function validatedSource(
  source: unknown,
  state: StreamingState,
): Promise<string> {
  const response = await safePromise(source);
  if (!state.active) return "avril-wasm-stream-inactive";
  const responseState = messageState(response, "response");
  if (responseState.status < 200 || responseState.status > 299) {
    throw new TypeError("WebAssembly response status is not successful");
  }
  const mime = extractMime(getHeader(responseState.headers, "content-type"));
  if (mime?.essence !== "application/wasm") {
    throw new TypeError("WebAssembly response has an invalid MIME type");
  }
  if (nextToken >= maxToken) {
    throw new TypeError("WebAssembly streaming token space exhausted");
  }
  const token = `avril-wasm-stream-${++nextToken}`;
  state.token = token;
  mapSet(pending, token, response);
  return token;
}

async function streamingCall<T>(
  source: unknown,
  invoke: (sourcePromise: unknown) => T,
): Promise<Awaited<T>> {
  const state: StreamingState = { active: true, token: undefined };
  const sourcePromise = validatedSource(source, state);
  markPromiseHandled(sourcePromise);
  try {
    return await safePromise(invoke(safePromise(sourcePromise)));
  } finally {
    state.active = false;
    if (state.token !== undefined) mapDelete(pending, state.token);
  }
}
const methods: {
  compileStreaming: CompileStreaming;
  instantiateStreaming: InstantiateStreaming;
} = {
  compileStreaming(source: unknown, options?: CompileOptions) {
    return streamingCall(
      source,
      (sourcePromise) =>
        apply(nativeCompileStreaming, wasm, [sourcePromise, options]) as Promise<
          WasmModule
        >,
    );
  },
  instantiateStreaming(
    source: unknown,
    imports?: Imports,
    options?: CompileOptions,
  ) {
    return streamingCall(
      source,
      (sourcePromise) =>
        apply(nativeInstantiateStreaming, wasm, [
          sourcePromise,
          imports,
          options,
        ]) as Promise<InstantiatedSource>,
    );
  },
};
const { compileStreaming, instantiateStreaming } = methods;

onRuntimeStart(() => {
  // V8 exposes WebAssembly only in live contexts, not snapshot-building contexts.
  wasm = WebAssembly;
  nativeCompileStreaming = wasm.compileStreaming;
  nativeInstantiateStreaming = wasm.instantiateStreaming;
  register(onNativeStart);
  define(wasm, "compileStreaming", {
    value: setLength(compileStreaming, "compileStreaming", 1),
    writable: true,
    enumerable: true,
    configurable: true,
  });
  define(wasm, "instantiateStreaming", {
    value: setLength(instantiateStreaming, "instantiateStreaming", 2),
    writable: true,
    enumerable: true,
    configurable: true,
  });
});