/**
 * Chrome App - SOCKS5 Proxy with Bidirectional Tunneling
 * Supports both HTTP and HTTPS through true SOCKS5 tunneling
 */

// Configuration
var PROXY_HOST = '127.0.0.1';
var PROXY_PORT = 1080;

// TCP Server Management
var serverSocketId = null;
var isListening = false;

// SOCKS5 Connection States
var ConnectionState = {
  GREETING: 'greeting',
  REQUEST: 'request',
  TUNNEL: 'tunnel',  // NEW: Tunneling state
  CLOSED: 'closed'
};

// Track connections: clientSocketId -> { state, buffer, destinationSocketId }
var connections = {};

// Track reverse mapping: destinationSocketId -> clientSocketId
var reverseConnections = {};

// ============================================================================
// TCP SERVER MANAGEMENT
// ============================================================================

function createServer(host, port, callback) {
  chrome.sockets.tcpServer.create({}, function(createInfo) {
    if (chrome.runtime.lastError) {
      callback(chrome.runtime.lastError, null);
      return;
    }

    serverSocketId = createInfo.socketId;
    console.log('[TCP Server] Created server socket: ' + serverSocketId);

    chrome.sockets.tcpServer.listen(serverSocketId, host, port, function(result) {
      if (result < 0) {
        callback(new Error('Failed to listen on ' + host + ':' + port + ', code: ' + result), null);
        return;
      }

      isListening = true;
      console.log('[TCP Server] Listening on ' + host + ':' + port);
      callback(null, serverSocketId);
    });
  });
}

// ============================================================================
// SOCKS5 PROTOCOL HANDLER
// ============================================================================

// SOCKS5 constants
var SOCKS_VERSION = 0x05;
var SOCKS_CMD_CONNECT = 0x01;
var SOCKS_ATYP_IPV4 = 0x01;
var SOCKS_ATYP_DOMAIN = 0x03;
var SOCKS_ATYP_IPV6 = 0x04;
var SOCKS_AUTH_NONE = 0x00;
var SOCKS_REPLY_SUCCESS = 0x00;
var SOCKS_REPLY_GENERAL_FAILURE = 0x01;
var SOCKS_REPLY_CONNECTION_NOT_ALLOWED = 0x02;
var SOCKS_REPLY_NETWORK_UNREACHABLE = 0x03;
var SOCKS_REPLY_HOST_UNREACHABLE = 0x04;
var SOCKS_REPLY_CONNECTION_REFUSED = 0x05;

function initConnection(socketId) {
  connections[socketId] = {
    state: ConnectionState.GREETING,
    buffer: new Uint8Array(0),
    destinationSocketId: null
  };
  console.log('[SOCKS] Initialized connection ' + socketId);
}

function removeConnection(socketId) {
  var conn = connections[socketId];
  if (conn && conn.destinationSocketId) {
    delete reverseConnections[conn.destinationSocketId];
    try {
      chrome.sockets.tcp.close(conn.destinationSocketId);
    } catch(e) {}
  }
  delete connections[socketId];
  console.log('[SOCKS] Removed connection ' + socketId);
}

function handleSocksData(clientSocketId, data) {
  var conn = connections[clientSocketId];
  if (!conn) {
    console.error('[SOCKS] No connection state for ' + clientSocketId);
    return;
  }

  // If we're in tunnel mode, forward data to destination
  if (conn.state === ConnectionState.TUNNEL && conn.destinationSocketId) {
    console.log('[Tunnel] Client -> Destination: ' + data.byteLength + ' bytes');
    chrome.sockets.tcp.send(conn.destinationSocketId, data, function(sendInfo) {
      if (sendInfo.resultCode < 0) {
        console.error('[Tunnel] Failed to send to destination: ' + sendInfo.resultCode);
        closeConnection(clientSocketId);
      }
    });
    return;
  }

  // Append new data to buffer
  var newData = new Uint8Array(data);
  var newBuffer = new Uint8Array(conn.buffer.length + newData.length);
  newBuffer.set(conn.buffer);
  newBuffer.set(newData, conn.buffer.length);
  conn.buffer = newBuffer;

  // Process based on current state
  if (conn.state === ConnectionState.GREETING) {
    handleGreeting(clientSocketId, conn);
  } else if (conn.state === ConnectionState.REQUEST) {
    handleRequest(clientSocketId, conn);
  }
}

function handleGreeting(socketId, conn) {
  var data = conn.buffer;

  if (data.length < 3) {
    return; // Need more data
  }

  var version = data[0];
  var numMethods = data[1];

  if (version !== SOCKS_VERSION) {
    console.error('[SOCKS] Invalid SOCKS version: ' + version);
    closeConnection(socketId);
    return;
  }

  if (data.length < 2 + numMethods) {
    return; // Need more data
  }

  // Check if no authentication (method 0x00) is supported
  var methodSupported = false;
  for (var i = 0; i < numMethods; i++) {
    if (data[2 + i] === SOCKS_AUTH_NONE) {
      methodSupported = true;
      break;
    }
  }

  if (methodSupported) {
    sendGreetingResponse(socketId, SOCKS_AUTH_NONE);
    conn.state = ConnectionState.REQUEST;
    // Clear buffer after greeting
    conn.buffer = new Uint8Array(0);
  } else {
    closeConnection(socketId);
  }
}

