'use strict';

const path    = require('path');
const fs      = require('fs');
const http    = require('http');
const crypto  = require('crypto');
const express = require('express');
const session = require('express-session');
const bcrypt  = require('bcryptjs');
const helmet  = require('helmet');
const rateLimit = require('express-rate-limit');
const WebSocket = require('ws');
const Database  = require('better-sqlite3');

const PORT     = process.env.PORT || 3012;
const DATA_DIR = path.join(__dirname, 'data');

if (!fs.existsSync(DATA_DIR)) fs.mkdirSync(DATA_DIR, { recursive: true });

// ── Stratagems (mirrored in public/stratagems.js) ─────────────────────────────
const STRATAGEMS = [
  // Patriotic Administration Center
  { name: 'Reinforce',          category: 'Patriotic Administration Center', sequence: ['up','down','right','left','up'] },
  { name: 'Resupply',           category: 'Patriotic Administration Center', sequence: ['down','down','up','right'] },
  { name: 'SOS Beacon',         category: 'Patriotic Administration Center', sequence: ['up','down','right','up'] },
  { name: 'Hellbomb',           category: 'Patriotic Administration Center', sequence: ['down','up','left','down','up','right','down','up'] },
  { name: 'SEAF Artillery',     category: 'Patriotic Administration Center', sequence: ['right','up','up','down'] },
  { name: 'Upload Data',        category: 'Patriotic Administration Center', sequence: ['right','right','left','up','up'] },
  { name: 'Eagle Rearm',        category: 'Patriotic Administration Center', sequence: ['up','up','left','up','right'] },
  { name: 'Prospecting Drill',  category: 'Patriotic Administration Center', sequence: ['down','down','left','right','down'] },
  // Orbital Cannons
  { name: 'Orbital Gatling Barrage',    category: 'Orbital Cannons', sequence: ['right','down','left','up','up'] },
  { name: 'Orbital Airburst Strike',    category: 'Orbital Cannons', sequence: ['right','right','right'] },
  { name: 'Orbital 120MM HE Barrage',   category: 'Orbital Cannons', sequence: ['right','right','down','left','right','down'] },
  { name: 'Orbital 380MM HE Barrage',   category: 'Orbital Cannons', sequence: ['right','down','up','up','left','down','down'] },
  { name: 'Orbital Walking Barrage',    category: 'Orbital Cannons', sequence: ['right','down','right','down','right','down'] },
  { name: 'Orbital Laser',              category: 'Orbital Cannons', sequence: ['right','down','up','right','down'] },
  { name: 'Orbital Railcannon Strike',  category: 'Orbital Cannons', sequence: ['right','up','down','down','right'] },
  { name: 'Orbital Precision Strike',   category: 'Orbital Cannons', sequence: ['right','right','up'] },
  { name: 'Orbital Gas Strike',         category: 'Orbital Cannons', sequence: ['right','right','down','right'] },
  { name: 'Orbital EMS Strike',         category: 'Orbital Cannons', sequence: ['right','right','left','down'] },
  { name: 'Orbital Smoke Strike',       category: 'Orbital Cannons', sequence: ['right','right','down','up'] },
  { name: 'Orbital Illumination Flare', category: 'Orbital Cannons', sequence: ['right','right','left','left'] },
  // Hangar
  { name: 'Eagle Strafing Run',      category: 'Hangar', sequence: ['up','right','right'] },
  { name: 'Eagle Airstrike',         category: 'Hangar', sequence: ['up','right','down','right'] },
  { name: 'Eagle Cluster Bomb',      category: 'Hangar', sequence: ['up','right','down','down','right'] },
  { name: 'Eagle Napalm Airstrike',  category: 'Hangar', sequence: ['up','right','down','up'] },
  { name: 'LIFT-850 Jump Pack',      category: 'Hangar', sequence: ['down','up','up','down','up'] },
  { name: 'Eagle Smoke Strike',      category: 'Hangar', sequence: ['up','right','up','down'] },
  { name: 'Eagle 110MM Rocket Pods', category: 'Hangar', sequence: ['up','right','up','left'] },
  { name: 'Eagle 500KG Bomb',        category: 'Hangar', sequence: ['up','right','down','down','down'] },
  // Bridge
  { name: 'Patriot Exosuit',     category: 'Bridge', sequence: ['left','down','right','up','left','down','right'] },
  { name: 'Emancipator Exosuit', category: 'Bridge', sequence: ['left','down','right','up','left','down','down'] },
  // Engineering Bay – Support Weapons
  { name: 'Machine Gun',               category: 'Engineering Bay', sequence: ['down','left','down','up','right'] },
  { name: 'Anti-Materiel Rifle',       category: 'Engineering Bay', sequence: ['down','left','right','up','down'] },
  { name: 'Stalwart',                  category: 'Engineering Bay', sequence: ['down','left','down','up','up','left'] },
  { name: 'Expendable Anti-Tank',      category: 'Engineering Bay', sequence: ['down','down','left','up'] },
  { name: 'Recoilless Rifle',          category: 'Engineering Bay', sequence: ['down','left','right','right','left'] },
  { name: 'Flamethrower',              category: 'Engineering Bay', sequence: ['down','left','up','down','up'] },
  { name: 'Autocannon',                category: 'Engineering Bay', sequence: ['down','left','down','up','up','right'] },
  { name: 'Heavy Machine Gun',         category: 'Engineering Bay', sequence: ['down','left','up','down','down'] },
  { name: 'Airburst Rocket Launcher',  category: 'Engineering Bay', sequence: ['down','up','up','left','right'] },
  { name: 'Commando',                  category: 'Engineering Bay', sequence: ['down','left','up','down','right'] },
  { name: 'Railgun',                   category: 'Engineering Bay', sequence: ['down','right','down','up','left','right'] },
  { name: 'Spear',                     category: 'Engineering Bay', sequence: ['down','down','up','down','down'] },
  { name: 'Quasar Cannon',             category: 'Engineering Bay', sequence: ['down','down','up','left','right'] },
  { name: 'Arc Thrower',               category: 'Engineering Bay', sequence: ['down','right','down','up','left','left'] },
  { name: 'Laser Cannon',              category: 'Engineering Bay', sequence: ['down','left','down','up','left'] },
  { name: 'Grenade Launcher',          category: 'Engineering Bay', sequence: ['down','left','up','left','down'] },
  // Engineering Bay – Equipment
  { name: 'Supply Pack',               category: 'Engineering Bay', sequence: ['down','left','down','up','up','down'] },
  { name: 'Guard Dog Rover',           category: 'Engineering Bay', sequence: ['down','up','left','up','right','right'] },
  { name: 'Guard Dog',                 category: 'Engineering Bay', sequence: ['down','up','left','up','right','down'] },
  { name: 'Ballistic Shield Backpack', category: 'Engineering Bay', sequence: ['down','left','down','down','up','left'] },
  { name: 'Shield Generator Pack',     category: 'Engineering Bay', sequence: ['down','up','left','right','left','right'] },
  { name: 'Directional Shield',        category: 'Engineering Bay', sequence: ['down','left','up','up','right'] },
  // Engineering Bay – Mines
  { name: 'Anti-Personnel Minefield',  category: 'Engineering Bay', sequence: ['down','left','up','right'] },
  { name: 'Incendiary Mines',          category: 'Engineering Bay', sequence: ['down','left','left','down'] },
  { name: 'Anti-Tank Mines',           category: 'Engineering Bay', sequence: ['down','down','left','left'] },
  // Robotics Workshop
  { name: 'Machine Gun Sentry',  category: 'Robotics Workshop', sequence: ['down','up','right','right','up'] },
  { name: 'Gatling Sentry',      category: 'Robotics Workshop', sequence: ['down','up','right','left'] },
  { name: 'Mortar Sentry',       category: 'Robotics Workshop', sequence: ['down','up','right','right','down'] },
  { name: 'Autocannon Sentry',   category: 'Robotics Workshop', sequence: ['down','up','right','up','left','up'] },
  { name: 'Rocket Sentry',       category: 'Robotics Workshop', sequence: ['down','up','right','right','left'] },
  { name: 'EMS Mortar Sentry',   category: 'Robotics Workshop', sequence: ['down','up','right','down','right'] },
  { name: 'Tesla Tower',         category: 'Robotics Workshop', sequence: ['down','up','right','up','left','up','up'] },
  // Defensive
  { name: 'Shield Generator Relay',   category: 'Defensive', sequence: ['down','up','left','right','left','down'] },
  { name: 'Anti-Tank Emplacement',    category: 'Defensive', sequence: ['down','right','right','up','left'] },
  { name: 'Orbital Shield Generator', category: 'Defensive', sequence: ['right','right','left','down','left','down'] },
];

