#!/usr/bin/env node

/**
 * PSK HTTP/HTTPS Proxy (Client)
 *
 * - Runs a local HTTP proxy (supports CONNECT for HTTPS and absolute/origin-form HTTP for port 80)
 * - Maintains a single TLS-PSK tunnel to a remote out-node (proxy-server.js)
 * - Multiplexes many local client connections over one TLS tunnel using frames:
 *
 * Frame: [1 byte type][4 bytes connection id][4 bytes data length][data...]
 * Types:
 *   DATA (2)
 *   CLOSE (3)
 *   OPEN (4)         payload = [2B hostLen][host utf8][2B port]
 *   OPEN_RESULT (5)  payload = [1B status] 1=success,0=failure
 *
 * HTTPS: Use CONNECT method. We respond 200 after OPEN_RESULT success, then tunnel raw TLS.
 * HTTP:  Convert absolute-form to origin-form and forward over opened upstream (default port 80 if not specified).
 */

const net = require('net');
const tls = require('tls');
const fs = require('fs');
const { URL } = require('url');
const { program } = require('commander');

program
  .requiredOption('--server-host <host>', 'Proxy out-node server host')
  .requiredOption('--server-port <port>', 'Proxy out-node server port')
  .requiredOption('--psk-file <path>', 'Path to PSK key file')
  .requiredOption('--identity <identity>', 'PSK identity string')
  .requiredOption('--proxy-port <port>', 'Local HTTP proxy port to listen on')
  .option('--bind-host <host>', 'Local bind host', '127.0.0.1')
  .option('--connect-timeout <ms>', 'Timeout waiting OPEN_RESULT (ms)', '10000')
  .option('--idle-timeout <ms>', 'Idle timeout for local sockets (ms)', '60000')
  .parse();

const options = program.opts();

let pskKey;
try {
  pskKey = fs.readFileSync(options.pskFile, 'utf8').trim();
} catch (error) {
  console.error(`Error reading PSK file: ${error.message}`);
  process.exit(1);
}

const OPEN_RESULT_TIMEOUT = parseInt(options.connectTimeout, 10) || 10000;
const IDLE_TIMEOUT = parseInt(options.idleTimeout, 10) || 60000;

// Message Types
const MSG_TYPES = {
  DATA: 2,
  CLOSE: 3,
  OPEN: 4,
  OPEN_RESULT: 5,
};

function buildOpenPayload(host, port) {
  const hostBuf = Buffer.from(host, 'utf8');
  const buf = Buffer.allocUnsafe(2 + hostBuf.length + 2);
  let off = 0;
  buf.writeUInt16BE(hostBuf.length, off); off += 2;
  hostBuf.copy(buf, off); off += hostBuf.length;
  buf.writeUInt16BE(port, off); off += 2;
  return buf;
}

function writeMessage(socket, type, connectionId, data = Buffer.alloc(0)) {
  if (!socket || socket.destroyed) return;
  const header = Buffer.allocUnsafe(9);
  header.writeUInt8(type, 0);
  header.writeUInt32BE(connectionId >>> 0, 1);
  header.writeUInt32BE(data.length >>> 0, 5);
  socket.write(header);
  if (data.length > 0) socket.write(data);
}

function createMessageReader() {
  let buffer = Buffer.alloc(0);
  let expectedLength = 9;
  let currentMessage = null;

  return function onData(data, callback) {
    buffer = Buffer.concat([buffer, data]);

    while (buffer.length >= expectedLength) {
      if (!currentMessage) {
        const type = buffer.readUInt8(0);
        const connectionId = buffer.readUInt32BE(1);
        const dataLength = buffer.readUInt32BE(5);
        currentMessage = { type, connectionId, dataLength };
        expectedLength = 9 + dataLength;

        if (dataLength === 0) {
          callback(currentMessage.type, currentMessage.connectionId, Buffer.alloc(0));
          buffer = buffer.subarray(9);
          currentMessage = null;
          expectedLength = 9;
        }
      } else {
        const messageData = buffer.subarray(9, expectedLength);
        callback(currentMessage.type, currentMessage.connectionId, messageData);
        buffer = buffer.subarray(expectedLength);
        currentMessage = null;
        expectedLength = 9;
      }
    }
  };
}

