'use strict';

// Splices a browser WebSocket to a client machine's VNC server, keeping a
// registry of what is live so operators can see and kill active sessions.

const net = require('net');
const { Transform, pipeline } = require('stream');
const { createWebSocketStream } = require('ws');

const { decrypt } = require('../crypto');
const { sessions, clients, audit } = require('../db');
const { handshakeWithServer, handshakeWithBrowser, ViewOnlyFilter, ByteReader } = require('./rfb');
const { canControl } = require('../auth');
const hub = require('./hub');

/** sessionId -> live session handle */
const live = new Map();

class Counter extends Transform {
  constructor() {
    super();
    this.bytes = 0;
  }
  _transform(chunk, _enc, cb) {
    this.bytes += chunk.length;
    cb(null, chunk);
  }
}

/** Drops input-bearing RFB messages so a viewer physically cannot control. */
class ViewOnlyTransform extends Transform {
  constructor() {
    super();
    this.filter = new ViewOnlyFilter();
  }
  _transform(chunk, _enc, cb) {
    let out;
    try {
      out = this.filter.push(chunk);
    } catch (err) {
      return cb(err);
    }
    cb(null, out || undefined);
  }
  get blocked() {
    return this.filter.blocked;
  }
}

function connectDirect(client, timeoutMs = 10_000) {
  return new Promise((resolve, reject) => {
    if (!client.host) return reject(new Error('this client has no host configured'));
    const socket = net.connect({ host: client.host, port: client.port || 5900 });
    socket.setNoDelay(true);

    const timer = setTimeout(() => {
      socket.destroy();
      reject(new Error(`timed out connecting to ${client.host}:${client.port || 5900}`));
    }, timeoutMs);

    socket.once('connect', () => {
      clearTimeout(timer);
      socket.removeListener('error', onError);
      resolve(socket);
    });
    const onError = (err) => {
      clearTimeout(timer);
      reject(new Error(`cannot reach ${client.host}:${client.port || 5900} (${err.code || err.message})`));
    };
    socket.once('error', onError);
  });
}

async function connectUpstream(client, meta) {
  if (client.mode === 'agent') {
    return hub.openTunnel(client.id, {
      requireConsent: !!client.require_consent,
      operator: meta.username,
      role: meta.role,
      sessionId: meta.sessionId,
    });
  }
  return connectDirect(client);
}

/**
 * Take over a browser WebSocket and run a VNC session on it.
 * `ctx` = { sessionId, client, username, role, source, inviteId, remoteIp }
 */
async function startSession(browserWs, ctx) {
  const { sessionId, client, username, role } = ctx;
  const control = canControl(role);

  let upstream;
  try {
    upstream = await connectUpstream(client, { username, role, sessionId });
    await handshakeWithServer(upstream, decrypt(client.vnc_password_enc));
  } catch (err) {
    if (upstream) upstream.destroy();
    audit(username, 'session.failed', client.id, { error: err.message });
    // The browser has not started RFB yet, so a close reason is still readable
    // by the viewer page. WebSocket close reasons are capped at 123 bytes.
    try { browserWs.close(4500, String(err.message).slice(0, 120)); } catch { /* gone */ }
    return;
  }

  sessions.start({
    id: sessionId,
    client_id: client.id,
    client_name: client.name,
    username,
    source: ctx.source || 'web',
    invite_id: ctx.inviteId || null,
    role,
    remote_ip: ctx.remoteIp || null,
  });

  const browserStream = createWebSocketStream(browserWs, { allowHalfOpen: false });

  try {
    await handshakeWithBrowser(browserStream);
    // ClientInit is a single shared-desktop flag and belongs to the handshake,
    // not to the message stream. Relay it by hand: the view-only filter would
    // otherwise try to read it as a message type and lose the framing.
    const clientInit = await new ByteReader(browserStream, 20_000).read(1);
    upstream.write(clientInit);
  } catch (err) {
    upstream.destroy();
    browserStream.destroy();
    sessions.end(sessionId, { reason: `handshake: ${err.message}` });
    return;
  }

  const toClient = new Counter();      // browser -> VNC server
  const toBrowser = new Counter();     // VNC server -> browser
  const guard = control ? null : new ViewOnlyTransform();

  let ended = false;
  const finish = (reason) => {
    if (ended) return;
    ended = true;
    live.delete(sessionId);
    sessions.end(sessionId, {
      bytesIn: toClient.bytes,
      bytesOut: toBrowser.bytes,
      reason,
    });
    hub.notifySessionEnded(client.id, sessionId);
    upstream.destroy();
    browserStream.destroy();
  };

  const outbound = guard ? [browserStream, guard, toClient, upstream] : [browserStream, toClient, upstream];
  pipeline(...outbound, (err) => finish(err ? `client stream: ${err.message}` : 'closed by viewer'));
  pipeline(upstream, toBrowser, browserStream, (err) => finish(err ? `server stream: ${err.message}` : 'closed by host'));

  live.set(sessionId, {
    id: sessionId,
    clientId: client.id,
    clientName: client.name,
    username,
    role,
    source: ctx.source || 'web',
    remoteIp: ctx.remoteIp || null,
    startedAt: Date.now(),
    get bytesIn() { return toClient.bytes; },
    get bytesOut() { return toBrowser.bytes; },
    get blockedInputs() { return guard ? guard.blocked : 0; },
    kill(reason) {
      try { browserWs.close(4008, String(reason).slice(0, 120)); } catch { /* gone */ }
      finish(reason);
    },
  });

  clients.touch(client.id, ctx.remoteIp);
  audit(username, 'session.start', client.id, { sessionId, role, source: ctx.source || 'web' });
}

function listLive() {
  return Array.from(live.values()).map((s) => ({
    id: s.id,
    clientId: s.clientId,
    clientName: s.clientName,
    username: s.username,
    role: s.role,
    source: s.source,
    remoteIp: s.remoteIp,
    startedAt: s.startedAt,
    bytesIn: s.bytesIn,
    bytesOut: s.bytesOut,
    blockedInputs: s.blockedInputs,
  }));
}

function killSession(sessionId, reason = 'disconnected by an administrator') {
  const s = live.get(sessionId);
  if (!s) return false;
  s.kill(reason);
  return true;
}

function killSessionsForClient(clientId, reason) {
  let n = 0;
  for (const s of Array.from(live.values())) {
    if (s.clientId === clientId) {
      s.kill(reason);
      n++;
    }
  }
  return n;
}

module.exports = { startSession, listLive, killSession, killSessionsForClient };