const VALID_NAMES = new Set(STRATAGEMS.map(s => s.name));

// ── SQLite ────────────────────────────────────────────────────────────────────
const db = new Database(path.join(DATA_DIR, 'helldivers.db'));
db.pragma('journal_mode = WAL');
db.pragma('foreign_keys = ON');

function initDB() {
  db.exec(`
    CREATE TABLE IF NOT EXISTS users (
      username   TEXT PRIMARY KEY,
      hash       TEXT NOT NULL,
      role       TEXT NOT NULL DEFAULT 'user',
      mustChange INTEGER NOT NULL DEFAULT 1
    );
    CREATE TABLE IF NOT EXISTS practice_sessions (
      id         INTEGER PRIMARY KEY AUTOINCREMENT,
      username   TEXT NOT NULL,
      stratagem  TEXT NOT NULL,
      category   TEXT NOT NULL,
      time_ms    INTEGER NOT NULL,
      score      INTEGER NOT NULL,
      created_at TEXT NOT NULL
    );
    CREATE TABLE IF NOT EXISTS matches (
      id            INTEGER PRIMARY KEY AUTOINCREMENT,
      winner        TEXT NOT NULL,
      loser         TEXT NOT NULL,
      winner_rounds INTEGER NOT NULL,
      loser_rounds  INTEGER NOT NULL,
      created_at    TEXT NOT NULL
    );
    CREATE INDEX IF NOT EXISTS idx_ps_user ON practice_sessions(username);
    CREATE INDEX IF NOT EXISTS idx_m_winner ON matches(winner);
    CREATE INDEX IF NOT EXISTS idx_m_loser  ON matches(loser);
  `);
}

