import { ArrayPrototypeMap } from "ext:core/primordials";
import {
  op_ffi_address_of,
  op_ffi_await,
  op_ffi_call,
  op_ffi_call_async_begin,
  op_ffi_call_async_finish,
  op_ffi_call_ptr,
  op_ffi_close,
  op_ffi_dlopen,
  op_ffi_read,
  op_ffi_read_bytes,
  op_ffi_read_cstring,
  op_ffi_write,
  op_ffi_write_bytes,
} from "ext:core/ops";

import type { NativeDef, NativeValue } from "./types.js";

export type NativeType =
  | "void"
  | "bool"
  | "i8" | "u8" | "i16" | "u16" | "i32" | "u32" | "i64" | "u64"
  | "isize" | "usize" | "f32" | "f64"
  | "pointer" | "buffer" | "function";

export interface ForeignFunction {
  parameters: readonly NativeType[];
  result: NativeType;
  /** Run on a native thread and return a Promise; buffer arguments are copied in and back. */
  nonblocking?: boolean;
  /** Symbol name in the library, if different from the property key. */
  name?: string;
}

/**
 * The JS-side argument for a native parameter. Pointer and function
 * arguments accept any view (not just `Uint8Array`) and may be `undefined`.
 */
export type JsArg<T extends NativeType> = T extends "i64" | "u64" | "isize" | "usize"
  ? bigint | number
  : T extends "pointer" | "function" ?
      | Pointer
      | UnsafeFnPointer
      | ArrayBufferView
      | ArrayBuffer
      | bigint
      | null
      | undefined
  : T extends "buffer" ? ArrayBuffer | ArrayBufferView
  : T extends "bool" ? boolean
  : number;

export type JsResult<T extends NativeType> = T extends "void" ? undefined
  : T extends "i64" | "u64" | "isize" | "usize" ? bigint
  : T extends "pointer" | "function" ? Pointer | null
  : T extends "bool" ? boolean
  : number;

export type CallArgs<P extends readonly NativeType[]> = P extends readonly [...infer K]
  ? { [I in keyof K]: K[I] extends NativeType ? JsArg<K[I]> : never }
  : JsArg<P[number]>[];

/**
 * A bound symbol. A `boolean`-typed `nonblocking` flag resolves to a
 * sync-or-async result.
 */
export type Callable<F extends ForeignFunction> = (
  ...args: CallArgs<F["parameters"]>
) => F extends { nonblocking: true } ? Promise<JsResult<F["result"]>>
  : "nonblocking" extends keyof F
    ? true extends F["nonblocking"]
      ? JsResult<F["result"]> | Promise<JsResult<F["result"]>>
    : JsResult<F["result"]>
  : JsResult<F["result"]>;

// A structurally empty instance type, exactly like the deleted declaration's
// `class Pointer { private constructor(); }`. Only pointers created by this
// module satisfy it at runtime; anything else throws on use.
export interface Pointer {}

export interface UnsafePointerView {
  readonly address: Pointer;
  getBool(offset?: number): boolean;
  getInt8(offset?: number): number;
  getUint8(offset?: number): number;
  getInt16(offset?: number): number;
  getUint16(offset?: number): number;
  getInt32(offset?: number): number;
  getUint32(offset?: number): number;
  getBigInt64(offset?: number): bigint;
  getBigUint64(offset?: number): bigint;
  getFloat32(offset?: number): number;
  getFloat64(offset?: number): number;
  getPointer(offset?: number): Pointer | null;
  setBool(value: boolean, offset?: number): void;
  setInt8(value: number, offset?: number): void;
  setUint8(value: number, offset?: number): void;
  setInt16(value: number, offset?: number): void;
  setUint16(value: number, offset?: number): void;
  setInt32(value: number, offset?: number): void;
  setUint32(value: number, offset?: number): void;
  setBigInt64(value: bigint | number, offset?: number): void;
  setBigUint64(value: bigint | number, offset?: number): void;
  setFloat32(value: number, offset?: number): void;
  setFloat64(value: number, offset?: number): void;
  setPointer(value: Pointer | UnsafeFnPointer | bigint | null, offset?: number): void;
  getArrayBuffer(byteLength: number, offset?: number): ArrayBuffer;
  getCString(offset?: number): string;
  copyInto(destination: ArrayBuffer | ArrayBufferView, offset?: number): void;
  writeFrom(source: ArrayBuffer | ArrayBufferView, offset?: number): void;
}

