server.js

const http = require('http');
const fs = require('fs');
const path = require('path');
const { WebSocketServer, WebSocket } = require('ws');
const { v4: uuidv4 } = require('uuid');
const crypto = require('crypto');

// ─── Auth ─────────────────────────────────────────────────────────────────────
// Secret key from environment variable "Key"
const SECRET = process.env.Key || '';
if (!SECRET) console.warn('[Auth] WARNING: Environment variable "Key" is not set. Admin login will be disabled.');

// HMAC-based session tokens: token = hmac(SECRET, nonce+timestamp)
// Sessions are stored server-side; possession of token alone is not enough —
// the server must have issued it.
const activeSessions = new Set(); // Set<token>

function createSession() {
  const nonce = crypto.randomBytes(32).toString('hex');
  const token = crypto.createHmac('sha256', SECRET).update(nonce).digest('hex');
  activeSessions.add(token);
  return token;
}

function isValidSession(token) {
  return typeof token === 'string' && activeSessions.has(token);
}

// ─── In-Memory State ──────────────────────────────────────────────────────────
const sounds = {};
const notifications = {};
const devices = {};

const deviceClients = {};
const adminClients = new Set(); // only authenticated WS connections

const MIME = {
  '.html': 'text/html',
  '.js':   'application/javascript',
  '.css':  'text/css',
  '.json': 'application/json',
  '.png':  'image/png',
  '.ico':  'image/x-icon',
};

// ─── HTTP File Server ─────────────────────────────────────────────────────────
const httpServer = http.createServer((req, res) => {
  let urlPath = req.url.split('?')[0];
  if (urlPath === '/' || urlPath === '') urlPath = '/index.html';
  if (urlPath === '/admin' || urlPath === '/admin/') urlPath = '/admin/index.html';

  const filePath = path.join(__dirname, 'public', urlPath);

  // Prevent path traversal
  const publicDir = path.resolve(__dirname, 'public');
  const resolved = path.resolve(filePath);
  if (!resolved.startsWith(publicDir)) {
    res.writeHead(403); res.end('Forbidden'); return;
  }

  const ext = path.extname(filePath);
  const mime = MIME[ext] || 'application/octet-stream';

  fs.readFile(filePath, (err, data) => {
    if (err) { res.writeHead(404, { 'Content-Type': 'text/plain' }); res.end('Not found'); return; }
    res.writeHead(200, { 'Content-Type': mime });
    res.end(data);
  });
});

// ─── WebSocket Server ─────────────────────────────────────────────────────────
const wss = new WebSocketServer({ server: httpServer });

wss.on('connection', (ws, req) => {
  const urlPath = req.url || '/';
  if (urlPath.startsWith('/admin-ws')) {
    handleAdminConnection(ws);
  } else {
    handleDeviceConnection(ws);
  }
});

// ─── Helpers ──────────────────────────────────────────────────────────────────
function send(ws, msg) {
  if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(msg));
}

function broadcastAdmin(msg) {
  for (const ws of adminClients) send(ws, msg);
}

function broadcastDevice(uuid, msg) {
  const ws = deviceClients[uuid];
  if (ws) send(ws, msg);
}

function broadcastAllDevices(msg) {
  for (const uuid of Object.keys(deviceClients)) broadcastDevice(uuid, msg);
}

function getFullState() {
  return {
    sounds: Object.values(sounds),
    notifications: Object.values(notifications),
    devices: Object.values(devices).map(devicePublic),
  };
}

function devicePublic(d) {
  return {
    uuid: d.uuid,
    name: d.name,
    notifications: d.notifications,
    lastConnection: d.lastConnection,
    schedule: d.schedule || [],
    pendingChanges: d.pendingChanges || [],
    online: !!deviceClients[d.uuid],
    cachedSounds: d.cachedSounds || [],
  };
}