async function initUsers() {
  const defaults = [
    { username: 'admin',  role: 'admin' },
    { username: 'jeremy', role: 'user'  },
  ];
  for (const { username, role } of defaults) {
    const exists = db.prepare('SELECT username FROM users WHERE username = ?').get(username);
    if (!exists) {
      const tempPw = crypto.randomBytes(6).toString('hex');
      const hash   = await bcrypt.hash(tempPw, 12);
      db.prepare('INSERT INTO users (username, hash, role, mustChange) VALUES (?, ?, ?, 1)')
        .run(username, hash, role);
      console.log(`[INIT] Created user '${username}' – temp password: ${tempPw}`);
    }
  }
}

// ── Session secret ────────────────────────────────────────────────────────────
function getSessionSecret() {
  const file = path.join(DATA_DIR, '.session-secret');
  if (fs.existsSync(file)) return fs.readFileSync(file, 'utf8').trim();
  const secret = crypto.randomBytes(32).toString('hex');
  fs.writeFileSync(file, secret, { mode: 0o600 });
  return secret;
}

// ── WebSocket shared state (module-level so routes can read userSockets) ─────
const userSockets      = new Map(); // userId → ws
const pendingChallenges = new Map(); // challengerId → targetId
const rooms            = new Map(); // roomId → roomState

// ── Express ───────────────────────────────────────────────────────────────────
const app = express();
app.set('trust proxy', 1);

app.use(helmet({
  contentSecurityPolicy: {
    directives: {
      defaultSrc: ["'self'"],
      scriptSrc:  ["'self'"],
      styleSrc:   ["'self'", "'unsafe-inline'", 'https://fonts.googleapis.com'],
      fontSrc:    ["'self'", 'https://fonts.gstatic.com'],
      imgSrc:     ["'self'", 'data:'],
      connectSrc: ["'self'", 'ws:', 'wss:'],
    },
  },
}));

const generalLimiter = rateLimit({ windowMs: 60_000,       max: 300, standardHeaders: true, legacyHeaders: false });
const loginLimiter   = rateLimit({ windowMs: 15 * 60_000,  max: 10,  message: { error: 'Too many login attempts, try again later' } });

app.use(generalLimiter);
app.use(express.json({ limit: '10kb' }));

