server.js

const express = require('express');
const http = require('http');
const WebSocket = require('ws');
const path = require('path');
const helmet = require('helmet');
const rateLimit = require('express-rate-limit');

const Room = require('./lib/Room');
const User = require('./lib/User');
const InviteURL = require('./lib/InviteURL');
const MessageHandler = require('./lib/MessageHandler');
const TorService = require('./services/torService');
const { securityMiddleware, apiLimiter, roomCreationLimiter, validateInput, checkWSRateLimit, checkWSMessageRate } = require('./middleware/security');

const app = express();
const server = http.createServer(app);
const wss = new WebSocket.Server({ server });

const PORT = process.env.PORT || 8000;

// Security middleware
app.use(helmet({
  contentSecurityPolicy: {
    directives: {
      defaultSrc: ["'self'"],
      scriptSrc: ["'self'"],
      scriptSrcAttr: ["'unsafe-inline'"],
      styleSrc: ["'self'", "'unsafe-inline'"],
      imgSrc: ["'self'"],
      connectSrc: ["'self'"],
      fontSrc: ["'self'"],
      objectSrc: ["'none'"],
      mediaSrc: ["'self'", "data:"],
      frameSrc: ["'none'"],
      upgradeInsecureRequests: null
    }
  }
}));

// Apply security middleware
app.use(apiLimiter);
app.use(securityMiddleware);
app.use(validateInput);

// Parse JSON bodies
app.use(express.json({ limit: '10mb' }));

// Serve static files
app.use(express.static('public'));

// Global state
const rooms = new Map();
const messageHandler = new MessageHandler();
const torService = new TorService();

// WebSocket connection handling
wss.on('connection', (ws, req) => {
  // Check WebSocket rate limit
  if (!checkWSRateLimit(req)) {
    ws.close(1008, 'Rate limit exceeded');
    return;
  }

  const user = new User();
  let currentRoom = null;

  ws.on('message', async (data) => {
    try {
      // Check message rate limit
      if (!checkWSMessageRate(user.id)) {
        ws.send(JSON.stringify({
          type: 'error',
          message: 'Message rate limit exceeded'
        }));
        return;
      }

      const message = JSON.parse(data);
      
      switch (message.type) {
        case 'join_room':
          const room = rooms.get(message.roomId);
          if (room && room.canJoin(message.inviteCode)) {
            // Burn the invite code after successful join
            InviteURL.burn(message.inviteCode);
            
            currentRoom = room;
            room.addUser(user, ws, message.inviteCode);
            
            // Send current user list to new joiner
            const userList = Array.from(room.users.values()).map(userData => ({
              userId: userData.user.id,
              username: userData.user.username,
              joinedAt: userData.joinedAt,
              inviteCode: userData.inviteCode
            }));
            
            ws.send(JSON.stringify({
              type: 'room_joined',
              roomId: message.roomId,
              userId: user.id,
              username: user.username,
              persistenceMode: room.persistenceMode,
              encryptionKey: room.encryptionKey,
              userList: userList
            }));
          } else {
            ws.send(JSON.stringify({
              type: 'error',
              message: 'Invalid room or invite code'
            }));
          }
          break;

        case 'join_room_admin':
          const adminRoom = rooms.get(message.roomId);
          if (adminRoom && adminRoom.verifyAdmin(message.adminToken)) {
            currentRoom = adminRoom;
            adminRoom.addUser(user, ws, null); // Admin doesn't use invite code
            
            // Send current user list to admin
            const userList = Array.from(adminRoom.users.values()).map(userData => ({
              userId: userData.user.id,
              username: userData.user.username,
              joinedAt: userData.joinedAt,
              isAdmin: userData.user.id === user.id,
              inviteCode: userData.inviteCode
            }));
            
            ws.send(JSON.stringify({
              type: 'room_joined',
              roomId: message.roomId,
              userId: user.id,
              username: user.username,
              persistenceMode: adminRoom.persistenceMode,
              encryptionKey: adminRoom.encryptionKey,
              userList: userList,
              isAdmin: true
            }));
          } else {
            ws.send(JSON.stringify({
              type: 'error',
              message: 'Invalid room or admin token'
            }));
          }
          break;

        case 'send_message':
          if (currentRoom && message.encrypted) {
            // Only handle encrypted messages - no plaintext accepted
            await messageHandler.handleMessage(currentRoom, user, null, message.encrypted);
          }
          break;

        case 'ratchet_key_exchange':
          if (currentRoom && message.publicKey && message.userId) {
            // Broadcast Double Ratchet public key to other room members
            currentRoom.broadcast({
              type: 'ratchet_key_exchange',
              publicKey: message.publicKey,
              userId: message.userId,
              timestamp: Date.now()
            }, user.id); // Exclude the sender
          }
          break;

        case 'leave_room':
          if (currentRoom) {
            currentRoom.removeUser(user);
            currentRoom = null;
          }
          break;
      }
    } catch (error) {
      console.error('WebSocket message processing failed');
      ws.send(JSON.stringify({
        type: 'error',
        message: 'Invalid message format'
      }));
    }
  });

  ws.on('close', () => {
    if (currentRoom) {
      currentRoom.removeUser(user);
    }
  });
});