// ─── Admin Connection Handler ─────────────────────────────────────────────────
function handleAdminConnection(ws) {
  let authenticated = false;

  // Give them 10 seconds to authenticate, then close
  const authTimeout = setTimeout(() => {
    if (!authenticated) {
      send(ws, { type: 'auth_timeout' });
      ws.close();
    }
  }, 10000);

  ws.on('message', (raw) => {
    let msg;
    try { msg = JSON.parse(raw); } catch { return; }

    // ── Authentication handshake ──────────────────────────────────────────────
    if (!authenticated) {
      if (msg.type === 'auth_resume' && isValidSession(msg.token)) {
        // Client is resuming a valid session
        authenticated = true;
        clearTimeout(authTimeout);
        adminClients.add(ws);
        send(ws, { type: 'auth_ok', token: msg.token });
        send(ws, { type: 'full_state', ...getFullState() });
        console.log('[Admin] session resumed');
        return;
      }

      if (msg.type === 'auth_login') {
        if (!SECRET) {
          send(ws, { type: 'auth_error', reason: 'Server has no Key configured.' });
          return;
        }
        // Constant-time comparison to prevent timing attacks
        const provided = Buffer.from(String(msg.password || ''));
        const expected = Buffer.from(SECRET);
        const match = provided.length === expected.length &&
                      crypto.timingSafeEqual(provided, expected);
        if (match) {
          authenticated = true;
          clearTimeout(authTimeout);
          const token = createSession();
          adminClients.add(ws);
          send(ws, { type: 'auth_ok', token });
          send(ws, { type: 'full_state', ...getFullState() });
          console.log('[Admin] authenticated');
        } else {
          send(ws, { type: 'auth_error', reason: 'Invalid password.' });
          console.warn('[Admin] failed login attempt');
        }
        return;
      }

      // Any other message before auth: reject silently
      send(ws, { type: 'auth_required' });
      return;
    }

    // ── Authenticated messages ─────────────────────────────────────────────────
    handleAdminMessage(ws, msg);
  });

  ws.on('close', () => {
    adminClients.delete(ws);
    clearTimeout(authTimeout);
    console.log(`[Admin] disconnected, total=${adminClients.size}`);
  });

  ws.on('error', (e) => console.error('[Admin] WS error', e.message));

  // Prompt client to authenticate
  send(ws, { type: 'auth_required' });
}

function handleAdminMessage(ws, msg) {
  switch (msg.type) {

    case 'create_sound': {
      const id = uuidv4();
      const sound = { id, name: msg.name, data: msg.data };
      sounds[id] = sound;
      broadcastAdmin({ type: 'sound_added', sound });
      break;
    }

    case 'create_notification': {
      const id = uuidv4();
      const notif = {
        id, name: msg.name, heading: msg.heading, body: msg.body,
        hyperlink: msg.hyperlink || '', displayed: false, soundId: msg.soundId || null,
      };
      notifications[id] = notif;
      for (const uuid of Object.keys(devices)) {
        if (!devices[uuid].notifications.find(n => n.id === id))
          devices[uuid].notifications.push({ ...notif });
      }
      broadcastAdmin({ type: 'notification_added', notification: notif });
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      broadcastAllDevices({ type: 'notification_added', notification: notif });
      break;
    }

    case 'schedule_notification': {
      const device = devices[msg.uuid];
      if (!device) return;
      device.schedule = device.schedule || [];
      const existing = device.schedule.find(s => s.notificationId === msg.notificationId);
      if (existing) { existing.scheduledAt = msg.scheduledAt; }
      else { device.schedule.push({ notificationId: msg.notificationId, scheduledAt: msg.scheduledAt }); }
      const scheduleMsg = { type: 'schedule_update', schedule: device.schedule };
      if (deviceClients[msg.uuid]) { broadcastDevice(msg.uuid, scheduleMsg); }
      else { (device.pendingChanges = device.pendingChanges || []).push(scheduleMsg); }
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }

    case 'play_now': {
      const playMsg = { type: 'play_now', notificationId: msg.notificationId };
      if (deviceClients[msg.uuid]) { broadcastDevice(msg.uuid, playMsg); }
      else {
        const device = devices[msg.uuid];
        if (device) (device.pendingChanges = device.pendingChanges || []).push(playMsg);
      }
      break;
    }

    case 'update_device_name': {
      const device = devices[msg.uuid];
      if (!device) return;
      device.name = msg.name;
      const nameMsg = { type: 'name_update', name: msg.name };
      if (deviceClients[msg.uuid]) { broadcastDevice(msg.uuid, nameMsg); }
      else { (device.pendingChanges = device.pendingChanges || []).push(nameMsg); }
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }

    case 'remove_schedule': {
      const device = devices[msg.uuid];
      if (!device) return;
      device.schedule = (device.schedule || []).filter(s => s.notificationId !== msg.notificationId);
      const rmMsg = { type: 'schedule_update', schedule: device.schedule };
      if (deviceClients[msg.uuid]) { broadcastDevice(msg.uuid, rmMsg); }
      else { (device.pendingChanges = device.pendingChanges || []).push(rmMsg); }
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }

    case 'request_sound': {
      const sound = sounds[msg.id];
      if (sound) send(ws, { type: 'sound_data', sound });
      break;
    }

    default:
      console.warn('[Admin] unknown message type:', msg.type);
  }
}