function sendGreetingResponse(socketId, method) {
  var response = new Uint8Array([SOCKS_VERSION, method]);
  chrome.sockets.tcp.send(socketId, response.buffer, function(sendInfo) {
    if (sendInfo.resultCode < 0) {
      console.error('[SOCKS] Failed to send greeting response: ' + sendInfo.resultCode);
    } else {
      console.log('[SOCKS] Sent greeting response, method: ' + method);
    }
  });
}

function handleRequest(clientSocketId, conn) {
  var data = conn.buffer;

  if (data.length < 4) {
    return; // Need more data
  }

  var version = data[0];
  var command = data[1];
  var atyp = data[3];

  if (version !== SOCKS_VERSION) {
    console.error('[SOCKS] Invalid SOCKS version in request: ' + version);
    sendRequestResponse(clientSocketId, SOCKS_REPLY_GENERAL_FAILURE);
    closeConnection(clientSocketId);
    return;
  }

  if (command !== SOCKS_CMD_CONNECT) {
    console.error('[SOCKS] Unsupported command: ' + command);
    sendRequestResponse(clientSocketId, SOCKS_REPLY_COMMAND_NOT_SUPPORTED);
    closeConnection(clientSocketId);
    return;
  }

  // Parse address
  var host = '';
  var port = 0;
  var consumed = 0;

  if (atyp === SOCKS_ATYP_IPV4) {
    if (data.length < 10) {
      return; // Need more data
    }
    host = data[4] + '.' + data[5] + '.' + data[6] + '.' + data[7];
    port = (data[8] << 8) | data[9];
    consumed = 10;
  } else if (atyp === SOCKS_ATYP_DOMAIN) {
    var domainLen = data[4];
    if (data.length < 7 + domainLen) {
      return; // Need more data
    }
    host = '';
    for (var i = 0; i < domainLen; i++) {
      host += String.fromCharCode(data[5 + i]);
    }
    port = (data[5 + domainLen] << 8) | data[6 + domainLen];
    consumed = 7 + domainLen;
  } else if (atyp === SOCKS_ATYP_IPV6) {
    if (data.length < 22) {
      return; // Need more data
    }
    // Parse IPv6 (simplified)
    var parts = [];
    for (var j = 0; j < 16; j += 2) {
      parts.push(((data[4 + j] << 8) | data[5 + j]).toString(16));
    }
    host = parts.join(':');
    port = (data[20] << 8) | data[21];
    consumed = 22;
  } else {
    console.error('[SOCKS] Unsupported address type: ' + atyp);
    sendRequestResponse(clientSocketId, SOCKS_REPLY_ADDRESS_TYPE_NOT_SUPPORTED);
    closeConnection(clientSocketId);
    return;
  }

  console.log('[SOCKS] CONNECT request: ' + host + ':' + port);

  // Create tunnel to destination
  createTunnel(clientSocketId, host, port, consumed);
}

function createTunnel(clientSocketId, host, port, consumed) {
  // Clear the request buffer
  var conn = connections[clientSocketId];
  conn.buffer = new Uint8Array(0);

  // Create TCP connection to destination
  chrome.sockets.tcp.create({}, function(createInfo) {
    if (chrome.runtime.lastError || !createInfo) {
      console.error('[Tunnel] Failed to create socket:', chrome.runtime.lastError);
      sendRequestResponse(clientSocketId, SOCKS_REPLY_GENERAL_FAILURE);
      closeConnection(clientSocketId);
      return;
    }

    var destSocketId = createInfo.socketId;
    conn.destinationSocketId = destSocketId;
    reverseConnections[destSocketId] = clientSocketId;

    console.log('[Tunnel] Created destination socket: ' + destSocketId);

    // Connect to destination
    chrome.sockets.tcp.connect(destSocketId, host, port, function(result) {
      if (result < 0) {
        console.error('[Tunnel] Failed to connect to ' + host + ':' + port + ', code: ' + result);

        // Map error codes
        var reply = SOCKS_REPLY_GENERAL_FAILURE;
        if (result === -109) { // Host unreachable
          reply = SOCKS_REPLY_HOST_UNREACHABLE;
        } else if (result === -111) { // Connection refused
          reply = SOCKS_REPLY_CONNECTION_REFUSED;
        } else if (result === -102) { // Network unreachable
          reply = SOCKS_REPLY_NETWORK_UNREACHABLE;
        }

        sendRequestResponse(clientSocketId, reply);
        closeConnection(clientSocketId);
        return;
      }

      console.log('[Tunnel] Connected to ' + host + ':' + port);

      // Send success response to client
      sendRequestResponse(clientSocketId, SOCKS_REPLY_SUCCESS);

      // Set to tunnel mode
      conn.state = ConnectionState.TUNNEL;

      // Unpause destination socket to start receiving data
      chrome.sockets.tcp.setPaused(destSocketId, false);

      console.log('[Tunnel] Tunnel established: client=' + clientSocketId + ' -> dest=' + destSocketId);
    });
  });
}

