// API-level adversarial stress tests for the quiz portal.
// Each scenario boots a fresh server on its own port, runs assertions,
// and logs JSON lines to stdout. Designed to be run repeatedly with
// different seeds; see run_loop.sh for the wrapper.

import { bootServer, setupSession, Student, Admin, STRESS_POOL, sleep, logLine, rand, pickRandom } from "./lib.mjs";
import WebSocket from "ws";

const SEED = parseInt(process.argv[2] || Date.now(), 10);
const PORT = parseInt(process.argv[3] || (8200 + (SEED % 100)), 10);

let mulberry32 = (s) => () => {
  s |= 0; s = s + 0x6D2B79F5 | 0;
  let t = Math.imul(s ^ s >>> 15, 1 | s);
  t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t;
  return ((t ^ t >>> 14) >>> 0) / 4294967296;
};
const rng = mulberry32(SEED);

let pass = 0, fail = 0, warn = 0;
const failures = [];

function expect(cond, scenario, msg, extra = {}) {
  if (cond) { pass++; logLine(scenario, "pass", msg, extra); }
  else      { fail++; failures.push({ scenario, msg, extra }); logLine(scenario, "fail", msg, extra); }
}
function note(scenario, msg, extra = {}) { warn++; logLine(scenario, "note", msg, extra); }

async function runScenario(name, fn) {
  logLine(name, "start", `seed=${SEED}`);
  let server;
  try {
    server = await bootServer({ port: PORT });
    await fn(server);
    logLine(name, "ok", "scenario completed");
  } catch (err) {
    fail++;
    failures.push({ scenario: name, msg: "uncaught", extra: { err: err.message, stack: err.stack?.slice(0, 600) } });
    logLine(name, "fail", "uncaught exception", { err: err.message });
  } finally {
    if (server) await server.stop();
  }
}

// ---------- Scenarios ----------

async function happyPath(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const N = 20;
  const students = await Promise.all(Array.from({ length: N }, async (_, i) => {
    const s = new Student(server.url, sid, `S${i.toString().padStart(3, "0")}`, `Student${i}`);
    await s.join();
    await s.connect();
    return s;
  }));
  const admin = new Admin(server.url, sid, jar);
  await admin.connect();
  await admin.waitFor("lobby_update");

  for (let q = 0; q < STRESS_POOL.questions.length; q++) {
    // Pre-register waiters so we don't lose the broadcast in the race window
    const studentOpenWaits = students.map(s => s.waitFor("question_open"));
    const adminOpenWait = admin.waitFor("question_open");
    admin.open(q, 5);
    await adminOpenWait;
    await Promise.all(studentOpenWaits);
    // Each student picks a random answer (mostly correct)
    await Promise.all(students.map(async (s, i) => {
      try {
        await sleep(rand(50, 800));
        const correct = STRESS_POOL.questions[q].correct;
        const ans = rng() < 0.7 ? correct : pickRandom(["A","B","C","D"]);
        const ackWait = s.waitFor("submit_ack", { timeoutMs: 3000 });
        s.submit(q, ans);
        const ack = await ackWait;
        expect(ack.question_idx === q, "happy", `student${i} q${q} ack idx==q`, { ack_idx: ack.question_idx, expected: q });
      } catch (e) {
        note("happy", `student${i} q${q}: ${e.message}`);
      }
    }));
    const studentClosedWaits = students.map(s => s.waitFor("question_closed", { timeoutMs: 3000 }).catch(() => null));
    const adminClosedWait = admin.waitFor("question_closed", { timeoutMs: 3000 });
    admin.close();
    await adminClosedWait;
    await Promise.all(studentClosedWaits);
    if (q < STRESS_POOL.questions.length - 1) {
      admin.next();
      await sleep(150);
    }
  }
  const sessionEndedWait = admin.waitFor("session_ended", { timeoutMs: 3000 });
  admin.end();
  await sessionEndedWait;
  expect(true, "happy", "session ended cleanly");
  students.forEach(s => s.disconnect());
  admin.disconnect();
}