// ─── Device Connection Handler ────────────────────────────────────────────────
function handleDeviceConnection(ws) {
  let deviceUUID = null;

  ws.on('message', (raw) => {
    let msg;
    try { msg = JSON.parse(raw); } catch { return; }

    if (msg.type === 'hello') {
      if (msg.uuid && devices[msg.uuid]) {
        deviceUUID = msg.uuid;
      } else {
        deviceUUID = msg.uuid || uuidv4();
        devices[deviceUUID] = {
          uuid: deviceUUID,
          name: `Device ${Object.keys(devices).length + 1}`,
          notifications: Object.values(notifications).map(n => ({ ...n })),
          lastConnection: null, schedule: [], pendingChanges: [],
        };
      }

      deviceClients[deviceUUID] = ws;
      const device = devices[deviceUUID];
      device.lastConnection = null;

      if (device.pendingChanges && device.pendingChanges.length > 0) {
        for (const change of device.pendingChanges) send(ws, change);
        device.pendingChanges = [];
      }

      send(ws, {
        type: 'device_init', uuid: deviceUUID,
        notifications: device.notifications, schedule: device.schedule, name: device.name,
      });

      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      console.log(`[Device] ${deviceUUID} connected (${device.name})`);
      return;
    }

    if (!deviceUUID) return;
    handleDeviceMessage(ws, deviceUUID, msg);
  });

  ws.on('close', () => {
    if (deviceUUID && devices[deviceUUID]) {
      devices[deviceUUID].lastConnection = Date.now();
      delete deviceClients[deviceUUID];
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      console.log(`[Device] ${deviceUUID} disconnected`);
    }
  });

  ws.on('error', (e) => console.error('[Device] WS error', e.message));
}

function handleDeviceMessage(ws, uuid, msg) {
  const device = devices[uuid];
  if (!device) return;

  switch (msg.type) {
    case 'mark_displayed': {
      const notif = device.notifications.find(n => n.id === msg.notificationId);
      if (notif) notif.displayed = true;
      if (notifications[msg.notificationId]) notifications[msg.notificationId].displayed = true;
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }
    case 'request_sound': {
      const sound = sounds[msg.soundId];
      if (sound) send(ws, { type: 'sound_data', sound });
      break;
    }
    case 'sync_schedule': {
      device.schedule = msg.schedule || [];
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }
    case 'cached_sounds': {
      device.cachedSounds = msg.soundIds || [];
      broadcastAdmin({ type: 'devices_updated', devices: Object.values(devices).map(devicePublic) });
      break;
    }
    default:
      console.warn('[Device] unknown message type:', msg.type);
  }
}

// ─── Start ────────────────────────────────────────────────────────────────────
const PORT = process.env.PORT || 7860;
httpServer.listen(PORT, () => {
  console.log(`Server running on http://0.0.0.0:${PORT}`);
});