function sendRequestResponse(socketId, reply) {
  var response = new Uint8Array([
    SOCKS_VERSION,      // VER
    reply,              // REP
    0x00,               // RSV
    SOCKS_ATYP_IPV4,    // ATYP
    0x00, 0x00, 0x00, 0x00, // BND.ADDR
    0x00, 0x00          // BND.PORT
  ]);

  chrome.sockets.tcp.send(socketId, response.buffer, function(sendInfo) {
    if (sendInfo.resultCode < 0) {
      console.error('[SOCKS] Failed to send request response: ' + sendInfo.resultCode);
    } else {
      console.log('[SOCKS] Sent request response, reply: ' + reply);
    }
  });
}

function closeConnection(socketId) {
  chrome.sockets.tcp.close(socketId);
  removeConnection(socketId);
}

// ============================================================================
// BIDIRECTIONAL TUNNEL DATA FORWARDING
// ============================================================================

// Handle data from destination - forward to client
function handleDestinationData(destSocketId, data) {
  var clientSocketId = reverseConnections[destSocketId];

  if (!clientSocketId) {
    console.warn('[Tunnel] No client for destination socket: ' + destSocketId);
    return;
  }

  console.log('[Tunnel] Destination -> Client: ' + data.byteLength + ' bytes');

  // Pause destination to prevent event spam while sending
  chrome.sockets.tcp.setPaused(destSocketId, true);

  chrome.sockets.tcp.send(clientSocketId, data, function(sendInfo) {
    // Resume destination after sending
    chrome.sockets.tcp.setPaused(destSocketId, false);

    if (sendInfo.resultCode < 0) {
      console.error('[Tunnel] Failed to send to client: ' + sendInfo.resultCode);
      // Close both connections
      closeConnection(clientSocketId);
    }
  });
}

// ============================================================================
// MAIN INITIALIZATION
// ============================================================================

function setupConnectionHandlers() {
  // Handle new connections from clients
  chrome.sockets.tcpServer.onAccept.addListener(function(acceptInfo) {
    console.log('[Proxy] New client connected: ' + acceptInfo.clientSocketId);

    // Initialize SOCKS state for this connection
    initConnection(acceptInfo.clientSocketId);

    // Unpause socket to start receiving data
    chrome.sockets.tcp.setPaused(acceptInfo.clientSocketId, false);
  });

  // Handle accept errors
  chrome.sockets.tcpServer.onAcceptError.addListener(function(errorInfo) {
    console.error('[Proxy] Accept error:', errorInfo);
  });

  // Handle receive errors
  chrome.sockets.tcp.onReceiveError.addListener(function(info) {
    var socketId = info.socketId;
    var resultCode = info.resultCode;
    console.error('[Proxy] Receive error on socket ' + socketId + ': ' + resultCode);

    // Clean up connections
    if (connections[socketId]) {
      closeConnection(socketId);
    } else if (reverseConnections[socketId]) {
      closeConnection(reverseConnections[socketId]);
    }
  });

  // Handle disconnects (EOF detection - zero byte receive)
  chrome.sockets.tcp.onReceive.addListener(function(receiveInfo) {
    if (receiveInfo.data.byteLength === 0) {
      var socketId = receiveInfo.socketId;
      console.log('[Proxy] Socket ' + socketId + ' disconnected (EOF)');
      if (connections[socketId]) {
        closeConnection(socketId);
      } else if (reverseConnections[socketId]) {
        closeConnection(reverseConnections[socketId]);
      }
      return;
    }

    // If not EOF, let the main onReceive handler handle the data
    var socketId = receiveInfo.socketId;

    // Check if this is data from a client or destination
    if (connections[socketId]) {
      // Data from client
      console.log('[Proxy] Received from client ' + socketId + ': ' + receiveInfo.data.byteLength + ' bytes');
      handleSocksData(socketId, receiveInfo.data);
    } else if (reverseConnections[socketId]) {
      // Data from destination
      handleDestinationData(socketId, receiveInfo.data);
    } else {
      console.warn('[Proxy] Received data from unknown socket: ' + socketId);
    }
  });
}

function initializeProxy() {
  console.log('[Proxy] Initializing Chrome SOCKS Proxy (with tunneling)...');

  createServer(PROXY_HOST, PROXY_PORT, function(error, socketId) {
    if (error) {
      console.error('[Proxy] Failed to initialize:', error);
      return;
    }

    console.log('[Proxy] SOCKS proxy is running on ' + PROXY_HOST + ':' + PROXY_PORT);
    setupConnectionHandlers();
  });
}

// Chrome App lifecycle events
chrome.app.runtime.onLaunched.addListener(function() {
  console.log('[Proxy] App launched');
  initializeProxy();
});

// Also initialize on startup (for apps that are already running)
initializeProxy();

console.log('[Proxy] Background script loaded');