const sessionMiddleware = session({
  secret: getSessionSecret(),
  resave: false,
  saveUninitialized: false,
  cookie: {
    secure:   process.env.NODE_ENV === 'production',
    httpOnly: true,
    sameSite: 'strict',
    maxAge:   7 * 24 * 60 * 60 * 1000,
  },
});
app.use(sessionMiddleware);

// ── Middleware ────────────────────────────────────────────────────────────────
function requireAuth(req, res, next) {
  if (!req.session.user) return res.status(401).json({ error: 'Not logged in' });
  if (req.session.mustChange) {
    const allowed = ['/api/change-password', '/api/logout', '/api/me'];
    if (!allowed.includes(req.path)) {
      return res.status(403).json({ error: 'Password change required', mustChange: true });
    }
  }
  next();
}

function requireAdmin(req, res, next) {
  if (req.session.role !== 'admin') return res.status(403).json({ error: 'Admin only' });
  next();
}

// ── Auth ──────────────────────────────────────────────────────────────────────
app.post('/api/login', loginLimiter, async (req, res) => {
  const { username, password } = req.body || {};
  if (!username || !password) return res.status(400).json({ error: 'Username and password required' });

  const user = db.prepare('SELECT * FROM users WHERE username = ?').get(username);
  if (!user) return res.status(401).json({ error: 'Invalid credentials' });

  const valid = await bcrypt.compare(password, user.hash);
  if (!valid) return res.status(401).json({ error: 'Invalid credentials' });

  req.session.regenerate((err) => {
    if (err) return res.status(500).json({ error: 'Session error' });
    req.session.user      = user.username;
    req.session.role      = user.role;
    req.session.mustChange = user.mustChange === 1;
    res.json({ ok: true, user: user.username, role: user.role, mustChange: user.mustChange === 1 });
  });
});

app.post('/api/logout', (req, res) => {
  req.session.destroy(() => res.json({ ok: true }));
});

app.get('/api/me', (req, res) => {
  if (!req.session.user) return res.json({ user: null });
  res.json({ user: req.session.user, role: req.session.role, mustChange: !!req.session.mustChange });
});

app.post('/api/change-password', requireAuth, async (req, res) => {
  const { oldPassword, newPassword } = req.body || {};
  if (!oldPassword || !newPassword) return res.status(400).json({ error: 'Both passwords required' });
  if (newPassword.length < 8) return res.status(400).json({ error: 'Password must be at least 8 characters' });

  const user  = db.prepare('SELECT * FROM users WHERE username = ?').get(req.session.user);
  const valid = await bcrypt.compare(oldPassword, user.hash);
  if (!valid) return res.status(401).json({ error: 'Current password incorrect' });

  const hash = await bcrypt.hash(newPassword, 12);
  db.prepare('UPDATE users SET hash = ?, mustChange = 0 WHERE username = ?').run(hash, req.session.user);
  req.session.mustChange = false;
  res.json({ ok: true });
});

// ── User management (admin) ───────────────────────────────────────────────────
app.get('/api/users', requireAuth, requireAdmin, (req, res) => {
  const users = db.prepare('SELECT username, role, mustChange FROM users ORDER BY username').all();
  res.json(users.map(u => ({ ...u, mustChange: u.mustChange === 1 })));
});

app.post('/api/users', requireAuth, requireAdmin, async (req, res) => {
  const { username, role } = req.body || {};
  if (!username || !['admin', 'user'].includes(role))
    return res.status(400).json({ error: 'Valid username and role required' });
  if (!/^[a-zA-Z0-9_-]{2,32}$/.test(username))
    return res.status(400).json({ error: 'Username: 2-32 alphanumeric chars, _ or -' });

  const exists = db.prepare('SELECT username FROM users WHERE username = ?').get(username);
  if (exists) return res.status(409).json({ error: 'User already exists' });

  const tempPw = crypto.randomBytes(6).toString('hex');
  const hash   = await bcrypt.hash(tempPw, 12);
  db.prepare('INSERT INTO users (username, hash, role, mustChange) VALUES (?, ?, ?, 1)').run(username, hash, role);
  res.json({ ok: true, tempPassword: tempPw });
});

