const WebSocket = require('ws');
const net = require('net');
const process = require('process');

// Configuration
const WS_HOST = process.argv[2] || '127.0.0.1:8090';
const [HOST, PORT] = WS_HOST.split(':');
const START_PORT = 9022;

// Client tracking
const clients = new Map(); // identifier -> { ws, port, server, socket, lastActive }
const portAllocation = new Map(); // port -> identifier
const PORT_TIMEOUT = 24 * 60 * 60 * 1000; // 24 hours

const wss = new WebSocket.Server({ host: HOST, port: PORT });
console.log(`WebSocket server running on ws://${HOST}:${PORT}`);

// Cleanup old ports
setInterval(() => {
  const now = Date.now();
  for (const [identifier, client] of clients.entries()) {
    if (now - client.lastActive > PORT_TIMEOUT) {
      console.log(`Releasing port ${client.port} from ${identifier} (inactive 24h)`);
      client.server.close();
      portAllocation.delete(client.port);
      clients.delete(identifier);
    }
  }
}, 60 * 60 * 1000); // Check hourly

wss.on('connection', (ws) => {
  ws.on('message', (message) => {
    try {
      const msg = JSON.parse(message);

      // Handle registration
      if (msg.type === 'register' && msg.identifier) {
        registerClient(ws, msg.identifier);
        return;
      }

      // Handle data forwarding
      const client = clients.get(msg.identifier);
      if (client && client.ws === ws) {
        client.lastActive = Date.now();
        if (msg.type === 'data' && client.socket) {
          client.socket.write(Buffer.from(msg.data, 'base64'));
        }
      }
    } catch (err) {
      console.error('Message processing error:', err);
    }
  });

  ws.on('close', () => {
    for (const [identifier, client] of clients.entries()) {
      if (client.ws === ws) {
        console.log(`Client ${identifier} disconnected`);
        // Don't close the server, just mark WS as disconnected and close socket
        client.ws = null;
        if (client.socket) {
          client.socket.end();
          client.socket = null;
        }
      }
    }
  });
});

function registerClient(ws, identifier) {
  // Check if this identifier already has a port assigned
  const existingClient = clients.get(identifier);
  
  if (existingClient) {
    // Reuse the existing port
    existingClient.ws = ws;
    existingClient.lastActive = Date.now();
    console.log(`[${identifier}] Reconnected, using existing port ${existingClient.port}`);
    
    ws.send(JSON.stringify({ type: 'registered', port: existingClient.port }));
    return;
  }

  // Find available port for new client
  let port = START_PORT;
  while (portAllocation.has(port)) port++;

  const server = net.createServer((socket) => {
    const client = clients.get(identifier);
    if (client) {
      client.socket = socket;
      console.log(`[${identifier}] New connection on port ${port}`);
    }

    socket.on('data', (data) => {
      const client = clients.get(identifier);
      if (client && client.ws) {
        client.ws.send(JSON.stringify({
          type: 'data',
          identifier,
          data: data.toString('base64')
        }));
      }
    });

    socket.on('close', () => {
      const client = clients.get(identifier);
      if (client) {
        client.socket = null;
      }
    });
  });

  server.listen(port, '0.0.0.0', () => {
    clients.set(identifier, { ws, port, server, socket: null, lastActive: Date.now() });
    portAllocation.set(port, identifier);
    console.log(`[${identifier}] Assigned port ${port}`);
    ws.send(JSON.stringify({ type: 'registered', port }));
  });
}