// HTTP API endpoints
app.post('/api/room/create', roomCreationLimiter, async (req, res) => {
  try {
    const { persistenceMode = 'ephemeral', adminPassword } = req.body;
    
    const room = new Room(persistenceMode);
    const inviteUrl = new InviteURL(room.id);
    
    rooms.set(room.id, room);
    
    res.json({
      roomId: room.id,
      inviteUrl: inviteUrl.url,
      inviteCode: inviteUrl.code,
      adminToken: room.adminToken,
      persistenceMode: room.persistenceMode
    });
  } catch (error) {
    res.status(500).json({ error: 'Failed to create room' });
  }
});

app.post('/api/room/:roomId/invite', async (req, res) => {
  try {
    const room = rooms.get(req.params.roomId);
    if (!room || !room.verifyAdmin(req.body.adminToken)) {
      return res.status(403).json({ error: 'Unauthorized' });
    }

    const inviteUrl = new InviteURL(room.id);
    res.json({ inviteUrl: inviteUrl.url });
  } catch (error) {
    res.status(500).json({ error: 'Failed to create invite' });
  }
});

app.delete('/api/room/:roomId/user/:userId', async (req, res) => {
  try {
    const room = rooms.get(req.params.roomId);
    if (!room || !room.verifyAdmin(req.body.adminToken)) {
      return res.status(403).json({ error: 'Unauthorized' });
    }

    room.kickUser(req.params.userId);
    res.json({ success: true });
  } catch (error) {
    res.status(500).json({ error: 'Failed to kick user' });
  }
});

// API endpoint to get room ID from invite code
app.get('/api/invite/:inviteCode/room', (req, res) => {
  try {
    const { inviteCode } = req.params;
    const roomId = InviteURL.getRoomId(inviteCode);
    
    if (!roomId || !rooms.has(roomId)) {
      return res.status(404).json({ error: 'Invalid or expired invite code' });
    }
    
    // Check if invite is still valid (not expired or used)
    const invite = InviteURL.inviteCodes.get(inviteCode);
    if (!invite || invite.used || (Date.now() - invite.createdAt > invite.expirationMs)) {
      return res.status(404).json({ error: 'Invite code has expired or been used' });
    }
    
    res.json({ roomId: roomId });
  } catch (error) {
    res.status(500).json({ error: 'Server error' });
  }
});

// Join room via invite URL
app.get('/join/:inviteCode', (req, res) => {
  try {
    const { inviteCode } = req.params;
    const roomId = InviteURL.getRoomId(inviteCode);
    
    if (!roomId || !rooms.has(roomId)) {
      return res.status(404).send('Invalid or expired invite link');
    }
    
    // Serve the main page with invite code in URL for client-side handling
    res.sendFile(path.join(__dirname, 'public', 'index.html'));
  } catch (error) {
    res.status(500).send('Server error');
  }
});

// Health check endpoint
app.get('/health', (req, res) => {
  res.json({ status: 'ok', timestamp: Date.now() });
});

// Cleanup expired rooms every 5 minutes
setInterval(() => {
  for (const [roomId, room] of rooms.entries()) {
    if (room.isExpired()) {
      room.cleanup();
      rooms.delete(roomId);
    }
  }
}, 5 * 60 * 1000);

// Initialize Tor service and start server
async function startServer() {
  try {
    console.log('Starting Tor hidden service...');
    await torService.start();
    
    const onionAddress = await torService.getOnionAddress();
    console.log(`Tor hidden service running at: ${onionAddress}`);
    
    server.listen(PORT, '127.0.0.1', () => {
      console.log(`Anonymous messaging server running on http://127.0.0.1:${PORT}`);
      console.log(`Tor hidden service accessible at: ${onionAddress}`);
    });
  } catch (error) {
    console.error('Failed to start services:', error.message);
    console.log('Starting server without Tor service...');
    
    server.listen(PORT, '127.0.0.1', () => {
      console.log(`Anonymous messaging server running on http://127.0.0.1:${PORT}`);
      console.log('Warning: Tor service not available');
    });
  }
}

startServer();

// Graceful shutdown
process.on('SIGTERM', async () => {
  console.log('Shutting down gracefully...');
  
  // Cleanup all rooms
  for (const room of rooms.values()) {
    room.cleanup();
  }
  
  // Stop Tor service
  await torService.stop();
  
  server.close(() => {
    process.exit(0);
  });
});

process.on('SIGINT', async () => {
  console.log('\nReceived SIGINT, shutting down gracefully...');
  
  // Cleanup all rooms
  for (const room of rooms.values()) {
    room.cleanup();
  }
  
  // Stop Tor service
  await torService.stop();
  
  server.close(() => {
    process.exit(0);
  });
});

module.exports = { app, server, wss };