#!/usr/bin/env node

import { webcrypto } from "node:crypto";
import assert from "node:assert/strict";

const wsUrl = process.argv[2] ?? process.env.WS_URL ?? "ws://localhost:8080/ws";
const apiBaseUrl = process.argv[3] ?? process.env.API_BASE_URL;
const protocolVersion = 1;

const encoder = new TextEncoder();
const decoder = new TextDecoder();

function bytesToBase64(bytes) {
  return Buffer.from(bytes).toString("base64");
}

function base64ToBytes(value) {
  return Buffer.from(value, "base64");
}

async function createSession() {
  const key = await webcrypto.subtle.generateKey(
    { name: "AES-GCM", length: 256 },
    true,
    ["encrypt", "decrypt"],
  );
  const rawKey = await webcrypto.subtle.exportKey("raw", key);

  return {
    key,
    keyId: webcrypto.randomUUID(),
    exportedKey: bytesToBase64(new Uint8Array(rawKey)),
  };
}

async function encryptMessage(session, message) {
  const nonce = webcrypto.getRandomValues(new Uint8Array(12));
  const plaintext = encoder.encode(JSON.stringify(message));
  const ciphertext = await webcrypto.subtle.encrypt(
    { name: "AES-GCM", iv: nonce },
    session.key,
    plaintext,
  );

  return {
    enc: `aes-256-gcm-v${protocolVersion}`,
    key_id: session.keyId,
    nonce: bytesToBase64(nonce),
    ciphertext: bytesToBase64(new Uint8Array(ciphertext)),
  };
}

async function decryptMessage(session, envelope) {
  const plaintext = await webcrypto.subtle.decrypt(
    { name: "AES-GCM", iv: base64ToBytes(envelope.nonce) },
    session.key,
    base64ToBytes(envelope.ciphertext),
  );

  return JSON.parse(decoder.decode(plaintext));
}

function waitForOpen(socket) {
  return new Promise((resolve, reject) => {
    socket.addEventListener("open", resolve, { once: true });
    socket.addEventListener("error", reject, { once: true });
  });
}

function waitForRawMessage(socket, timeoutMs = 5000) {
  return new Promise((resolve, reject) => {
    const timeout = setTimeout(() => {
      socket.removeEventListener("message", onMessage);
      reject(new Error(`timeout after ${timeoutMs}ms`));
    }, timeoutMs);

    function onMessage(event) {
      clearTimeout(timeout);
      resolve(String(event.data));
    }

    socket.addEventListener("message", onMessage, { once: true });
  });
}

async function waitForDecryptedType(client, expectedType, timeoutMs = 5000) {
  const deadline = Date.now() + timeoutMs;

  while (Date.now() < deadline) {
    const raw = await waitForRawMessage(client.socket, Math.max(100, deadline - Date.now()));
    const wire = JSON.parse(raw);

    assert.equal(wire.type, "encrypted", `${client.name}: expected encrypted wire message`);
    assert.equal(wire.payload.key_id, client.session.keyId, `${client.name}: key id mismatch`);
    assertNoPlaintext(raw, client.name);

    const message = await decryptMessage(client.session, wire.payload);
    if (message.type === expectedType) {
      return message;
    }
  }

  throw new Error(`${client.name}: did not receive ${expectedType}`);
}

async function waitForRecordChanged(client, expectedTitle, timeoutMs = 5000) {
  const deadline = Date.now() + timeoutMs;

  while (Date.now() < deadline) {
    const message = await waitForDecryptedType(
      client,
      "record_changed",
      Math.max(100, deadline - Date.now()),
    );
    if (!expectedTitle || message.payload.record?.title === expectedTitle) {
      return message;
    }
  }

  throw new Error(`${client.name}: did not receive record_changed ${expectedTitle}`);
}

function assertNoPlaintext(raw, clientName) {
  for (const forbidden of ["snapshot", "record_changed", "Erster Datensatz", "records"]) {
    assert.equal(
      raw.includes(forbidden),
      false,
      `${clientName}: raw frame leaked plaintext marker ${forbidden}`,
    );
  }
}