async function lateJoiners(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar);
  await admin.connect();
  // 5 students join in lobby
  const early = await Promise.all([0,1,2,3,4].map(async i => {
    const s = new Student(server.url, sid, `E${i}`, `Early${i}`);
    await s.join(); await s.connect(); return s;
  }));
  const adminOpenWait = admin.waitFor("question_open");
  admin.open(0, 8);
  await adminOpenWait;
  await sleep(2000);
  // 3 late joiners
  for (let i = 0; i < 3; i++) {
    const s = new Student(server.url, sid, `L${i}`, `Late${i}`);
    await s.join();
    // Pre-register waiters BEFORE connect so we catch the snapshot on connect
    const stateWait = s.waitFor("state");
    const qopenWait = s.waitFor("question_open", { timeoutMs: 2000 });
    await s.connect();
    const m = await qopenWait.catch(() => null);
    if (!m) { fail++; failures.push({ scenario: "late_join", msg: `late${i} got no question_open on connect`, extra: {} }); logLine("late_join", "fail", `late${i} got no question_open on connect`); continue; }
    expect(m.question_idx === 0, "late_join", `late${i} sees correct idx`);
    expect(m.remaining_ms < 8000, "late_join", `late${i} remaining_ms reduced`, { remaining_ms: m.remaining_ms });
    expect(m.remaining_ms > 0, "late_join", `late${i} remaining_ms > 0`, { remaining_ms: m.remaining_ms });
    const ackWait = s.waitFor("submit_ack", { timeoutMs: 3000 });
    s.submit(0, STRESS_POOL.questions[0].correct);
    const ack = await ackWait.catch(() => null);
    expect(ack && ack.score > 0, "late_join", `late${i} got positive score from late submit`, { score: ack?.score });
    s.disconnect();
  }
  const adminClosedWait = admin.waitFor("question_closed");
  admin.close();
  await adminClosedWait;
  early.forEach(s => s.disconnect());
  admin.disconnect();
}

async function disconnectMidQuestion(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar);
  await admin.connect();
  const s = new Student(server.url, sid, "D1", "Drop");
  await s.join();
  const stateWaitInitial = s.waitFor("state");
  await s.connect();
  await stateWaitInitial;
  const sQopenWait = s.waitFor("question_open");
  const aOpenWait = admin.waitFor("question_open");
  admin.open(0, 10);
  await aOpenWait;
  await sQopenWait;
  s.disconnect();
  await sleep(500);
  // While dropped, instructor closes
  const closedWait = admin.waitFor("question_closed");
  admin.close();
  await closedWait;
  // Reconnect — should get state=question_closed (or current state)
  const reconnectStateWait = s.waitFor("state", { timeoutMs: 2000 });
  await s.reconnect();
  const state = await reconnectStateWait;
  expect(["question_closed", "between_questions", "lobby"].includes(state.state), "disconnect_midq", `state on reconnect = ${state.state}`);
  s.disconnect();
  admin.disconnect();
}

