import { randomBytes } from 'node:crypto';
import { lookup } from 'node:dns/promises';
import http from 'node:http';
import net, { BlockList } from 'node:net';
import tls from 'node:tls';
import { WebSocket, WebSocketServer } from 'ws';

const PORT = numberFromEnv('WS_PORT', 3001, 1, 65535);
const MAX_SESSIONS_PER_IP = numberFromEnv('PROXY_MAX_SESSIONS_PER_IP', 3, 1, 20);
const MAX_GLOBAL_SESSIONS = numberFromEnv('PROXY_MAX_GLOBAL_SESSIONS', 200, 1, 10_000);
const NEW_CONNECTION_LIMIT = numberFromEnv('PROXY_NEW_CONNECTIONS_PER_10_MIN', 20, 1, 1_000);
const CONNECT_TIMEOUT_MS = 10_000;
const PERSISTENCE_TIMEOUT_MS = 5 * 60_000;
const MAX_BUFFER_MESSAGES = 250;
const MAX_BUFFER_BYTES = 256 * 1024;
const MAX_CONTROL_BYTES = 4 * 1024;
const MAX_BYTES_PER_SECOND = 2 * 1024 * 1024;
const DENIED_PORTS = new Set([25, 465, 587, 2525]);
const production = process.env.NODE_ENV === 'production';
const allowedOrigins = new Set(
  (process.env.ALLOWED_ORIGINS || (production ? '' : 'http://localhost:5173,http://127.0.0.1:5173'))
    .split(',').map((origin) => origin.trim()).filter(Boolean)
);
const deniedAddresses = createDeniedAddressList();
const sessions = new Map();
const connectionAttempts = new Map();

const server = http.createServer((request, response) => {
  if (request.url === '/health') {
    response.writeHead(200, { 'content-type': 'application/json' });
    response.end(JSON.stringify({ ok: true, sessions: sessions.size }));
    return;
  }
  response.writeHead(404);
  response.end();
});
const wss = new WebSocketServer({ noServer: true, maxPayload: 64 * 1024, perMessageDeflate: false });

function numberFromEnv(name, fallback, minimum, maximum) {
  const parsed = Number.parseInt(process.env[name] || '', 10);
  return Number.isFinite(parsed) && parsed >= minimum && parsed <= maximum ? parsed : fallback;
}

function createDeniedAddressList() {
  const list = new BlockList();
  for (const [address, prefix] of [
    ['0.0.0.0', 8], ['10.0.0.0', 8], ['100.64.0.0', 10], ['127.0.0.0', 8],
    ['169.254.0.0', 16], ['172.16.0.0', 12], ['192.0.0.0', 24], ['192.0.2.0', 24],
    ['192.168.0.0', 16], ['198.18.0.0', 15], ['198.51.100.0', 24], ['203.0.113.0', 24],
    ['224.0.0.0', 4], ['240.0.0.0', 4]
  ]) list.addSubnet(address, prefix, 'ipv4');
  for (const [address, prefix] of [
    ['::', 128], ['::1', 128], ['64:ff9b::', 96], ['2001::', 32],
    ['2001:db8::', 32], ['2002::', 16], ['fc00::', 7], ['fe80::', 10], ['ff00::', 8]
  ]) list.addSubnet(address, prefix, 'ipv6');
  return list;
}

function getClientIp(request) {
  if (process.env.TRUST_PROXY === '1') {
    const forwarded = request.headers['x-forwarded-for'];
    if (typeof forwarded === 'string' && forwarded.length > 0) return forwarded.split(',')[0].trim();
  }
  return request.socket.remoteAddress || 'unknown';
}

function sendControl(ws, payload) {
  if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(payload));
}

function failUpgrade(socket, status, message) {
  socket.write(`HTTP/1.1 ${status}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${message}`);
  socket.destroy();
}

function countSessionsForIp(clientIp) {
  let count = 0;
  for (const session of sessions.values()) if (session.clientIp === clientIp) count += 1;
  return count;
}

function recordConnectionAttempt(clientIp) {
  const cutoff = Date.now() - 10 * 60_000;
  const recent = (connectionAttempts.get(clientIp) || []).filter((timestamp) => timestamp > cutoff);
  if (recent.length >= NEW_CONNECTION_LIMIT) return false;
  recent.push(Date.now());
  connectionAttempts.set(clientIp, recent);
  return true;
}

function validateConnectMessage(value) {
  if (!value || typeof value !== 'object' || value.type !== 'connect') throw new Error('The first frame must be a connect control message.');
  if (typeof value.host !== 'string' || value.host.length < 1 || value.host.length > 253) throw new Error('Invalid host.');
  if (!/^[a-zA-Z0-9._:-]+$/.test(value.host)) throw new Error('Invalid host characters.');
  if (!Number.isInteger(value.port) || value.port < 1 || value.port > 65535 || DENIED_PORTS.has(value.port)) throw new Error('Invalid or denied port.');
  if (typeof value.tls !== 'boolean') throw new Error('Invalid TLS setting.');
  if (value.resumeToken !== undefined && (typeof value.resumeToken !== 'string' || value.resumeToken.length > 128)) throw new Error('Invalid resume token.');
  return { host: value.host, port: value.port, useTls: value.tls, resumeToken: value.resumeToken };
}

