// Tests the hand-rolled WebSocket codec directly, using a raw TCP client so
// the low-level frames the browser API never exposes (ping/pong, fragmentation,
// binary frames, close) can be exercised.

import { connect } from "node:net";
import { randomBytes } from "node:crypto";
import { TestCase } from "./framework/test_case.js";
import { withServer } from "./framework/helpers.js";

const OP_CONTINUATION = 0x0;
const OP_TEXT = 0x1;
const OP_BINARY = 0x2;
const OP_CLOSE = 0x8;
const OP_PING = 0x9;

function encodeFrame(opcode, payload, fin = true) {
  const mask = randomBytes(4);
  const length = payload.length;
  let header;
  if (length < 126) {
    header = Buffer.alloc(2);
    header[1] = 0x80 | length;
  } else if (length < 65536) {
    header = Buffer.alloc(4);
    header[1] = 0x80 | 126;
    header.writeUInt16BE(length, 2);
  } else {
    header = Buffer.alloc(10);
    header[1] = 0x80 | 127;
    header.writeBigUInt64BE(BigInt(length), 2);
  }
  header[0] = (fin ? 0x80 : 0) | opcode;
  const masked = Buffer.alloc(length);
  for (let i = 0; i < length; i++) masked[i] = payload[i] ^ mask[i & 3];
  return Buffer.concat([header, mask, masked]);
}

function parseFrames(state) {
  while (true) {
    const buffer = state.buffer;
    if (buffer.length < 2) return;
    const b0 = buffer[0];
    const b1 = buffer[1];
    const fin = (b0 & 0x80) !== 0;
    const opcode = b0 & 0x0f;
    const masked = (b1 & 0x80) !== 0;
    let length = b1 & 0x7f;
    let offset = 2;
    if (length === 126) {
      if (buffer.length < 4) return;
      length = buffer.readUInt16BE(2);
      offset = 4;
    } else if (length === 127) {
      if (buffer.length < 10) return;
      length = Number(buffer.readBigUInt64BE(2));
      offset = 10;
    }
    let maskKey = null;
    if (masked) {
      if (buffer.length < offset + 4) return;
      maskKey = buffer.subarray(offset, offset + 4);
      offset += 4;
    }
    if (buffer.length < offset + length) return;
    let payload = buffer.subarray(offset, offset + length);
    if (masked) {
      const out = Buffer.alloc(length);
      for (let i = 0; i < length; i++) out[i] = payload[i] ^ maskKey[i & 3];
      payload = out;
    }
    state.buffer = buffer.subarray(offset + length);
    state.frames.push({ fin, opcode, payload });
  }
}

function rawClient(port, path = "/ws") {
  const socket = connect(port, "127.0.0.1");
  const state = { buffer: Buffer.alloc(0), frames: [] };
  const client = {
    socket,
    frames: state.frames,
    handshakeComplete: false,
    responseHeaders: "",
    send: (buffer) => socket.write(buffer),
    nextFrame: (timeout = 2000) =>
      new Promise((resolve, reject) => {
        const deadline = Date.now() + timeout;
        const tick = () => {
          parseFrames(state);
          if (state.frames.length > 0) {
            resolve(state.frames.shift());
            return;
          }
          if (Date.now() > deadline) {
            reject(new Error("timed out waiting for a frame"));
            return;
          }
          setTimeout(tick, 5);
        };
        tick();
      }),
    closeRaw: () => socket.destroy(),
  };

  client.ready = new Promise((resolve) => {
    const key = randomBytes(16).toString("base64");
    socket.on("connect", () => {
      socket.write(
        `GET ${path} HTTP/1.1\r\n` +
          `Host: 127.0.0.1:${port}\r\n` +
          "Upgrade: websocket\r\n" +
          "Connection: Upgrade\r\n" +
          `Sec-WebSocket-Key: ${key}\r\n` +
          "Sec-WebSocket-Version: 13\r\n\r\n"
      );
    });
    socket.on("data", (chunk) => {
      state.buffer = Buffer.concat([state.buffer, chunk]);
      if (!client.handshakeComplete) {
        const text = state.buffer.toString("latin1");
        const end = text.indexOf("\r\n\r\n");
        if (end < 0) return;
        client.responseHeaders = text.slice(0, end);
        state.buffer = state.buffer.subarray(end + 4);
        client.handshakeComplete = true;
        resolve();
      }
      parseFrames(state);
    });
  });
  client.closed = new Promise((resolve) => {
    socket.on("close", resolve);
    socket.on("end", resolve);
  });
  socket.on("error", () => {});
  return client;
}