// THE KEY scenario from the user: phone screen sleeps mid-quiz, instructor
// advances to a new question, phone wakes — the student MUST see the
// LATEST question, not a stale screen.
async function sleepWakeNextQuestion(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar);
  await admin.connect();
  const s = new Student(server.url, sid, "SW1", "SleepWake");
  await s.join();
  const initState = s.waitFor("state");
  await s.connect();
  await initState;
  // Open Q0
  let sWait = s.waitFor("question_open");
  let aWait = admin.waitFor("question_open");
  admin.open(0, 5);
  await aWait; await sWait;
  await sleep(500);
  const ackWait = s.waitFor("submit_ack");
  s.submit(0, "B");
  await ackWait;
  // Phone "sleeps" — drop WS hard
  s.disconnect();
  await sleep(800);
  // Instructor closes, advances, opens Q2 (skip Q1 to make wake state non-trivial)
  let closedW = admin.waitFor("question_closed");
  admin.close();
  await closedW;
  admin.next();
  await sleep(150);
  aWait = admin.waitFor("question_open");
  admin.open(2, 8);
  await aWait;
  await sleep(300);
  // Phone wakes — reconnect; pre-register both expected events
  const stateOnWake = s.waitFor("state", { timeoutMs: 2000 });
  const qopenOnWake = s.waitFor("question_open", { timeoutMs: 2000 });
  await s.reconnect();
  const state = await stateOnWake;
  expect(state.state === "question_open", "sleep_wake", `state on wake = ${state.state}`, { state });
  expect(state.current_question_idx === 2, "sleep_wake", `current_question_idx on wake = ${state.current_question_idx}`);
  const qopen = await qopenOnWake.catch(() => null);
  expect(qopen && qopen.question_idx === 2, "sleep_wake", "reconnect emits question_open for current idx", { qopen_idx: qopen?.question_idx });
  expect(qopen && qopen.text === STRESS_POOL.questions[2].text, "sleep_wake", "question text matches latest Q");
  if (qopen) {
    const ack2W = s.waitFor("submit_ack", { timeoutMs: 3000 });
    s.submit(2, "B");
    const ack = await ack2W.catch(() => null);
    expect(ack && ack.question_idx === 2, "sleep_wake", "submit on Q2 after wake succeeded", { ack_idx: ack?.question_idx });
  }
  s.disconnect();
  admin.disconnect();
}

// Cookie tampering: try to flip student_id by mangling the signed cookie
async function cookieTampering(server) {
  const { sid } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const s = new Student(server.url, sid, "ORIG", "Original");
  await s.join();
  // Tamper: append a junk byte
  const original = s.jar.get("qz_student");
  s.jar.set("qz_student", original + "X");
  // Try to connect WS
  let connected = false, closeCode = null;
  try {
    await new Promise((res, rej) => {
      const wsUrl = server.url.replace(/^http/, "ws") + `/ws/student/${sid}`;
      const w = new WebSocket(wsUrl, { headers: { Cookie: s.jar.header() } });
      w.on("open", () => { connected = true; w.close(); res(); });
      w.on("close", (c) => { closeCode = c; res(); });
      w.on("error", () => res());
      setTimeout(res, 1500);
    });
  } catch {}
  expect(!connected || closeCode === 4001 || closeCode === 1006 || closeCode === 1008, "cookie_tamper", "tampered cookie rejected on WS", { connected, closeCode });
  // Reset cookie and retry
  s.jar.set("qz_student", original);
  await s.connect();
  await s.waitFor("state");
  expect(true, "cookie_tamper", "valid cookie still works after tamper attempt");
  s.disconnect();
}

// Cross-session cookie: cookie from session A should not work on session B.
async function crossSessionCookie(server) {
  const { sid: sidA, jar: jarA } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const { sid: sidB } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const s = new Student(server.url, sidA, "X1", "CrossUser");
  await s.join();
  // Try to use sidA's cookie to access sidB
  const wsUrl = server.url.replace(/^http/, "ws") + `/ws/student/${sidB}`;
  let opened = false;
  await new Promise(res => {
    const w = new WebSocket(wsUrl, { headers: { Cookie: s.jar.header() } });
    w.on("open", () => { opened = true; w.close(); res(); });
    w.on("close", () => res());
    w.on("error", () => res());
    setTimeout(res, 1500);
  });
  expect(!opened, "cross_session", "cookie from sidA rejected when used against sidB", { opened });
}

// Duplicate student_id: two browsers join with same student_id (different cookies)
async function duplicateStudentId(server) {
  const { sid } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const s1 = new Student(server.url, sid, "DUP", "FirstSession");
  const s2 = new Student(server.url, sid, "DUP", "SecondSession");
  await s1.join();
  const s1State = s1.waitFor("state");
  await s1.connect();
  await s1State;
  await s2.join();
  const s2State = s2.waitFor("state");
  await s2.connect();
  await s2State;
  expect(s1.ws.readyState === WebSocket.OPEN, "dup_id", "first DUP session open");
  expect(s2.ws.readyState === WebSocket.OPEN, "dup_id", "second DUP session open");
  s1.disconnect(); s2.disconnect();
}