app.delete('/api/users/:username', requireAuth, requireAdmin, (req, res) => {
  const { username } = req.params;
  if (username === req.session.user) return res.status(400).json({ error: 'Cannot delete yourself' });

  const user = db.prepare('SELECT * FROM users WHERE username = ?').get(username);
  if (!user) return res.status(404).json({ error: 'User not found' });

  if (user.role === 'admin') {
    const adminCount = db.prepare("SELECT COUNT(*) AS c FROM users WHERE role = 'admin'").get().c;
    if (adminCount <= 1) return res.status(400).json({ error: 'Cannot delete the last admin' });
  }

  db.prepare('DELETE FROM users WHERE username = ?').run(username);
  res.json({ ok: true });
});

// ── Leaderboard cache ─────────────────────────────────────────────────────────
let lbCache = null;
let lbCacheTime = 0;
const LB_TTL = 30_000;

function invalidateLB() { lbCache = null; }

function getLeaderboard() {
  if (lbCache && Date.now() - lbCacheTime < LB_TTL) return lbCache;
  lbCache = db.prepare(`
    SELECT
      u.username,
      COALESCE(ps.sessions, 0)    AS sessions,
      COALESCE(ps.totalScore, 0)  AS totalScore,
      COALESCE(ps.fastestTime, 0) AS fastestTime,
      COALESCE(mw.wins, 0)        AS wins,
      COALESCE(mw.wins, 0) + COALESCE(ml.losses, 0) AS matches
    FROM users u
    LEFT JOIN (
      SELECT username, COUNT(*) AS sessions, SUM(score) AS totalScore, MIN(time_ms) AS fastestTime
      FROM practice_sessions GROUP BY username
    ) ps ON ps.username = u.username
    LEFT JOIN (SELECT winner AS username, COUNT(*) AS wins   FROM matches GROUP BY winner) mw ON mw.username = u.username
    LEFT JOIN (SELECT loser  AS username, COUNT(*) AS losses FROM matches GROUP BY loser)  ml ON ml.username = u.username
    WHERE COALESCE(ps.sessions,0) > 0 OR COALESCE(mw.wins,0) > 0
    ORDER BY totalScore DESC, fastestTime ASC
    LIMIT 20
  `).all();
  lbCacheTime = Date.now();
  return lbCache;
}

// ── Dashboard ─────────────────────────────────────────────────────────────────
app.get('/api/dashboard', requireAuth, (req, res) => {
  const u = req.session.user;

  const stats = db.prepare(`
    SELECT COUNT(*) AS sessions,
           COALESCE(SUM(score),0)    AS totalScore,
           COALESCE(MAX(score),0)    AS bestScore,
           COALESCE(MIN(time_ms),0)  AS fastestTime
    FROM practice_sessions WHERE username = ?
  `).get(u);

  const matchStats = db.prepare(`
    SELECT COUNT(*) AS matches,
           SUM(CASE WHEN winner = ? THEN 1 ELSE 0 END) AS wins
    FROM matches WHERE winner = ? OR loser = ?
  `).get(u, u, u);

  const lb    = getLeaderboard();
  const rankIdx = lb.findIndex(r => r.username === u);
  const rank  = rankIdx >= 0 ? { position: rankIdx + 1 } : null;

  const recent = db.prepare(`
    SELECT stratagem, category, score, time_ms, created_at
    FROM practice_sessions WHERE username = ?
    ORDER BY created_at DESC LIMIT 5
  `).all(u);

  const dayOfYear   = Math.floor((new Date() - new Date(new Date().getFullYear(), 0, 0)) / 86_400_000);
  const dailyStrat  = STRATAGEMS[dayOfYear % STRATAGEMS.length];
  const dailyBest   = db.prepare(`
    SELECT MIN(time_ms) AS bestTime FROM practice_sessions
    WHERE stratagem = ? AND username = ?
  `).get(dailyStrat.name, u);

  res.json({
    stats:    { ...stats, ...matchStats },
    rank,
    online:   [...userSockets.keys()],
    recent,
    daily:    { stratagem: dailyStrat, bestTime: dailyBest?.bestTime ?? null },
  });
});

