const net = require('net');
const http = require('http');

const listenPort = Number(process.env.SOCKS5_PORT || 1080);
const metricsPort = Number(process.env.SOCKS5_METRICS_PORT || 1081);
const stats = { accepted: 0, connected: 0, failed: 0, active: 0 };

function reply(socket, code, addressType = 1) {
  const address = addressType === 4 ? Buffer.alloc(16) : Buffer.alloc(4);
  socket.write(Buffer.concat([
    Buffer.from([5, code, 0, addressType]),
    address,
    Buffer.from([0, 0]),
  ]));
}

function consumeRequest(socket, state) {
  if (state.phase === 'greeting') {
    if (state.buffer.length < 2) return;
    const methodsLength = state.buffer[1];
    if (state.buffer.length < 2 + methodsLength) return;
    const methods = state.buffer.subarray(2, 2 + methodsLength);
    state.buffer = state.buffer.subarray(2 + methodsLength);
    if (state.buffer[0] !== undefined && state.buffer[0] !== 5) {
      socket.destroy();
      return;
    }
    if (!methods.includes(0)) {
      socket.end(Buffer.from([5, 255]));
      return;
    }
    socket.write(Buffer.from([5, 0]));
    state.phase = 'request';
  }

  if (state.phase !== 'request' || state.buffer.length < 4) return;
  const [version, command, , addressType] = state.buffer;
  if (version !== 5 || command !== 1) {
    reply(socket, command === 1 ? 1 : 7);
    socket.end();
    return;
  }

  let offset = 4;
  let host;
  if (addressType === 1) {
    if (state.buffer.length < offset + 4 + 2) return;
    host = Array.from(state.buffer.subarray(offset, offset + 4)).join('.');
    offset += 4;
  } else if (addressType === 3) {
    if (state.buffer.length < offset + 1) return;
    const length = state.buffer[offset++];
    if (state.buffer.length < offset + length + 2) return;
    host = state.buffer.subarray(offset, offset + length).toString('utf8');
    offset += length;
  } else if (addressType === 4) {
    if (state.buffer.length < offset + 16 + 2) return;
    const words = [];
    for (let index = 0; index < 16; index += 2) {
      words.push(state.buffer.readUInt16BE(offset + index).toString(16));
    }
    host = words.join(':');
    offset += 16;
  } else {
    reply(socket, 8);
    socket.end();
    return;
  }

  const port = state.buffer.readUInt16BE(offset);
  state.buffer = state.buffer.subarray(offset + 2);
  state.phase = 'connecting';

  const upstream = net.createConnection({ host, port });
  upstream.setTimeout(30000);
  upstream.once('connect', () => {
    stats.connected += 1;
    stats.active += 1;
    console.log(JSON.stringify({ event: 'connect', host, port, addressType }));
    reply(socket, 0, addressType === 4 ? 4 : 1);
    if (state.buffer.length) upstream.write(state.buffer);
    socket.pipe(upstream);
    upstream.pipe(socket);
  });
  upstream.once('error', (error) => {
    stats.failed += 1;
    console.error(JSON.stringify({ event: 'connect_error', host, port, code: error.code }));
    if (!socket.destroyed) {
      reply(socket, error.code === 'ECONNREFUSED' ? 5 : 4);
      socket.end();
    }
  });
  upstream.once('close', () => {
    if (stats.active > 0) stats.active -= 1;
  });
  socket.once('close', () => upstream.destroy());
}

const server = net.createServer((socket) => {
  stats.accepted += 1;
  const state = { phase: 'greeting', buffer: Buffer.alloc(0) };
  socket.setTimeout(30000);
  socket.on('data', (chunk) => {
    if (state.phase === 'connecting') return;
    state.buffer = Buffer.concat([state.buffer, chunk]);
    consumeRequest(socket, state);
  });
  socket.on('timeout', () => socket.destroy());
  socket.on('error', () => {});
});

server.listen(listenPort, '0.0.0.0', () => {
  console.log(`SOCKS5 test proxy listening on ${listenPort}`);
});

http.createServer((request, response) => {
  if (request.url !== '/stats') {
    response.writeHead(404).end();
    return;
  }
  response.writeHead(200, { 'content-type': 'application/json' });
  response.end(JSON.stringify(stats));
}).listen(metricsPort, '0.0.0.0');