/**
 * E2E crypto for the sueta p2p stack. WebCrypto only. See docs/PROTOCOL.md §2-3.
 *
 * Key structure: every key is derived per-SENDER per-SESSION (peer ids are fresh
 * random per page load), so no two parties ever encrypt under the same key.
 * IVs are owned by a Sealer (epoch ‖ counter); there is no API that accepts a
 * caller-supplied IV, making nonce reuse structurally impossible.
 */

const td = new TextDecoder();

/** Bytes always backed by a plain ArrayBuffer (what WebCrypto's types demand). */
export type Bytes = Uint8Array<ArrayBuffer>;

const utf8 = (s: string): Bytes => new TextEncoder().encode(s) as Bytes;

export function randomBytes(n: number): Bytes {
  const b = new Uint8Array(n);
  crypto.getRandomValues(b);
  return b;
}

export function toB64(bytes: Bytes): string {
  let s = "";
  for (const b of bytes) s += String.fromCharCode(b);
  return btoa(s);
}

export function fromB64(s: string): Bytes {
  const bin = atob(s);
  const b = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) b[i] = bin.charCodeAt(i);
  return b;
}

export function toB64url(bytes: Bytes): string {
  return toB64(bytes).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/, "");
}

export function fromB64url(s: string): Bytes {
  return fromB64(s.replaceAll("-", "+").replaceAll("_", "/"));
}

export function newRoomSecret(): string {
  return toB64url(randomBytes(32));
}

export function newPeerId(): string {
  return Array.from(randomBytes(16), (b) => b.toString(16).padStart(2, "0")).join("");
}

export interface Envelope {
  v: 1;
  iv: string;
  ct: string;
}

/** Owns (key, epoch, counter). The only way to encrypt in this library. */
export class Sealer {
  #key: CryptoKey;
  #epoch: Bytes;
  #counter = 0n;

  constructor(key: CryptoKey) {
    this.#key = key;
    this.#epoch = randomBytes(4);
  }