function pskCallback(/* hint */) {
  return {
    identity: options.identity,
    psk: Buffer.from(pskKey, 'hex')
  };
}

// Global tunnel socket (single multiplexed connection)
let tunnelSocket = null;
let reader = null;

const connections = new Map(); // connectionId -> ConnState
let nextConnId = 1;

function genConnId() {
  let id = nextConnId >>> 0;
  do {
    id = (id + 1) >>> 0;
    if (id === 0) id = 1;
  } while (connections.has(id));
  nextConnId = id;
  return id;
}

function connectTunnel() {
  const host = options.serverHost;
  const port = parseInt(options.serverPort, 10);
  console.log(`Connecting to proxy out-node ${host}:${port} via TLS-PSK...`);

  const sock = tls.connect({
    host,
    port,
    pskCallback,
    ciphers: 'PSK-AES256-GCM-SHA384:PSK-AES128-GCM-SHA256',
    checkServerIdentity: () => undefined
  }, () => {
    console.log('Proxy tunnel connected');
  });

  sock.setNoDelay(true);
  sock.setKeepAlive(true, 30000);

  tunnelSocket = sock;
  reader = createMessageReader();

  sock.on('data', (data) => {
    reader(data, (type, connectionId, payload) => {
      const st = connections.get(connectionId);
      if (!st) {
        // Unknown or already closed
        if (type === MSG_TYPES.OPEN_RESULT) {
          // No state; ignore
        }
        return;
      }

      if (type === MSG_TYPES.OPEN_RESULT) {
        const ok = payload.length > 0 ? payload.readUInt8(0) === 1 : false;
        if (st.opened) return;
        clearTimeout(st.openTimer);
        if (!ok) {
          // Notify client error
          if (st.mode === 'CONNECT') {
            safeWrite(st.localSocket, Buffer.from(
              'HTTP/1.1 502 Bad Gateway\r\nProxy-Agent: PSK-Proxy\r\n\r\n'
            ));
          } else {
            safeWrite(st.localSocket, Buffer.from(
              'HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\nContent-Length: 11\r\nContent-Type: text/plain\r\n\r\nBad Gateway'
            ));
          }
          destroyLocal(connectionId);
          return;
        }

        st.opened = true;
        // For CONNECT: send 200 OK to client, then flush any buffered data (likely none)
        if (st.mode === 'CONNECT') {
          safeWrite(st.localSocket, Buffer.from(
            'HTTP/1.1 200 Connection Established\r\nProxy-Agent: PSK-Proxy\r\n\r\n'
          ));
          // If there was any extra data after headers (shouldn't for CONNECT), flush it
          flushBufferedToTunnel(connectionId);
        } else {
          // HTTP mode: first send the rewritten initial request, then any buffered tail
          if (st.initialUpstream && st.initialUpstream.length > 0) {
            writeMessage(tunnelSocket, MSG_TYPES.DATA, connectionId, st.initialUpstream);
            st.initialUpstream = null;
          }
          flushBufferedToTunnel(connectionId);
        }
      } else if (type === MSG_TYPES.DATA) {
        if (!st.localSocket.destroyed) {
          st.localSocket.write(payload);
        }
      } else if (type === MSG_TYPES.CLOSE) {
        // Upstream closed
        destroyLocal(connectionId);
      }
    });
  });

  sock.on('close', () => {
    console.log('Proxy tunnel closed. Cleaning up local connections and retrying in 2s...');
    // Close all local connections
    for (const [id] of connections) {
      destroyLocal(id);
    }
    tunnelSocket = null;
    setTimeout(connectTunnel, 2000);
  });

  sock.on('error', (err) => {
    console.error('Tunnel error:', err.message);
  });
}

function safeWrite(socket, buf) {
  if (!socket.destroyed) {
    try { socket.write(buf); } catch (_) {}
  }
}

function flushBufferedToTunnel(connectionId) {
  const st = connections.get(connectionId);
  if (!st || !tunnelSocket || tunnelSocket.destroyed) return;
  while (st.bufferQueue.length > 0) {
    const chunk = st.bufferQueue.shift();
    writeMessage(tunnelSocket, MSG_TYPES.DATA, connectionId, chunk);
  }
}

function destroyLocal(connectionId) {
  const st = connections.get(connectionId);
  if (!st) return;
  connections.delete(connectionId);
  clearTimeout(st.openTimer);
  try { st.localSocket.destroy(); } catch (_) {}
}