// Submit out-of-order / wrong question_idx
async function badSubmits(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar); await admin.connect();
  const s = new Student(server.url, sid, "BAD", "Bad");
  await s.join();
  const initState = s.waitFor("state");
  await s.connect();
  await initState;
  // Submit before any question opened
  s.submit(0, "A");
  await sleep(300);
  expect(!s.lastMsgByType.submit_ack, "bad_submit", "submit before open ignored / no ack", { last: Object.keys(s.lastMsgByType) });
  let sWait = s.waitFor("question_open");
  let aWait = admin.waitFor("question_open");
  admin.open(0, 5);
  await aWait; await sWait;
  // Wrong idx submit (give it 600 ms to arrive if it does)
  const ackBefore = s.lastMsgByType.submit_ack;
  s.submit(99, "A");
  await sleep(600);
  const ackAfter = s.lastMsgByType.submit_ack;
  expect(ackAfter === ackBefore || (ackAfter && ackAfter.question_idx !== 99), "bad_submit", "wrong-idx submit not acked", { ackAfter });
  // Valid submit
  const okWait = s.waitFor("submit_ack", { timeoutMs: 2000 });
  s.submit(0, "B");
  const ok = await okWait;
  expect(ok.question_idx === 0, "bad_submit", "valid submit acked");
  // Resubmit (already submitted) — should NOT change stored answer
  s.submit(0, "A");
  await sleep(400);
  const closedWait = s.waitFor("question_closed");
  admin.close();
  const closed = await closedWait;
  expect(closed.your_answer === "B", "bad_submit", "your_answer remained 'B' after resubmit attempt", { your_answer: closed.your_answer });
  s.disconnect(); admin.disconnect();
}

// Race: many students submit at the moment of close (within ms)
async function closeBoundaryRace(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar); await admin.connect();
  const N = 30;
  const students = await Promise.all(Array.from({ length: N }, async (_, i) => {
    const s = new Student(server.url, sid, `R${i}`, `Race${i}`);
    await s.join(); await s.connect(); return s;
  }));
  // Pre-register everyone's question_open + question_closed waits
  const sOpens = students.map(s => s.waitFor("question_open", { timeoutMs: 3000 }).catch(() => null));
  const aOpen = admin.waitFor("question_open", { timeoutMs: 3000 });
  // Pre-register the closed wait too — auto-close fires even without manual close()
  const aClosed = admin.waitFor("question_closed", { timeoutMs: 6000 });
  admin.open(0, 1); // 1-second window — auto-close should fire ~1s later
  await aOpen;
  await Promise.all(sOpens);
  // Have all students fire submit at random times spanning the window edge
  const fires = students.map(async (s, i) => {
    const ackW = s.waitFor("submit_ack", { timeoutMs: 2000 });
    await sleep(rand(800, 1200)); // some fire after auto-close
    s.submit(0, "B");
    return ackW.catch(() => null);
  });
  const acks = await Promise.all(fires);
  const acked = acks.filter(Boolean).length;
  const closed = await aClosed.catch(() => null);
  logLine("close_race", "info", `race results`, { acked, total: N, hist: closed?.histogram });
  expect(closed !== null, "close_race", "question_closed broadcast received (auto-close or manual)");
  expect(acked >= 1 && acked <= N, "close_race", "no crash, some submits succeeded");
  students.forEach(s => s.disconnect()); admin.disconnect();
}