export interface UnsafeFnPointer<F extends ForeignFunction = ForeignFunction> {
  call(...args: CallArgs<F["parameters"]>): JsResult<F["result"]>;
}

export interface DynamicLibrary<S extends Record<string, ForeignFunction>> {
  readonly symbols: { readonly [K in keyof S]: Callable<S[K]> };
  readonly closed: boolean;
  close(): void;
  [Symbol.dispose](): void;
}

type WrappedValue = Pointer | null | NativeValue;

const TYPES = new Set<unknown>([
  "void",
  "bool",
  "i8",
  "u8",
  "i16",
  "u16",
  "i32",
  "u32",
  "i64",
  "u64",
  "isize",
  "usize",
  "f32",
  "f64",
  "pointer",
  "buffer",
  "function",
]);

function normalizeDef(key: string, def: unknown): NativeDef {
  if (def === null || typeof def !== "object" || Array.isArray(def)) {
    throw new TypeError(`Symbol ${key} must be a definition object`);
  }
  const { parameters, result, nonblocking = false, name } = def as {
    parameters: unknown;
    result: unknown;
    nonblocking?: unknown;
    name?: unknown;
  };
  if (!Array.isArray(parameters)) {
    throw new TypeError(`Symbol ${key} must have a parameters array`);
  }
  for (const type of parameters as unknown[]) {
    if (!TYPES.has(type)) throw new TypeError(`Unknown native type: ${type}`);
    if (type === "void") throw new TypeError("void is not a parameter type");
  }
  if (!TYPES.has(result)) throw new TypeError(`Unknown native type: ${result}`);
  if (result === "buffer") throw new TypeError("buffer is not a result type");
  if (typeof nonblocking !== "boolean") {
    throw new TypeError("nonblocking must be a boolean");
  }
  return {
    symbol: name ?? key,
    params: [...parameters] as NativeDef["params"],
    result: result as NativeDef["result"],
    nonblocking,
  };
}

const addresses = new WeakMap<object, bigint>();

const PointerImplementation = class Pointer {
  declare private readonly pointerBrand: void;
  declare static readonly [Symbol.hasInstance]: (value: unknown) => value is Pointer;
  constructor(address: bigint | number) {
    addresses.set(this, BigInt(address));
    Object.freeze(this);
  }
};

export const Pointer: {
  prototype: Pointer;
  [Symbol.hasInstance](value: unknown): value is Pointer;
} = PointerImplementation;

function addressOf(pointer: Pointer): bigint {
  const address = addresses.get(pointer);
  if (address === undefined) throw new TypeError("Expected a Pointer");
  return address;
}

function unwrap(value: unknown): unknown {
  if (value instanceof PointerImplementation) return addressOf(value);
  if (value instanceof UnsafeFnPointerImplementation) return value.address;
  return value;
}

export const UnsafePointer: Readonly<{
  /** The pointer does not keep the view alive; it must not outlive it. */
  of(view: ArrayBuffer | ArrayBufferView): Pointer;
  value(pointer: Pointer): bigint;
  offset(pointer: Pointer, amount: number | bigint): Pointer;
  create(length: number): [Pointer, Uint8Array];
  equals(a: Pointer, b: Pointer): boolean;
}> = Object.freeze({
  of(view: ArrayBuffer | ArrayBufferView): Pointer {
    // The resulting Pointer does not keep the view alive; it must not
    // outlive the view it was taken from.
    return new PointerImplementation(op_ffi_address_of(view));
  },
  value(pointer: Pointer): bigint {
    return addressOf(pointer);
  },
  offset(pointer: Pointer, amount: number | bigint): Pointer {
    return new PointerImplementation(addressOf(pointer) + BigInt(amount));
  },
  create(length: number): [Pointer, Uint8Array<ArrayBuffer>] {
    const view = new Uint8Array(length);
    return [new PointerImplementation(op_ffi_address_of(view)), view];
  },
  equals(a: Pointer, b: Pointer): boolean {
    return addressOf(a) === addressOf(b);
  },
});

