/**
 * SOCKS5 proxy server — built on Node.js net module, no external deps.
 *
 * Implements:
 *   - SOCKS5 handshake (RFC 1928)
 *   - No-auth method (0x00)
 *   - CONNECT command (TCP tunnel)
 *   - Byte tracking per connection
 */
import * as net from "net";
import { BandwidthTracker } from "./bandwidth";

const SOCKS_VER = 0x05;
const NO_AUTH = 0x00;
const CMD_CONNECT = 0x01;
const ATYP_IPV4 = 0x01;
const ATYP_DOMAIN = 0x03;
const ATYP_IPV6 = 0x04;
const REP_SUCCESS = 0x00;
const REP_FAIL = 0x01;
const CONN_TIMEOUT_MS = 30_000;

export interface Socks5ServerOptions {
  port: number;
  host?: string;
  onNewConnection?: () => boolean; // return false to reject (bandwidth limit)
  bandwidth?: BandwidthTracker;
}

function buildReply(rep: number, ip = "0.0.0.0", port = 0): Buffer {
  const parts = ip.split(".").map(Number);
  const buf = Buffer.alloc(10);
  buf[0] = SOCKS_VER;
  buf[1] = rep;
  buf[2] = 0x00; // reserved
  buf[3] = ATYP_IPV4;
  buf[4] = parts[0] ?? 0;
  buf[5] = parts[1] ?? 0;
  buf[6] = parts[2] ?? 0;
  buf[7] = parts[3] ?? 0;
  buf.writeUInt16BE(port, 8);
  return buf;
}

function handleClient(
  socket: net.Socket,
  opts: Socks5ServerOptions
): void {
  socket.setTimeout(CONN_TIMEOUT_MS);
  socket.on("timeout", () => socket.destroy());
  socket.on("error", () => socket.destroy());

  let state: "greeting" | "request" | "connected" = "greeting";
  let buf = Buffer.alloc(0);

  socket.on("data", (chunk) => {
    buf = Buffer.concat([buf, chunk]);

    if (state === "greeting") {
      // Need at least: VER(1) + NMETHODS(1) + methods
      if (buf.length < 2) return;
      const nmethods = buf[1] ?? 0;
      if (buf.length < 2 + nmethods) return;

      const methods = buf.slice(2, 2 + nmethods);
      buf = buf.slice(2 + nmethods);

      if (!methods.includes(NO_AUTH)) {
        socket.write(Buffer.from([SOCKS_VER, 0xFF])); // no acceptable method
        socket.destroy();
        return;
      }

      socket.write(Buffer.from([SOCKS_VER, NO_AUTH]));
      state = "request";
      return;
    }

    if (state === "request") {
      // Minimum: VER(1) CMD(1) RSV(1) ATYP(1) + address + port(2)
      if (buf.length < 4) return;

      const cmd = buf[1];
      const atyp = buf[3];

      let host: string;
      let port: number;
      let consumed: number;

      if (atyp === ATYP_IPV4) {
        if (buf.length < 10) return;
        host = `${buf[4]}.${buf[5]}.${buf[6]}.${buf[7]}`;
        port = buf.readUInt16BE(8);
        consumed = 10;
      } else if (atyp === ATYP_DOMAIN) {
        if (buf.length < 5) return;
        const domainLen = buf[4] ?? 0;
        if (buf.length < 5 + domainLen + 2) return;
        host = buf.slice(5, 5 + domainLen).toString("utf8");
        port = buf.readUInt16BE(5 + domainLen);
        consumed = 5 + domainLen + 2;
      } else if (atyp === ATYP_IPV6) {
        if (buf.length < 22) return;
        const parts: string[] = [];
        for (let i = 0; i < 8; i++) {
          parts.push(buf.readUInt16BE(4 + i * 2).toString(16));
        }
        host = parts.join(":");
        port = buf.readUInt16BE(20);
        consumed = 22;
      } else {
        socket.write(buildReply(0x08)); // address type not supported
        socket.destroy();
        return;
      }

      buf = buf.slice(consumed);

      if (cmd !== CMD_CONNECT) {
        socket.write(buildReply(0x07)); // command not supported
        socket.destroy();
        return;
      }

      // Check bandwidth limit before accepting
      if (opts.onNewConnection && !opts.onNewConnection()) {
        socket.write(buildReply(REP_FAIL));
        socket.destroy();
        return;
      }

      opts.bandwidth?.incrementRequests();

      const target = net.createConnection({ host, port }, () => {
        socket.write(buildReply(REP_SUCCESS));
        state = "connected";

        // Pipe remaining buffered bytes
        if (buf.length > 0) {
          target.write(buf);
          buf = Buffer.alloc(0);
        }

        // Bidirectional pipe with byte counting
        socket.on("data", (d) => {
          opts.bandwidth?.addBytes(d.length);
          if (!target.destroyed) target.write(d);
        });

        target.on("data", (d) => {
          opts.bandwidth?.addBytes(d.length);
          if (!socket.destroyed) socket.write(d);
        });

        socket.on("end", () => { if (!target.destroyed) target.end(); });
        target.on("end", () => { if (!socket.destroyed) socket.end(); });
      });

      target.setTimeout(CONN_TIMEOUT_MS);
      target.on("timeout", () => { target.destroy(); socket.destroy(); });
      target.on("error", () => {
        if (!socket.destroyed) {
          socket.write(buildReply(REP_FAIL));
          socket.destroy();
        }
      });
    }
  });
}

export function createSocks5Server(opts: Socks5ServerOptions): net.Server {
  const server = net.createServer((socket) => {
    handleClient(socket, opts);
  });

  server.on("error", (err) => {
    console.error("[SOCKS5] Server error:", err.message);
  });

  return server;
}