// Fuzz: malformed messages from a student WS
async function fuzzMessages(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar); await admin.connect();
  // First: a sentinel student that does NOT receive fuzz, to verify global server health
  const sentinel = new Student(server.url, sid, "SENT", "Sentinel");
  await sentinel.join();
  const sentinelStateW = sentinel.waitFor("state");
  await sentinel.connect();
  await sentinelStateW;
  // Second: a fuzz student that gets garbage shoved at it
  const s = new Student(server.url, sid, "FZ", "Fuzz");
  await s.join();
  const fzStateW = s.waitFor("state");
  await s.connect();
  await fzStateW;
  const garbages = [
    "not json",
    "{}",
    JSON.stringify({ type: "open_question" }), // student trying to act as instructor
    JSON.stringify({ type: "submit", question_idx: -1, answer: "Z" }),
    JSON.stringify({ type: "submit", question_idx: 0, answer: { nested: "obj" } }),
    JSON.stringify({ type: "💀" }),
    "x".repeat(50_000),
  ];
  for (const g of garbages) { try { s.ws.send(g); } catch {} ; await sleep(50); }
  await sleep(500);
  // Server should still serve OTHER clients regardless of what happened to fuzz student.
  const sentOpenW = sentinel.waitFor("question_open", { timeoutMs: 2000 });
  const adminOpenW = admin.waitFor("question_open", { timeoutMs: 2000 });
  admin.open(0, 3);
  await adminOpenW;
  const sm = await sentOpenW.catch(() => null);
  expect(sm && sm.question_idx === 0, "fuzz", "OTHER clients still served after fuzz on one student", { got: !!sm });
  // Did the fuzz student survive? (informational, not asserted as pass/fail)
  const survived = !s.closed && s.ws.readyState === WebSocket.OPEN;
  logLine("fuzz", "info", `fuzz student survival`, { survived, ws_state: s.ws.readyState });
  if (!survived) note("fuzz", "fuzz student WS was killed by malformed input — server lacks JSON-decode try/except in WS loop (codex room.py student_ws line ~87)");
  const adminClosedW = admin.waitFor("question_closed", { timeoutMs: 5000 });
  admin.close();
  await adminClosedW.catch(() => null);
  s.disconnect(); sentinel.disconnect(); admin.disconnect();
}

// Repeated rapid connect/disconnect (simulating flaky network)
async function flakyReconnect(server) {
  const { sid, jar } = await setupSession(server.url, server.adminPw, STRESS_POOL);
  const admin = new Admin(server.url, sid, jar); await admin.connect();
  const s = new Student(server.url, sid, "FK", "Flaky");
  await s.join();
  for (let i = 0; i < 10; i++) {
    const stateW = s.waitFor("state", { timeoutMs: 2000 });
    await s.connect();
    await stateW.catch(() => {});
    s.disconnect();
    await sleep(rand(50, 200));
  }
  // Final reconnect
  const finalStateW = s.waitFor("state", { timeoutMs: 2000 });
  await s.connect();
  await finalStateW;
  const sQopenW = s.waitFor("question_open");
  const aQopenW = admin.waitFor("question_open");
  admin.open(0, 5);
  await aQopenW; await sQopenW;
  const ackW = s.waitFor("submit_ack", { timeoutMs: 3000 });
  s.submit(0, "B");
  const ack = await ackW;
  expect(ack !== null, "flaky_reconnect", "post-flaky student can still submit");
  const closedW = admin.waitFor("question_closed");
  admin.close();
  await closedW;
  s.disconnect(); admin.disconnect();
}

const SCENARIOS = {
  happy: happyPath,
  late_join: lateJoiners,
  disconnect_midq: disconnectMidQuestion,
  sleep_wake: sleepWakeNextQuestion,
  cookie_tamper: cookieTampering,
  cross_session: crossSessionCookie,
  dup_id: duplicateStudentId,
  bad_submit: badSubmits,
  close_race: closeBoundaryRace,
  fuzz: fuzzMessages,
  flaky_reconnect: flakyReconnect,
};

// Pick scenario subset based on env or run all in random order
const wanted = process.env.SCENARIOS
  ? process.env.SCENARIOS.split(",")
  : Object.keys(SCENARIOS).sort(() => rng() - 0.5);

logLine("runner", "info", `starting api stress`, { seed: SEED, port: PORT, scenarios: wanted });

for (const name of wanted) {
  const fn = SCENARIOS[name];
  if (!fn) { logLine("runner", "warn", `unknown scenario ${name}`); continue; }
  await runScenario(name, fn);
}

logLine("runner", "summary", `done`, { pass, fail, warn, failures });
process.exit(fail > 0 ? 1 : 0);