export class WebSocketProtocolTest extends TestCase {
  async test_handshake_and_text_message() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      this.assertTrue(client.responseHeaders.startsWith("HTTP/1.1 101"));
      this.assertTrue(/sec-websocket-accept/i.test(client.responseHeaders));
      client.send(encodeFrame(OP_TEXT, Buffer.from(JSON.stringify({ t: "request_games" }))));
      const frame = await client.nextFrame();
      this.assertEqual(frame.opcode, OP_TEXT);
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_wrong_path_is_rejected() {
    await withServer(async (port) => {
      const client = rawClient(port, "/nope");
      await client.closed;
      this.assertFalse(client.handshakeComplete, "no upgrade on a foreign path");
      client.closeRaw();
    });
  }

  async test_ping_gets_pong() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      client.send(encodeFrame(OP_PING, Buffer.from("hi")));
      const frame = await client.nextFrame();
      this.assertEqual(frame.opcode, 0xa);
      this.assertEqual(frame.payload.toString("utf8"), "hi");
      client.closeRaw();
    });
  }

  async test_fragmented_text_is_reassembled() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      const payload = Buffer.from(JSON.stringify({ t: "request_games" }));
      const half = Math.floor(payload.length / 2);
      client.send(encodeFrame(OP_TEXT, payload.subarray(0, half), false));
      client.send(encodeFrame(OP_CONTINUATION, payload.subarray(half), true));
      const frame = await client.nextFrame();
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_binary_frame_is_delivered() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      client.send(encodeFrame(OP_BINARY, Buffer.from(JSON.stringify({ t: "request_games" }))));
      const frame = await client.nextFrame();
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_malformed_message_is_ignored() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      client.send(encodeFrame(OP_TEXT, Buffer.from("{not json")));
      client.send(encodeFrame(OP_TEXT, Buffer.from(JSON.stringify({ t: "request_games" }))));
      const frame = await client.nextFrame();
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_medium_frame_uses_the_16_bit_length() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      // 300-ish bytes exercises the 16-bit length branch; trailing whitespace
      // keeps the padding valid JSON, so the server simply ignores the type.
      const payload = Buffer.from(JSON.stringify({ t: "noop" }) + " ".repeat(300));
      client.send(encodeFrame(OP_TEXT, payload));
      client.send(encodeFrame(OP_TEXT, Buffer.from(JSON.stringify({ t: "request_games" }))));
      const frame = await client.nextFrame();
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_large_frame_is_parsed() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      // Larger than 64 KiB, exercising the 64-bit length branch. The body is
      // not valid JSON, so the server ignores it but must stay usable.
      client.send(encodeFrame(OP_TEXT, Buffer.alloc(70000, 0x78)));
      client.send(encodeFrame(OP_TEXT, Buffer.from(JSON.stringify({ t: "request_games" }))));
      const frame = await client.nextFrame();
      this.assertEqual(JSON.parse(frame.payload.toString("utf8")).t, "game_list");
      client.closeRaw();
    });
  }

  async test_close_frame_is_answered() {
    await withServer(async (port) => {
      const client = rawClient(port);
      await client.ready;
      client.send(encodeFrame(OP_CLOSE, Buffer.alloc(0)));
      const frame = await client.nextFrame();
      this.assertEqual(frame.opcode, OP_CLOSE);
      await client.closed;
      this.assertTrue(client.closed, "the server closed the socket");
      client.closeRaw();
    });
  }
}