// ── Scores ────────────────────────────────────────────────────────────────────
app.post('/api/scores/practice', requireAuth, (req, res) => {
  const { stratagem, category, time_ms, score } = req.body || {};
  if (!VALID_NAMES.has(stratagem))                           return res.status(400).json({ error: 'Invalid stratagem' });
  if (typeof time_ms !== 'number' || time_ms <= 0 || time_ms > 35_000) return res.status(400).json({ error: 'Invalid time' });
  if (typeof score   !== 'number' || score   <  0 || score   > 15_000) return res.status(400).json({ error: 'Invalid score' });

  db.prepare(`
    INSERT INTO practice_sessions (username, stratagem, category, time_ms, score, created_at)
    VALUES (?, ?, ?, ?, ?, ?)
  `).run(req.session.user, stratagem, category || '', time_ms, score, new Date().toISOString());

  invalidateLB();
  res.json({ ok: true });
});

app.get('/api/scores/leaderboard', requireAuth, (req, res) => {
  res.json(getLeaderboard());
});

app.get('/api/scores/me', requireAuth, (req, res) => {
  const u = req.session.user;
  const practice = db.prepare(`
    SELECT stratagem, category, score, time_ms, created_at
    FROM practice_sessions WHERE username = ?
    ORDER BY created_at DESC LIMIT 50
  `).all(u);
  const matches = db.prepare(`
    SELECT * FROM matches WHERE winner = ? OR loser = ?
    ORDER BY created_at DESC LIMIT 20
  `).all(u, u);
  res.json({ practice, matches });
});

// ── Static files ──────────────────────────────────────────────────────────────
app.use(express.static(path.join(__dirname, 'public'), {
  etag: false,
  setHeaders: (res) => res.setHeader('Cache-Control', 'no-store'),
}));
app.use((req, res) => res.sendFile(path.join(__dirname, 'public', 'index.html')));

