import {
  initSync,
  WasmNoiseSession,
} from "./noise-core-wasm/peach_noise_core";
import { PEACH_NOISE_CORE_WASM_BASE64 } from "./noise-core-wasm/peach_noise_core_bytes";

export interface RustNoiseSession {
  encrypt(plaintext: Uint8Array): Uint8Array;
  decrypt(ciphertext: Uint8Array): Uint8Array;
  readonly handshakeHash: Uint8Array;
  rekeyOutgoing(): void;
  rekeyIncoming(): void;
  close(): void;
}

export interface RustNoiseKeyPair {
  publicKey: Uint8Array;
  secretKey: Uint8Array;
}

let initialized = false;

function ensureInitialized(): void {
  if (initialized) return;
  const binary = atob(PEACH_NOISE_CORE_WASM_BASE64);
  const bytes = new Uint8Array(binary.length);
  for (let index = 0; index < binary.length; index += 1) {
    bytes[index] = binary.charCodeAt(index);
  }
  try {
    initSync({ module: bytes });
    initialized = true;
  } finally {
    bytes.fill(0);
  }
}

function session(native: WasmNoiseSession): RustNoiseSession {
  let closed = false;
  const requireOpen = () => {
    if (closed) throw new Error("Noise session is closed");
  };
  return {
    encrypt(plaintext) {
      requireOpen();
      return native.encrypt(plaintext);
    },
    decrypt(ciphertext) {
      requireOpen();
      return native.decrypt(ciphertext);
    },
    get handshakeHash() {
      requireOpen();
      return native.handshakeHash;
    },
    rekeyOutgoing() {
      requireOpen();
      native.rekeyOutgoing();
    },
    rekeyIncoming() {
      requireOpen();
      native.rekeyIncoming();
    },
    close() {
      if (closed) return;
      closed = true;
      native.free();
    },
  };
}

async function driveInitiator(
  native: WasmNoiseSession,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  try {
    await send(native.writeHandshake(new Uint8Array(0)));
    native.readHandshake(await receive());
    return session(native);
  } catch (error) {
    native.free();
    throw error;
  }
}

async function driveResponder(
  native: WasmNoiseSession,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  try {
    native.readHandshake(await receive());
    await send(native.writeHandshake(new Uint8Array(0)));
    return session(native);
  } catch (error) {
    native.free();
    throw error;
  }
}

export async function rustNoiseNkInitiator(
  responderStaticPublicKey: Uint8Array,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  ensureInitialized();
  const native = WasmNoiseSession.nkInitiator(responderStaticPublicKey);
  return driveInitiator(native, send, receive);
}

export async function rustNoiseNkResponder(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  ensureInitialized();
  const native = WasmNoiseSession.nkResponder(ourStaticPrivateKey, ourStaticPublicKey);
  return driveResponder(native, send, receive);
}

export async function rustNoiseKkInitiator(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  responderStaticPublicKey: Uint8Array,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  ensureInitialized();
  const native = WasmNoiseSession.kkInitiator(
    ourStaticPrivateKey,
    ourStaticPublicKey,
    responderStaticPublicKey,
  );
  return driveInitiator(native, send, receive);
}

export async function rustNoiseKkResponder(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  expectedInitiatorStaticPublicKey: Uint8Array,
  send: (message: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<RustNoiseSession> {
  ensureInitialized();
  const native = WasmNoiseSession.kkResponder(
    ourStaticPrivateKey,
    ourStaticPublicKey,
    expectedInitiatorStaticPublicKey,
  );
  return driveResponder(native, send, receive);
}