'use strict';

// Agent registry and tunnel broker.
//
// Agent-mode clients sit behind NAT, so they dial *out* to the hub and hold a
// control WebSocket open. When someone wants to view that machine, the hub asks
// the agent over that control channel to open a second, data-only WebSocket; the
// agent pipes it to the local VNC server. The hub pairs that data socket with
// the waiting browser session.

const { createWebSocketStream } = require('ws');
const config = require('../config');
const { clients, audit } = require('../db');
const { randomToken } = require('../crypto');

const HEARTBEAT_MS = 30_000;

class Hub {
  constructor() {
    /** clientId -> { ws, info, connectedAt, lastSeen, alive } */
    this.agents = new Map();
    /** tunnelId -> { clientId, resolve, reject, timer } */
    this.pending = new Map();

    this.heartbeat = setInterval(() => this._sweep(), HEARTBEAT_MS);
    this.heartbeat.unref?.();
  }

  /* ------------------------------------------------------------- control */

  handleAgentSocket(ws, client, remoteIp) {
    // A machine may only have one live control channel; a reconnect wins.
    const existing = this.agents.get(client.id);
    if (existing && existing.ws !== ws) {
      try { existing.ws.close(4001, 'replaced by a newer connection'); } catch { /* already gone */ }
    }

    const entry = { ws, info: {}, connectedAt: Date.now(), lastSeen: Date.now(), alive: true, lastIp: remoteIp };
    this.agents.set(client.id, entry);
    clients.touch(client.id, remoteIp);

    ws.on('pong', () => {
      entry.alive = true;
      entry.lastSeen = Date.now();
      clients.touch(client.id, remoteIp);
    });

    ws.on('message', (raw) => {
      let msg;
      try {
        msg = JSON.parse(raw.toString());
      } catch {
        return;
      }
      entry.lastSeen = Date.now();
      this._onAgentMessage(client, entry, msg, remoteIp);
    });

    ws.on('close', () => {
      if (this.agents.get(client.id) === entry) this.agents.delete(client.id);
      // Fail anything that was waiting on this agent.
      for (const [tunnelId, p] of this.pending) {
        if (p.clientId === client.id) this._rejectPending(tunnelId, new Error('agent disconnected'));
      }
    });

    ws.on('error', () => { /* close handler does the cleanup */ });

    this._send(ws, {
      type: 'welcome',
      clientId: client.id,
      name: client.name,
      requireConsent: !!client.require_consent,
      heartbeatMs: HEARTBEAT_MS,
    });
  }

  _onAgentMessage(client, entry, msg, remoteIp) {
    switch (msg.type) {
      case 'hello': {
        entry.info = {
          version: msg.version,
          os: msg.os,
          hostname: msg.hostname,
          vncPort: msg.vncPort,
        };
        clients.update(client.id, {
          os: msg.os ?? null,
          hostname: msg.hostname ?? null,
          agent_version: msg.version ?? null,
          last_seen_at: Date.now(),
          last_ip: remoteIp ?? null,
        });
        break;
      }
      case 'denied': {
        this._rejectPending(msg.tunnelId, new Error(msg.reason || 'the person at that machine declined'));
        break;
      }
      case 'error': {
        this._rejectPending(msg.tunnelId, new Error(msg.message || 'agent reported an error'));
        break;
      }
      default:
        break;
    }
  }

  _send(ws, obj) {
    if (ws.readyState === ws.OPEN) ws.send(JSON.stringify(obj));
  }

  _sweep() {
    const now = Date.now();
    for (const [clientId, entry] of this.agents) {
      if (!entry.alive) {
        try { entry.ws.terminate(); } catch { /* already gone */ }
        this.agents.delete(clientId);
        continue;
      }
      entry.alive = false;
      try { entry.ws.ping(); } catch { /* handled on next sweep */ }
      if (now - entry.lastSeen < config.agentOfflineAfterMs) clients.touch(clientId, entry.lastIp);
    }
  }

  /* -------------------------------------------------------------- status */

  isOnline(clientId) {
    return this.agents.has(clientId);
  }

  onlineIds() {
    return Array.from(this.agents.keys());
  }

  agentInfo(clientId) {
    const e = this.agents.get(clientId);
    return e ? { ...e.info, connectedAt: e.connectedAt, lastSeen: e.lastSeen } : null;
  }

  /* ------------------------------------------------------------- tunnels */

  /**
   * Ask an agent to open a data tunnel. Resolves with a Duplex carrying the raw
   * RFB byte stream from the client machine's local VNC server.
   */
  openTunnel(clientId, meta = {}) {
    const entry = this.agents.get(clientId);
    if (!entry) return Promise.reject(new Error('that machine is offline'));

    const tunnelId = randomToken(16);
    const requireConsent = !!meta.requireConsent;
    const timeoutMs = requireConsent ? config.consentTimeoutMs : 15_000;

    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pending.delete(tunnelId);
        reject(new Error(requireConsent
          ? 'no response to the connection request on that machine'
          : 'the agent did not open a tunnel in time'));
      }, timeoutMs);
      timer.unref?.();

      this.pending.set(tunnelId, { clientId, resolve, reject, timer });

      this._send(entry.ws, {
        type: 'open',
        tunnelId,
        requireConsent,
        operator: meta.operator || 'someone',
        role: meta.role || 'viewer',
        sessionId: meta.sessionId,
      });
    });
  }

  /** Called when the agent's data socket arrives and claims a pending tunnel. */
  handleTunnelSocket(ws, clientId, tunnelId) {
    const pending = this.pending.get(tunnelId);
    if (!pending || pending.clientId !== clientId) {
      try { ws.close(4004, 'unknown tunnel'); } catch { /* nothing to do */ }
      return false;
    }
    clearTimeout(pending.timer);
    this.pending.delete(tunnelId);
    pending.resolve(createWebSocketStream(ws, { allowHalfOpen: false }));
    return true;
  }

  _rejectPending(tunnelId, err) {
    const pending = this.pending.get(tunnelId);
    if (!pending) return;
    clearTimeout(pending.timer);
    this.pending.delete(tunnelId);
    pending.reject(err);
  }

  /** Tell an agent to drop whatever it is doing (used by force-disconnect). */
  notifySessionEnded(clientId, sessionId) {
    const entry = this.agents.get(clientId);
    if (entry) this._send(entry.ws, { type: 'session-ended', sessionId });
  }

  disconnectAgent(clientId, reason = 'removed') {
    const entry = this.agents.get(clientId);
    if (!entry) return false;
    try { entry.ws.close(4003, reason); } catch { /* already gone */ }
    this.agents.delete(clientId);
    audit(null, 'agent.disconnect', clientId, { reason });
    return true;
  }
}

module.exports = new Hub();