const WebSocket = require('ws');

const PORT = 33001;
const clients = new Map();   // uuid => ws (clients)
const webUIs = new Map();    // uuid => ws (web UIs)

const wss = new WebSocket.Server({ port: PORT });

console.log(`WebSocket server started on ws://localhost:${PORT}`);

wss.on('connection', (ws, req) => {
  let registeredId = null;
  let registeredType = null; // 'client' or 'web'

  console.log(`New WebSocket connection from ${req.socket.remoteAddress}:${req.socket.remotePort}`);

  ws.on('message', (msg) => {
    const senderLabel = registeredId
      ? (registeredType === 'client' ? `client: ${registeredId}` : `web: ${registeredId}`)
      : 'unregistered client';

    console.log(`Received message from ${senderLabel}: ${msg}`);

    let data;
    try {
      data = JSON.parse(msg);
    } catch (err) {
      ws.send(JSON.stringify({ type: 'error', message: 'Invalid JSON' }));
      return;
    }

    // Register client
    if (data.action === 'register' && data.apikey && data.equipment) {
      registeredId = data.equipment;
      registeredType = 'client';
      clients.set(registeredId, ws);
      ws.send(JSON.stringify({ answer: 'public', data: 'registered' }));
      console.log(`Client registered: ${registeredId}`);
      return;
    }

    // Register web UI
    if (data.action === 'registerweb' && data.apikey && data.uuid) {
      registeredId = data.uuid;
      registeredType = 'web';
      webUIs.set(registeredId, ws);
      ws.send(JSON.stringify({ answer: 'public', data: 'registered' }));
      console.log(`Web UI registered: ${registeredId}`);
      return;
    }

    if (!registeredId) {
      ws.send(JSON.stringify({ type: 'error', message: 'Not registered' }));
      return;
    }

    if (!data.uuid) {
      ws.send(JSON.stringify({ type: 'error', message: 'Missing target uuid (uuid)' }));
      return;
    }

    // Determine target ws based on sender type
    let targetWs = null;

    if (registeredType === 'web') {
      targetWs = clients.get(data.uuid);
      if (!targetWs) {
        ws.send(JSON.stringify({ type: 'error', message: `Target client ${data.uuid} not connected` }));
        return;
      }
    } else if (registeredType === 'client') {
      targetWs = webUIs.get(data.uuid);
      if (!targetWs) {
        ws.send(JSON.stringify({ type: 'error', message: `Target web UI ${data.uuid} not connected` }));
        return;
      }
    } else {
      ws.send(JSON.stringify({ type: 'error', message: 'Unknown registered type' }));
      return;
    }

    if (targetWs.readyState !== WebSocket.OPEN) {
      ws.send(JSON.stringify({ type: 'error', message: `Target ${data.uuid} connection not open` }));
      return;
    }

    // Handle your specific routing and modification logic here

    if (registeredType === 'web') {
      // From web UI to client
      if (data.action === 'reset') {
        // Modify and forward reset message
        const forwardMsg = JSON.stringify({
          run: 'reset',
          commands: 'reset',
          simtty: data.tty
        });
        targetWs.send(forwardMsg);
        ws.send(JSON.stringify({ type: 'success', message: `Sent reset command to client ${data.uuid}` }));
        console.log(`Forwarded reset command from web UI ${registeredId} to client ${data.uuid}`);
        return;
      } else if (data.action === 'send') {
        // Send a sequence of commands to client
        const cmds = [
          { run: 'emptyqueue', commands: 'emptyqueue', simtty: data.tty },
          { run: 'AT', simtty: data.tty, commands: `${data.at}\\r\\n` },
          { run: 'AT', simtty: data.tty, commands: 'wait\\r\\n' }
        ];
        try {
          cmds.forEach(cmd => targetWs.send(JSON.stringify(cmd)));
          ws.send(JSON.stringify({ type: 'success', message: `Sent AT commands sequence to client ${data.uuid}` }));
          console.log(`Sent AT command sequence from web UI ${registeredId} to client ${data.uuid}`);
        } catch (err) {
          ws.send(JSON.stringify({ type: 'error', message: 'Failed to send commands to client' }));
        }
        return;
      }
    }

    // For other message types or from client to web UI, forward as-is
    targetWs.send(msg);
    ws.send(JSON.stringify({ type: 'success', message: `Forwarded message to ${data.uuid}` }));
    console.log(`Forwarded message from ${registeredId} (${registeredType}) to ${data.uuid}`);
  });

  ws.on('close', () => {
    if (registeredId) {
      if (registeredType === 'client') {
        clients.delete(registeredId);
        console.log(`Client disconnected: ${registeredId}`);
      } else if (registeredType === 'web') {
        webUIs.delete(registeredId);
        console.log(`Web UI disconnected: ${registeredId}`);
      }
    }
  });
});