  #nextIV(): Bytes {
    const iv = new Uint8Array(12);
    iv.set(this.#epoch, 0);
    new DataView(iv.buffer).setBigUint64(4, this.#counter++);
    return iv;
  }

  async seal(obj: unknown, aad: Bytes): Promise<Envelope> {
    const iv = this.#nextIV();
    const ct = await crypto.subtle.encrypt(
      { name: "AES-GCM", iv, additionalData: aad },
      this.#key,
      utf8(JSON.stringify(obj)),
    );
    return { v: 1, iv: toB64(iv), ct: toB64(new Uint8Array(ct)) };
  }

  /** Binary envelope: [0x01 | 12-byte IV | ciphertext+tag] as one ArrayBuffer. */
  async sealBinary(data: Bytes, aad: Bytes): Promise<ArrayBuffer> {
    const iv = this.#nextIV();
    const ct = new Uint8Array(
      await crypto.subtle.encrypt({ name: "AES-GCM", iv, additionalData: aad }, this.#key, data as BufferSource),
    );
    const out = new Uint8Array(1 + 12 + ct.length);
    out[0] = 0x01;
    out.set(iv, 1);
    out.set(ct, 13);
    return out.buffer;
  }
}

async function open(key: CryptoKey, env: Envelope, aad: Bytes): Promise<unknown | null> {
  try {
    const pt = await crypto.subtle.decrypt(
      { name: "AES-GCM", iv: fromB64(env.iv) as BufferSource, additionalData: aad },
      key,
      fromB64(env.ct) as BufferSource,
    );
    return JSON.parse(td.decode(pt));
  } catch {
    return null; // authentication failure ⇒ drop silently (protocol §3)
  }
}

async function openBinary(key: CryptoKey, buf: ArrayBuffer, aad: Bytes): Promise<Uint8Array | null> {
  const b = new Uint8Array(buf);
  if (b.length < 1 + 12 + 16 || b[0] !== 0x01) return null;
  try {
    const pt = await crypto.subtle.decrypt(
      { name: "AES-GCM", iv: b.subarray(1, 13) as BufferSource, additionalData: aad },
      key,
      b.subarray(13) as BufferSource,
    );
    return new Uint8Array(pt);
  } catch {
    return null;
  }
}

/**
 * All derivations for one room. `appSalt` namespaces the application
 * (e.g. "chat.ardegazu.ro/v1") so identical secrets on different apps built on
 * this stack land in different rooms.
 */
export class RoomCrypto {
  readonly roomId: string;
  readonly peerId: string;
  #ikm: CryptoKey;
  #salt: Bytes;
  #keyCache = new Map<string, Promise<CryptoKey>>();
  #sigSealer: Sealer;
  #msgSealer: Sealer;

  private constructor(ikm: CryptoKey, salt: Bytes, roomId: string, peerId: string, sig: Sealer, msg: Sealer) {
    this.#ikm = ikm;
    this.#salt = salt;
    this.roomId = roomId;
    this.peerId = peerId;
    this.#sigSealer = sig;
    this.#msgSealer = msg;
  }

  static async create(secretB64url: string, appSalt: string): Promise<RoomCrypto> {
    const secret = fromB64url(secretB64url);
    if (secret.length < 16) throw new Error("room secret too short");
    const salt = utf8(appSalt);
    const ikm = await crypto.subtle.importKey("raw", secret as BufferSource, "HKDF", false, [
      "deriveBits",
      "deriveKey",
    ]);
    const roomIdBits = await crypto.subtle.deriveBits(
      { name: "HKDF", hash: "SHA-256", salt: salt as BufferSource, info: utf8("roomid") as BufferSource },
      ikm,
      256,
    );
    const peerId = newPeerId();
    const kSig = await deriveKey(ikm, salt, `signal|${peerId}`);
    const kMsg = await deriveKey(ikm, salt, `msg|${peerId}`);
    return new RoomCrypto(ikm, salt, toB64url(new Uint8Array(roomIdBits)), peerId, new Sealer(kSig), new Sealer(kMsg));
  }

  #keyFor(info: string): Promise<CryptoKey> {
    let p = this.#keyCache.get(info);
    if (!p) {
      p = deriveKey(this.#ikm, this.#salt, info);
      this.#keyCache.set(info, p);
    }
    return p;
  }

  #sigAAD(from: string, to: string): Bytes {
    return utf8(`v1|${this.roomId}|${from}|${to}|sig`);
  }

  #msgAAD(from: string): Bytes {
    return utf8(`v1|${this.roomId}|${from}|msg`);
  }

  /** Encrypt a signaling payload for a specific recipient. */
  sealSignal(to: string, payload: unknown): Promise<Envelope> {
    return this.#sigSealer.seal(payload, this.#sigAAD(this.peerId, to));
  }

  /** Decrypt a signaling payload from `from` addressed to me. */
  async openSignal(from: string, env: Envelope): Promise<unknown | null> {
    return open(await this.#keyFor(`signal|${from}`), env, this.#sigAAD(from, this.peerId));
  }

  /** Encrypt a datachannel JSON payload (same ciphertext broadcasts to everyone). */
  sealMsg(payload: unknown): Promise<Envelope> {
    return this.#msgSealer.seal(payload, this.#msgAAD(this.peerId));
  }

  async openMsg(from: string, env: Envelope): Promise<unknown | null> {
    return open(await this.#keyFor(`msg|${from}`), env, this.#msgAAD(from));
  }

  sealMsgBinary(data: Bytes): Promise<ArrayBuffer> {
    return this.#msgSealer.sealBinary(data, this.#msgAAD(this.peerId));
  }

  async openMsgBinary(from: string, buf: ArrayBuffer): Promise<Uint8Array | null> {
    return openBinary(await this.#keyFor(`msg|${from}`), buf, this.#msgAAD(from));
  }
}

function deriveKey(ikm: CryptoKey, salt: Bytes, info: string): Promise<CryptoKey> {
  return crypto.subtle.deriveKey(
    { name: "HKDF", hash: "SHA-256", salt: salt as BufferSource, info: utf8(info) as BufferSource },
    ikm,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"],
  );
}