function parseHeadersUntil(buffer) {
  const idx = buffer.indexOf('\r\n\r\n');
  if (idx === -1) return null;
  const headerPart = buffer.subarray(0, idx).toString('utf8');
  const rest = buffer.subarray(idx + 4);
  const lines = headerPart.split('\r\n');
  const requestLine = lines.shift() || '';
  const headers = {};
  for (const line of lines) {
    const p = line.indexOf(':');
    if (p > -1) {
      const name = line.slice(0, p).trim();
      const val = line.slice(p + 1).trim();
      const key = name.toLowerCase();
      if (headers[key] === undefined) headers[key] = val;
      else if (Array.isArray(headers[key])) headers[key].push(val);
      else headers[key] = [headers[key], val];
    }
  }
  return { requestLine, rawHeaderLines: lines, headers, rest };
}

function buildRequestLine(method, target, version) {
  return `${method} ${target} ${version}`;
}

function startOpen(connectionId, host, port, st) {
  if (!tunnelSocket || tunnelSocket.destroyed) {
    // No tunnel available
    if (st.mode === 'CONNECT') {
      safeWrite(st.localSocket, Buffer.from(
        'HTTP/1.1 503 Service Unavailable\r\nProxy-Agent: PSK-Proxy\r\n\r\n'
      ));
    } else {
      safeWrite(st.localSocket, Buffer.from(
        'HTTP/1.1 503 Service Unavailable\r\nConnection: close\r\nContent-Length: 19\r\nContent-Type: text/plain\r\n\r\nService Unavailable'
      ));
    }
    destroyLocal(connectionId);
    return;
  }

  // Set a timer for open result
  st.openTimer = setTimeout(() => {
    if (!st.opened) {
      // Timeout
      if (st.mode === 'CONNECT') {
        safeWrite(st.localSocket, Buffer.from(
          'HTTP/1.1 504 Gateway Timeout\r\nProxy-Agent: PSK-Proxy\r\n\r\n'
        ));
      } else {
        safeWrite(st.localSocket, Buffer.from(
          'HTTP/1.1 504 Gateway Timeout\r\nConnection: close\r\nContent-Length: 15\r\nContent-Type: text/plain\r\n\r\nGateway Timeout'
        ));
      }
      // Best-effort to close upstream
      writeMessage(tunnelSocket, MSG_TYPES.CLOSE, connectionId);
      destroyLocal(connectionId);
    }
  }, OPEN_RESULT_TIMEOUT);

  writeMessage(tunnelSocket, MSG_TYPES.OPEN, connectionId, buildOpenPayload(host, port));
}