// ── Boot ──────────────────────────────────────────────────────────────────────
async function main() {
  initDB();
  await initUsers();

  const server = app.listen(PORT, () => {
    console.log(`[helldivers] listening on http://localhost:${PORT}`);
  });

  // ── WebSocket server ────────────────────────────────────────────────────────
  const wss = new WebSocket.Server({ server });

  function send(ws, type, payload) {
    if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify({ type, payload }));
  }

  function broadcastLobbyUpdate() {
    const online = [...userSockets.keys()];
    wss.clients.forEach(ws => {
      if (ws.readyState !== WebSocket.OPEN || !ws.userId) return;
      const incoming = [...pendingChallenges.entries()]
        .filter(([, target]) => target === ws.userId)
        .map(([from]) => from);
      send(ws, 'lobby-update', { online, incoming });
    });
  }

  function getRoomForUser(userId) {
    for (const room of rooms.values()) {
      if (room.players.some(p => p.userId === userId)) return room;
    }
    return null;
  }

  function startRound(room) {
    room.state   = 'active';
    room.current = STRATAGEMS[Math.floor(Math.random() * STRATAGEMS.length)];
    room.players.forEach(p => { p.progress = 0; });
    broadcastToRoom(room, 'round-start', { stratagem: room.current });
  }

  function resolveRound(room, winnerId) {
    const loser = room.players.find(p => p.userId !== winnerId);
    room.matchScores[winnerId]++;
    const matchScores = { ...room.matchScores };

    broadcastToRoom(room, 'round-complete', { winner: winnerId, matchScores });

    if (room.matchScores[winnerId] >= 5) {
      broadcastToRoom(room, 'match-end', { winner: winnerId, matchScores });
      db.prepare(`
        INSERT INTO matches (winner, loser, winner_rounds, loser_rounds, created_at)
        VALUES (?, ?, ?, ?, ?)
      `).run(winnerId, loser.userId, room.matchScores[winnerId], room.matchScores[loser.userId], new Date().toISOString());
      invalidateLB();
      rooms.delete(room.roomId);
    } else {
      room.state = 'waiting';
      room.readyPlayers.clear();
    }
  }

  function broadcastToRoom(room, type, payload) {
    room.players.forEach(({ userId }) => {
      const ws = userSockets.get(userId);
      if (ws) send(ws, type, payload);
    });
  }

  function handleMessage(ws, type, payload) {
    const userId = ws.userId;

    switch (type) {
      case 'challenge-user': {
        const { targetUser } = payload;
        if (!userSockets.has(targetUser) || targetUser === userId) return;
        pendingChallenges.set(userId, targetUser);
        send(userSockets.get(targetUser), 'challenge-received', { from: userId });
        broadcastLobbyUpdate();
        break;
      }

      case 'accept-challenge': {
        const { challengerId } = payload;
        if (!userSockets.has(challengerId) || pendingChallenges.get(challengerId) !== userId) return;
        pendingChallenges.delete(challengerId);

        const roomId = crypto.randomUUID();
        const room = {
          roomId,
          players: [{ userId: challengerId, progress: 0 }, { userId, progress: 0 }],
          state:   'waiting',
          current: null,
          matchScores:  { [challengerId]: 0, [userId]: 0 },
          readyPlayers: new Set(),
        };
        rooms.set(roomId, room);

        send(userSockets.get(challengerId), 'room-joined', { roomId, opponent: userId,        matchScores: room.matchScores });
        send(ws,                             'room-joined', { roomId, opponent: challengerId, matchScores: room.matchScores });
        broadcastLobbyUpdate();
        break;
      }

      case 'decline-challenge': {
        const { challengerId } = payload;
        pendingChallenges.delete(challengerId);
        const cWs = userSockets.get(challengerId);
        if (cWs) send(cWs, 'challenge-declined', { by: userId });
        broadcastLobbyUpdate();
        break;
      }

      case 'player-ready': {
        const room = getRoomForUser(userId);
        if (!room || room.state !== 'waiting') return;
        room.readyPlayers.add(userId);
        if (room.readyPlayers.size >= 2) startRound(room);
        break;
      }

      case 'input-arrow': {
        const room = getRoomForUser(userId);
        if (!room || room.state !== 'active') return;
        const { direction } = payload;
        if (!['up','down','left','right'].includes(direction)) return;

        const player   = room.players.find(p => p.userId === userId);
        const expected = room.current.sequence[player.progress];

        if (direction === expected) {
          player.progress++;
          broadcastToRoom(room, 'input-result', { userId, correct: true, progress: player.progress });

          if (player.progress === room.current.sequence.length) {
            room.state = 'round-resolving'; // lock before any async work
            resolveRound(room, userId);
          }
        } else {
          player.progress = 0;
          broadcastToRoom(room, 'input-result', { userId, correct: false, progress: 0 });
        }
        break;
      }

      case 'leave-room': {
        const room = getRoomForUser(userId);
        if (!room) return;
        const opponent = room.players.find(p => p.userId !== userId);
        if (opponent) {
          const oWs = userSockets.get(opponent.userId);
          if (oWs) send(oWs, 'opponent-left', {});
        }
        rooms.delete(room.roomId);
        break;
      }
    }
  }

  wss.on('connection', (ws, req) => {
    // Re-use session middleware to authenticate the WS upgrade request
    sessionMiddleware(req, {}, () => {
      if (!req.session?.user) { ws.close(1008, 'Unauthorized'); return; }

      const userId = req.session.user;

      // Close any stale socket for this user
      const stale = userSockets.get(userId);
      if (stale && stale !== ws) stale.terminate();

      userSockets.set(userId, ws);
      ws.userId  = userId;
      ws.isAlive = true;

      broadcastLobbyUpdate();

      ws.on('message', (raw) => {
        try {
          const { type, payload } = JSON.parse(raw.toString());
          handleMessage(ws, type, payload || {});
        } catch { /* ignore malformed */ }
      });

      ws.on('pong', () => { ws.isAlive = true; });

      ws.on('close', () => {
        userSockets.delete(userId);
        pendingChallenges.delete(userId);

        // Notify opponent if in a room
        const room = getRoomForUser(userId);
        if (room) {
          const opponent = room.players.find(p => p.userId !== userId);
          if (opponent) {
            const oWs = userSockets.get(opponent.userId);
            if (oWs) send(oWs, 'opponent-left', {});
          }
          rooms.delete(room.roomId);
        }

        broadcastLobbyUpdate();
      });
    });
  });

  // Heartbeat – terminates stale connections every 30s
  const heartbeat = setInterval(() => {
    wss.clients.forEach(ws => {
      if (!ws.isAlive) { ws.terminate(); return; }
      ws.isAlive = false;
      ws.ping();
    });
  }, 30_000);

  wss.on('close', () => clearInterval(heartbeat));
}

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

process.on('SIGTERM', () => { db.close(); process.exit(0); });
process.on('SIGINT',  () => { db.close(); process.exit(0); });