async function connectClient(name) {
  const session = await createSession();
  const socket = new WebSocket(wsUrl);

  await waitForOpen(socket);
  socket.send(JSON.stringify({
    type: "hello",
    payload: {
      protocol_version: protocolVersion,
      key_id: session.keyId,
      session_key: session.exportedKey,
    },
  }));

  const ack = JSON.parse(await waitForRawMessage(socket));
  assert.equal(ack.type, "hello_ack", `${name}: expected hello_ack`);
  assert.equal(ack.payload.protocol_version, protocolVersion, `${name}: protocol mismatch`);
  assert.equal(ack.payload.key_id, session.keyId, `${name}: ack key mismatch`);

  const client = { name, socket, session };
  const snapshot = await waitForDecryptedType(client, "snapshot");
  assert.ok(Array.isArray(snapshot.payload.records), `${name}: snapshot records missing`);

  await sendEncrypted(client, {
    type: "subscribe",
    payload: { topic: "records" },
  });

  return client;
}

async function sendEncrypted(client, message) {
  const envelope = await encryptMessage(client.session, message);
  client.socket.send(JSON.stringify({ type: "encrypted", payload: envelope }));
}

function closeClient(client) {
  if (client.socket.readyState === WebSocket.OPEN) {
    client.socket.close();
  }
}

async function request(method, path, body, token) {
  const response = await fetch(`${apiBaseUrl}${path}`, {
    method,
    headers: {
      "Content-Type": "application/json",
      ...(token ? { Authorization: `Bearer ${token}` } : {}),
    },
    body: body === undefined ? undefined : JSON.stringify(body),
  });
  const text = await response.text();
  const data = text ? JSON.parse(text) : {};

  if (!response.ok) {
    throw new Error(`${method} ${path} failed: ${response.status} ${JSON.stringify(data)}`);
  }

  return data;
}

async function createLiveEventViaApi() {
  const email = `live-event-${Date.now()}@example.test`;
  const bootstrap = await request("POST", "/api/v1/dev/bootstrap-local", {
    organization_name: "Live Event Test GmbH",
    email,
  });
  const login = await request("POST", "/api/v1/auth/login", {
    email,
    password: bootstrap.password,
  });
  const token = login.access_token;
  await request("POST", "/api/v1/auth/select-organization", {
    organization_id: login.organization_id,
  }, token);

  await request("POST", "/api/v1/activities", {
    activity_number: null,
    activity_type: "task",
    title: "Live-Event testen",
    body: "Änderung muss an alle Clients gehen.",
    status: "open",
    priority: "normal",
    due_at: null,
  }, token);
}

async function main() {
  console.log(`testing encrypted communication via ${wsUrl}`);

  const clientA = await connectClient("client-a");
  console.log("client-a handshake, encrypted snapshot and subscribe ok");

  const clientB = await connectClient("client-b");
  console.log("client-b handshake, encrypted snapshot and subscribe ok");

  await sendEncrypted(clientA, { type: "ping" });
  const [pongA, pongB] = await Promise.all([
    waitForDecryptedType(clientA, "pong"),
    waitForDecryptedType(clientB, "pong"),
  ]);

  assert.equal(pongA.type, "pong");
  assert.equal(pongB.type, "pong");
  assert.notEqual(clientA.session.keyId, clientB.session.keyId, "clients must use different keys");

  if (apiBaseUrl) {
    console.log(`testing api-triggered live event via ${apiBaseUrl}`);
    const waitA = waitForRecordChanged(clientA, "Aktivität angelegt");
    const waitB = waitForRecordChanged(clientB, "Aktivität angelegt");
    await createLiveEventViaApi();
    await Promise.all([waitA, waitB]);
    console.log("api-triggered live event reached both clients");
  }

  closeClient(clientA);
  closeClient(clientB);

  console.log("encrypted multi-client communication test ok");
  process.exit(0);
}

main().catch((error) => {
  console.error(error);
  process.exit(1);
});