async function resolvePublicTarget(host) {
  const directFamily = net.isIP(host);
  const answers = directFamily ? [{ address: host, family: directFamily }] : await lookup(host, { all: true, verbatim: true });
  if (answers.length === 0) throw new Error('Host did not resolve.');
  for (const answer of answers) {
    if (isDeniedAddress(answer.address, answer.family)) throw new Error('Target resolves to a non-public address.');
  }
  return answers[0];
}

function isDeniedAddress(address, family = net.isIP(address)) {
  const mapped = address.match(/^::ffff:(?:(\d+\.\d+\.\d+\.\d+)|([0-9a-f]+):([0-9a-f]+))$/i);
  if (mapped) {
    let ipv4 = mapped[1];
    if (!ipv4) {
      const high = Number.parseInt(mapped[2], 16);
      const low = Number.parseInt(mapped[3], 16);
      ipv4 = `${high >> 8}.${high & 255}.${low >> 8}.${low & 255}`;
    }
    return deniedAddresses.check(ipv4, 'ipv4');
  }
  return deniedAddresses.check(address, family === 4 ? 'ipv4' : 'ipv6');
}

function consumeTraffic(session, direction, bytes) {
  if (Date.now() - session.traffic.startedAt >= 1_000) session.traffic = { startedAt: Date.now(), inbound: 0, outbound: 0 };
  session.traffic[direction] += bytes;
  return session.traffic[direction] <= MAX_BYTES_PER_SECOND;
}

function bufferMessage(session, data) {
  session.buffer.push(Buffer.from(data));
  session.bufferBytes += data.length;
  while (session.buffer.length > MAX_BUFFER_MESSAGES || session.bufferBytes > MAX_BUFFER_BYTES) {
    const removed = session.buffer.shift();
    if (removed) session.bufferBytes -= removed.length;
  }
}

function destroySession(session, reason = 'closed') {
  if (session.closed) return;
  session.closed = true;
  if (session.persistenceTimer) clearTimeout(session.persistenceTimer);
  sessions.delete(session.token);
  const attached = session.ws;
  session.ws = null;
  if (attached) sendControl(attached, { type: 'upstream_closed', reason });
  if (!session.socket.destroyed) session.socket.destroy();
}

function rotateSessionToken(session) {
  sessions.delete(session.token);
  session.token = randomBytes(32).toString('base64url');
  sessions.set(session.token, session);
}

function attachWebSocket(session, ws, resumed) {
  if (session.persistenceTimer) clearTimeout(session.persistenceTimer);
  session.persistenceTimer = null;
  session.ws = ws;
  rotateSessionToken(session);
  if (resumed) {
    sendControl(ws, { type: 'replay_started', messageCount: session.buffer.length });
    for (const message of session.buffer) if (ws.readyState === WebSocket.OPEN) ws.send(message);
    const messagesReplayed = session.buffer.length;
    session.buffer = [];
    session.bufferBytes = 0;
    sendControl(ws, { type: 'replay_finished', messagesReplayed });
  }
  sendControl(ws, { type: resumed ? 'session_resumed' : 'session_started', resumeToken: session.token });
}

async function createSession(ws, request, target) {
  const clientIp = getClientIp(request);
  const origin = request.headers.origin || '';
  if (sessions.size >= MAX_GLOBAL_SESSIONS) throw new Error('Proxy capacity reached.');
  if (countSessionsForIp(clientIp) >= MAX_SESSIONS_PER_IP) throw new Error('Per-client connection limit reached.');
  if (!recordConnectionAttempt(clientIp)) throw new Error('Connection rate limit reached.');
  const resolved = await resolvePublicTarget(target.host);
  const socketOptions = { host: resolved.address, port: target.port, family: resolved.family,
    ...(target.useTls ? { servername: net.isIP(target.host) ? undefined : target.host, rejectUnauthorized: true } : {}) };
  const socket = target.useTls ? tls.connect(socketOptions) : net.createConnection(socketOptions);
  const session = { token: randomBytes(32).toString('base64url'), clientIp, origin, target, socket, ws: null,
    buffer: [], bufferBytes: 0, persistenceTimer: null, closed: false,
    traffic: { startedAt: Date.now(), inbound: 0, outbound: 0 } };
  sessions.set(session.token, session);
  const connected = new Promise((resolve, reject) => {
    const connectedEvent = target.useTls ? 'secureConnect' : 'connect';
    const fail = (error) => reject(error instanceof Error ? error : new Error('Upstream connection failed.'));
    socket.once(connectedEvent, resolve);
    socket.once('error', fail);
    socket.once('close', () => fail(new Error('Upstream closed before connecting.')));
    socket.setTimeout(CONNECT_TIMEOUT_MS, () => {
      fail(new Error('Upstream connection timed out.'));
      destroySession(session, 'timeout');
    });
  });
  socket.on('data', (data) => {
    if (!consumeTraffic(session, 'inbound', data.length)) return destroySession(session, 'traffic_limit');
    if (session.ws?.readyState === WebSocket.OPEN) session.ws.send(data);
    else bufferMessage(session, data);
  });
  socket.on('error', (error) => { if (session.ws) sendControl(session.ws, { type: 'error', code: 'UPSTREAM_ERROR', message: error.message }); });
  socket.on('close', () => destroySession(session, 'upstream_closed'));
  socket.on('drain', () => session.ws?.resume());
  try {
    await connected;
  } catch (error) {
    destroySession(session, 'connect_failed');
    throw error;
  }
  socket.setTimeout(0);
  if (ws.readyState !== WebSocket.OPEN) {
    destroySession(session, 'client_closed_during_connect');
    throw new Error('Client closed while the upstream connection was opening.');
  }
  attachWebSocket(session, ws, false);
  return session;
}