export const UnsafePointerView: {
  prototype: UnsafePointerView;
  new (pointer: Pointer): UnsafePointerView;
} = class UnsafePointerView {
  #address: bigint;
  constructor(pointer: Pointer | bigint | number) {
    this.#address = pointer instanceof Pointer
      ? addressOf(pointer)
      : BigInt(pointer as bigint | number);
  }
  get address(): Pointer {
    return new PointerImplementation(this.#address);
  }
  getBool(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "bool");
  }
  getInt8(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "i8");
  }
  getUint8(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "u8");
  }
  getInt16(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "i16");
  }
  getUint16(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "u16");
  }
  getInt32(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "i32");
  }
  getUint32(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "u32");
  }
  getBigInt64(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "i64");
  }
  getBigUint64(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "u64");
  }
  getFloat32(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "f32");
  }
  getFloat64(offset = 0) {
    return op_ffi_read(this.#address + BigInt(offset), "f64");
  }
  getPointer(offset = 0) {
    const value = op_ffi_read(this.#address + BigInt(offset), "pointer");
    return value === 0n ? null : new PointerImplementation(value);
  }
  setBool(value: boolean, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "bool", value);
  }
  setInt8(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "i8", value);
  }
  setUint8(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "u8", value);
  }
  setInt16(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "i16", value);
  }
  setUint16(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "u16", value);
  }
  setInt32(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "i32", value);
  }
  setUint32(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "u32", value);
  }
  setBigInt64(value: bigint | number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "i64", value);
  }
  setBigUint64(value: bigint | number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "u64", value);
  }
  setFloat32(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "f32", value);
  }
  setFloat64(value: number, offset = 0) {
    op_ffi_write(this.#address + BigInt(offset), "f64", value);
  }
  setPointer(
    value: Pointer | UnsafeFnPointer | bigint | null,
    offset = 0,
  ) {
    op_ffi_write(
      this.#address + BigInt(offset),
      "pointer",
      unwrap(value) ?? 0n,
    );
  }
  getArrayBuffer(byteLength: number, offset = 0): ArrayBuffer {
    const bytes = op_ffi_read_bytes(this.#address + BigInt(offset), byteLength);
    return bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + byteLength);
  }
  getCString(offset = 0): string {
    const bytes = op_ffi_read_cstring(this.#address + BigInt(offset), 1048576);
    return new TextDecoder().decode(bytes);
  }
  copyInto(destination: ArrayBuffer | ArrayBufferView, offset = 0): void {
    const bytes = op_ffi_read_bytes(
      this.#address + BigInt(offset),
      destination.byteLength,
    );
    new Uint8Array(
      ((destination as ArrayBufferView).buffer ??
        destination) as ArrayBufferLike,
      (destination as ArrayBufferView).byteOffset ?? 0,
      destination.byteLength,
    ).set(bytes);
  }
  writeFrom(source: ArrayBuffer | ArrayBufferView, offset = 0): void {
    op_ffi_write_bytes(this.#address + BigInt(offset), source);
  }
};

const UnsafeFnPointerImplementation = class UnsafeFnPointer<F extends ForeignFunction = ForeignFunction> {
  #def: NativeDef;
  readonly address: bigint;
  constructor(pointer: Pointer | bigint | number, def: F) {
    if (def?.nonblocking) {
      throw new TypeError("Function pointers cannot be called nonblocking");
    }
    const normalized = normalizeDef("function pointer", {
      parameters: def.parameters,
      result: def.result,
      nonblocking: false,
    });
    this.#def = normalized;
    this.address = pointer instanceof Pointer
      ? addressOf(pointer)
      : BigInt(pointer as bigint | number);
    Object.freeze(this);
  }
  call(...args: CallArgs<F["parameters"]>): JsResult<F["result"]> {
    if (this.address === 0n) throw new TypeError("Null function pointer");
    return wrapResult(
      this.#def.result,
      op_ffi_call_ptr(this.address, this.#def, ArrayPrototypeMap(args, unwrap)),
    ) as JsResult<F["result"]>;
  }
};

export const UnsafeFnPointer: {
  prototype: UnsafeFnPointer;
  new <F extends ForeignFunction>(pointer: Pointer, def: F): UnsafeFnPointer<F>;
} = UnsafeFnPointerImplementation;

function wrapResult(
  result: NativeDef["result"],
  value: NativeValue,
): WrappedValue {
  if (result === "pointer" || result === "function") {
    return value === 0n ? null : new PointerImplementation(value as bigint);
  }
  return value;
}

const libraryIds = new WeakMap<object, number>();
const libraryToken = Symbol("library");

const DynamicLibraryImplementation = class DynamicLibrary<
  S extends Record<string, ForeignFunction> = Record<string, ForeignFunction>,
> {
  readonly symbols: { readonly [K in keyof S]: Callable<S[K]> };
  #closed = false;
  declare static readonly [Symbol.hasInstance]: (value: unknown) => value is DynamicLibrary;
  constructor(
    token: typeof libraryToken,
    id: number,
    symbols: { [K in keyof S]: Callable<S[K]> },
  ) {
    if (token !== libraryToken) {
      throw new TypeError("DynamicLibrary cannot be constructed directly");
    }
    libraryIds.set(this, id);
    this.symbols = Object.freeze({ ...symbols });
    Object.freeze(this);
  }
  close(): void {
    if (this.#closed) return;
    this.#closed = true;
    op_ffi_close(libraryIds.get(this)!);
  }
  get closed(): boolean {
    return this.#closed;
  }
  [Symbol.dispose](): void {
    this.close();
  }
};

export const DynamicLibrary: {
  prototype: DynamicLibrary<Record<string, ForeignFunction>>;
  [Symbol.hasInstance](value: unknown): value is DynamicLibrary<Record<string, ForeignFunction>>;
} = DynamicLibraryImplementation;

/**
 * Loads a shared library and binds its symbols. Native code runs
 * in-process and is not sandboxed; loading requires ffi authority for
 * the library path (a bare loader-searched name requires unrestricted
 * ffi authority).
 */
export function dlopen<const S extends Record<string, ForeignFunction>>(
  path: string | URL,
  symbols: S,
): DynamicLibrary<S> {
  if (symbols === null || typeof symbols !== "object") {
    throw new TypeError("dlopen requires a symbols object");
  }
  const keys = Object.keys(symbols);
  const defs = ArrayPrototypeMap(
    keys,
    (key) => normalizeDef(key, symbols[key]),
  );
  const id = op_ffi_dlopen(
    `${path instanceof URL ? path : String(path)}`,
    defs,
  );
  const wrapped: Record<
    string,
    (...args: unknown[]) => WrappedValue | Promise<WrappedValue>
  > = {};
  for (let index = 0; index < keys.length; index++) {
    const key = keys[index];
    const def = defs[index];
    wrapped[key] = (...args: unknown[]) => {
      if (args.length !== def.params.length) {
        throw new TypeError(
          `${key} expects ${def.params.length} arguments, got ${args.length}`,
        );
      }
      const unwrapped = ArrayPrototypeMap(args, unwrap);
      if (def.nonblocking) {
        return (async () => {
          const callId = op_ffi_call_async_begin(id, index, unwrapped);
          await op_ffi_await(callId);
          return wrapResult(
            def.result,
            op_ffi_call_async_finish(callId, def.result),
          );
        })();
      }
      return wrapResult(def.result, op_ffi_call(id, index, unwrapped));
    };
  }
  // Dynamic key iteration loses the relation between each definition and its wrapper.
  return new DynamicLibraryImplementation(
    libraryToken,
    id,
    wrapped as unknown as { [K in keyof S]: Callable<S[K]> },
  );
}