function handleLocalConnection(localSocket) {
  localSocket.setNoDelay(true);
  localSocket.setKeepAlive(true, 30000);
  if (IDLE_TIMEOUT > 0) {
    localSocket.setTimeout(IDLE_TIMEOUT, () => {
      try { localSocket.destroy(); } catch (_) {}
    });
  }

  const connectionId = genConnId();
  const state = {
    id: connectionId,
    localSocket,
    bufferQueue: [],
    opened: false,
    openTimer: null,
    mode: null,               // 'CONNECT' | 'HTTP'
    initialUpstream: null     // Buffer to send immediately after OPEN_RESULT in HTTP mode
  };
  connections.set(connectionId, state);

  let firstBuffer = Buffer.alloc(0);
  let headersParsed = false;

  function onData(chunk) {
    if (!headersParsed) {
      firstBuffer = Buffer.concat([firstBuffer, chunk]);

      // Limit header size to prevent abuse (64KB)
      if (firstBuffer.length > 64 * 1024) {
        safeWrite(localSocket, Buffer.from(
          'HTTP/1.1 431 Request Header Fields Too Large\r\nConnection: close\r\n\r\n'
        ));
        destroyLocal(connectionId);
        return;
      }

      const parsed = parseHeadersUntil(firstBuffer);
      if (!parsed) {
        // Wait for more
        return;
      }

      headersParsed = true;
      const { requestLine, rawHeaderLines, headers, rest } = parsed;

      const parts = requestLine.split(' ');
      if (parts.length < 3) {
        safeWrite(localSocket, Buffer.from(
          'HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n'
        ));
        destroyLocal(connectionId);
        return;
      }

      const method = parts[0];
      const target = parts[1];
      const version = parts[2];

      if (method.toUpperCase() === 'CONNECT') {
        state.mode = 'CONNECT';
        // target expected as host:port
        let host = target;
        let port = 443;
        const cidx = host.lastIndexOf(':');
        if (cidx !== -1) {
          port = parseInt(host.slice(cidx + 1), 10) || 443;
          host = host.slice(0, cidx);
        }
        // Initiate OPEN
        startOpen(connectionId, host, port, state);
        // Any rest after headers (rare for CONNECT) will be buffered until open
        if (rest.length > 0) {
          state.bufferQueue.push(rest);
        }
      } else {
        state.mode = 'HTTP';
        // Determine host, port, and path; rebuild request line to origin-form
        let host = null;
        let port = 80;
        let path = target;

        if (/^http:\/\//i.test(target)) {
          try {
            const u = new URL(target);
            host = u.hostname;
            port = u.port ? parseInt(u.port, 10) : 80;
            path = u.pathname + (u.search || '');
            if (path.length === 0) path = '/';
          } catch {
            safeWrite(localSocket, Buffer.from(
              'HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n'
            ));
            destroyLocal(connectionId);
            return;
          }
        } else if (/^https:\/\//i.test(target)) {
          // We do not support absolute-form HTTPS over plaintext proxy without CONNECT
          safeWrite(localSocket, Buffer.from(
            'HTTP/1.1 400 Bad Request\r\nConnection: close\r\nContent-Length: 49\r\nContent-Type: text/plain\r\n\r\nUse CONNECT method for HTTPS requests through this proxy'
          ));
          destroyLocal(connectionId);
          return;
        } else {
          // origin-form: must use Host header
          const hostHeader = headers['host'];
          if (!hostHeader) {
            safeWrite(localSocket, Buffer.from(
              'HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n'
            ));
            destroyLocal(connectionId);
            return;
          }
          const hostVal = Array.isArray(hostHeader) ? hostHeader[0] : hostHeader;
          const cidx = hostVal.lastIndexOf(':');
          if (cidx !== -1) {
            host = hostVal.slice(0, cidx);
            port = parseInt(hostVal.slice(cidx + 1), 10) || 80;
          } else {
            host = hostVal;
            port = 80;
          }
          // path is already origin-form target
        }

        // Rebuild request line to origin-form
        const newRequestLine = buildRequestLine(method, path, version);
        const headerStr = [newRequestLine, ...rawHeaderLines].join('\r\n') + '\r\n\r\n';
        state.initialUpstream = Buffer.concat([Buffer.from(headerStr, 'utf8'), rest]);

        // Initiate OPEN then send initialUpstream upon success
        startOpen(connectionId, host, port, state);
      }
      return;
    }

    // After headers parsed: if not yet opened, buffer; else forward
    if (!state.opened) {
      state.bufferQueue.push(chunk);
    } else {
      writeMessage(tunnelSocket, MSG_TYPES.DATA, connectionId, chunk);
    }
  }

  localSocket.on('data', onData);

  localSocket.on('close', () => {
    // Inform remote
    writeMessage(tunnelSocket, MSG_TYPES.CLOSE, connectionId);
    connections.delete(connectionId);
    clearTimeout(state.openTimer);
  });

  localSocket.on('error', () => {
    writeMessage(tunnelSocket, MSG_TYPES.CLOSE, connectionId);
    connections.delete(connectionId);
    clearTimeout(state.openTimer);
  });
}

// Start local HTTP proxy server
const proxyServer = net.createServer((socket) => {
  handleLocalConnection(socket);
});

proxyServer.on('error', (err) => {
  console.error('Local proxy server error:', err.message);
  process.exit(1);
});

proxyServer.listen(parseInt(options.proxyPort, 10), options.bindHost, () => {
  console.log(`Local HTTP proxy listening on ${options.bindHost}:${options.proxyPort}`);
});

// Connect tunnel and maintain it
connectTunnel();

// Graceful shutdown
process.on('SIGINT', () => {
  console.log('Shutting down...');
  try { proxyServer.close(); } catch (_) {}
  try { if (tunnelSocket) tunnelSocket.destroy(); } catch (_) {}
  for (const [id] of connections) {
    destroyLocal(id);
  }
  process.exit(0);
});