function tryResume(ws, request, target) {
  if (!target.resumeToken) return null;
  const session = sessions.get(target.resumeToken);
  const clientIp = getClientIp(request);
  const origin = request.headers.origin || '';
  if (!session || session.ws || session.clientIp !== clientIp || session.origin !== origin || session.target.host !== target.host ||
      session.target.port !== target.port || session.target.useTls !== target.useTls) return null;
  attachWebSocket(session, ws, true);
  return session;
}

wss.on('connection', (ws, request) => {
  let session = null;
  let explicitDisconnect = false;
  let initialized = false;
  const initializationTimer = setTimeout(() => { if (!initialized) ws.close(1008, 'Connect control timeout'); }, 5_000);
  ws.on('message', async (data, isBinary) => {
    try {
      if (!initialized) {
        if (isBinary || data.length > MAX_CONTROL_BYTES) throw new Error('Invalid connect control frame.');
        const target = validateConnectMessage(JSON.parse(data.toString('utf8')));
        initialized = true;
        clearTimeout(initializationTimer);
        session = tryResume(ws, request, target) || await createSession(ws, request, target);
        return;
      }
      if (!isBinary) {
        if (data.length > MAX_CONTROL_BYTES) throw new Error('Control frame too large.');
        const control = JSON.parse(data.toString('utf8'));
        if (control.type === 'disconnect') {
          explicitDisconnect = true;
          if (session) destroySession(session, 'client_disconnect');
          ws.close(1000, 'Disconnected');
          return;
        }
        throw new Error('Unknown control frame.');
      }
      if (!session || session.ws !== ws || !session.socket.writable) throw new Error('Upstream connection is not writable.');
      if (!consumeTraffic(session, 'outbound', data.length)) return destroySession(session, 'traffic_limit');
      if (!session.socket.write(data)) ws.pause();
    } catch (error) {
      sendControl(ws, { type: 'error', code: initialized ? 'PROTOCOL_ERROR' : 'CONNECT_ERROR', message: error instanceof Error ? error.message : 'Unknown proxy error.' });
      if (!session) ws.close(1008, 'Connection rejected');
    }
  });
  ws.on('close', () => {
    clearTimeout(initializationTimer);
    if (!session || explicitDisconnect || session.ws !== ws) return;
    session.ws = null;
    session.persistenceTimer = setTimeout(() => destroySession(session, 'resume_timeout'), PERSISTENCE_TIMEOUT_MS);
  });
  ws.on('error', () => ws.close());
});

server.on('upgrade', (request, socket, head) => {
  let pathname;
  try { pathname = new URL(request.url || '/', 'http://localhost').pathname; } catch { return failUpgrade(socket, '400 Bad Request', 'Invalid URL'); }
  if (pathname !== '/mud-ws') return failUpgrade(socket, '404 Not Found', 'Not found');
  const origin = request.headers.origin;
  if (typeof origin !== 'string' || !allowedOrigins.has(origin)) return failUpgrade(socket, '403 Forbidden', 'Origin not allowed');
  wss.handleUpgrade(request, socket, head, (client) => wss.emit('connection', client, request));
});

function shutdown() {
  wss.close();
  for (const session of [...sessions.values()]) destroySession(session, 'server_shutdown');
  server.close(() => process.exit(0));
  setTimeout(() => process.exit(1), 5_000).unref();
}
process.on('SIGINT', shutdown);
process.on('SIGTERM', shutdown);
if (process.env.NODE_ENV !== 'test') server.listen(PORT, () => console.log(`MUD WebSocket proxy listening on port ${PORT}`));
export { server, validateConnectMessage, resolvePublicTarget, isDeniedAddress };