#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <process.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>

#include <stdint.h>

typedef struct {
    uint8_t version;
    uint8_t num_methods;
} socks5_client_hello_t;

typedef struct {
    uint8_t version;
    uint8_t method;
} socks5_server_hello_t;

typedef struct {
    uint8_t version;
    uint8_t command;
    uint8_t reserved;
    uint8_t addr_type;
} socks5_request_t;

typedef struct {
    uint8_t version;
    uint8_t reply;
    uint8_t reserved;
    uint8_t addr_type;
} socks5_reply_t;

// Version
#define SOCKS5_VERSION 0x05

// Authentication
#define SOCKS5_AUTH_NO_AUTH 0x00
#define SOCKS5_AUTH_NOT_ACCEPT 0xff

// Command
#define SOCKS5_CMD_CONNECT 0x01
#define SOCKS5_CMD_BIND 0x02
#define SOCKS5_CMD_ASSOCIATE 0x03

// Address type
#define SOCKS5_ATYP_IPV4 0x01
#define SOCKS5_ATYP_DOMAIN_NAME 0x03
#define SOCKS5_ATYP_IPV6 0x04

// Reply
#define SOCKS5_REP_SUCCESS 0x00
#define SOCKS5_REP_GENERAL_FAILURE 0x01
#define SOCKS5_REP_CONNECTION_NOT_ALLOWED 0x02
#define SOCKS5_REP_NET_UNREACHABLE 0x03
#define SOCKS5_REP_HOST_UNREACHABLE 0x04
#define SOCKS5_REP_CONNECTION_REFUSED 0x05
#define SOCKS5_REP_TTL_EXPIRED 0x06
#define SOCKS5_REP_COMMAND_NOT_SUPPORTED 0x07
#define SOCKS5_REP_ADDR_TYPE_NOT_SUPPORTED 0x08

// Global variables for server control
SOCKET server_fd = INVALID_SOCKET;
HANDLE hServerThread = NULL;
volatile bool server_running = false;

// Forward declarations
unsigned server_loop(void* arg);
unsigned client_handler_func(void* arg);

// Start SOCKS5 server on specified port
InitSocks5_Server(char* result, int port) {

    // Check if already running
    if (server_running) {
        LOG("[PROXY] Server already running\n");
        sprintf(result, "Server already running\n");
        return;
    }

    server_running = true;
    
    // Pass port to thread
    int* port_arg = malloc(sizeof(int));
    *port_arg = port;
    
    // Start server thread
    hServerThread = (HANDLE)_beginthread(server_loop, 0, port_arg);
    if (hServerThread == NULL) {
        LOG_ERROR("[PROXY] Failed to create server thread\n");
        sprintf(result, "Failed to create server thread\n");
        free(port_arg);
        server_running = false;
        // WSACleanup();
        return;
    }
    sprintf(result, "Started proxy server\n");
    return;
}

// Stop SOCKS5 server
int StopSocks5_Server(char* result) {
    if (!server_running) {
        LOG("[PROXY] Server not running\n");
        sprintf(result, "Proxy server not running\n");
        return;
    }
    
    // Signal server to stop
    server_running = false;
    
    // Close server socket to break out of accept()
    if (server_fd != INVALID_SOCKET) {
        closesocket(server_fd);
        server_fd = INVALID_SOCKET;
    }
    
    // Wait for thread to finish
    if (hServerThread != NULL) {
        WaitForSingleObject(hServerThread, 5000); // Wait up to 5 seconds
        CloseHandle(hServerThread);
        hServerThread = NULL;
    }
    sprintf(result, "Proxy server stopped\n");
    // WSACleanup();
    return;
}

// Server main loop
unsigned server_loop(void* arg) {
    int port = *(int*)arg;
    free(arg);
    
    // Create socket
    server_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (server_fd == INVALID_SOCKET) {
        LOG_ERROR("socket failed: %d\n", WSAGetLastError());
        server_running = false;
        return 1;
    }
    
    // Set socket to reuse address
    BOOL opt = TRUE;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, (char*)&opt, sizeof(opt));
    
    // Bind to port
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons((u_short)port);
    
    if (bind(server_fd, (struct sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR) {
        LOG_ERROR("bind failed: %d\n", WSAGetLastError());
        closesocket(server_fd);
        server_running = false;
        return 1;
    }
    
    // Listen for connections
    if (listen(server_fd, SOMAXCONN) == SOCKET_ERROR) {
        LOG_ERROR("listen failed: %d\n", WSAGetLastError());
        closesocket(server_fd);
        server_running = false;
        return 1;
    }
    
    LOG("[PROXY] Server started on port %d\n", port);
    
    // Accept connections
    while (server_running) {
        fd_set readfds;
        FD_ZERO(&readfds);
        FD_SET(server_fd, &readfds);
        
        // Use select with timeout to check server_running flag
        struct timeval tv = {1, 0}; // 1 second timeout
        int ready = select(0, &readfds, NULL, NULL, &tv);
        
        if (ready == SOCKET_ERROR || !server_running) {
            break;
        }
        
        if (ready == 0) { // Timeout
            continue;
        }
        
        // Accept connection
        struct sockaddr_in client_addr;
        int addr_len = sizeof(client_addr);
        SOCKET client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &addr_len);
        
        if (client_fd == INVALID_SOCKET) {
            continue;
        }
        
        // Handle client in new thread
        SOCKET* client_sock = malloc(sizeof(SOCKET));
        *client_sock = client_fd;
        _beginthread(client_handler_func, 0, client_sock);
    }
    
    LOG("[PROXY] Server stopped\n");
    return 0;
}

// Handle SOCKS5 client connection
unsigned client_handler_func(void* arg) {
    SOCKET client_fd = *(SOCKET*)arg;
    free(arg);
    
    SOCKET remote_fd = INVALID_SOCKET;
    
    // Handle SOCKS5 handshake
    
    // 1. Authentication method negotiation
    char buffer[512];
    if (recv(client_fd, buffer, 2, 0) != 2) {
        goto cleanup;
    }
    
    // Check SOCKS version
    if (buffer[0] != 5) {
        goto cleanup;
    }
    
    // Get number of auth methods
    int nmethods = buffer[1];
    if (recv(client_fd, buffer, nmethods, 0) != nmethods) {
        goto cleanup;
    }
    
    // We only support no auth (0x00)
    char auth_reply[2] = {5, 0}; // VER=5, METHOD=NO_AUTH
    send(client_fd, auth_reply, 2, 0);
    
    // 2. Connection request
    if (recv(client_fd, buffer, 4, 0) != 4) {
        goto cleanup;
    }
    
    // Check SOCKS version and command
    if (buffer[0] != 5 || buffer[1] != 1) { // VER=5, CMD=CONNECT
        goto cleanup;
    }
    
    // Get address type
    int atyp = buffer[3];
    
    struct sockaddr_in remote_addr;
    memset(&remote_addr, 0, sizeof(remote_addr));
    remote_addr.sin_family = AF_INET;
    
    if (atyp == 1) { // IPv4
        // Get IPv4 address
        if (recv(client_fd, buffer, 4, 0) != 4) {
            goto cleanup;
        }
        memcpy(&remote_addr.sin_addr, buffer, 4);
        
        // Get port
        if (recv(client_fd, buffer, 2, 0) != 2) {
            goto cleanup;
        }
        memcpy(&remote_addr.sin_port, buffer, 2);
    }
    else if (atyp == 3) { // Domain name
        // Get domain length
        if (recv(client_fd, buffer, 1, 0) != 1) {
            goto cleanup;
        }
        int domain_len = buffer[0];
        
        // Get domain
        if (recv(client_fd, buffer, domain_len, 0) != domain_len) {
            goto cleanup;
        }
        buffer[domain_len] = '\0';
        
        // Get port
        if (recv(client_fd, buffer + domain_len, 2, 0) != 2) {
            goto cleanup;
        }
        
        // Resolve domain
        struct addrinfo *result, hints;
        memset(&hints, 0, sizeof(hints));
        hints.ai_family = AF_INET;
        hints.ai_socktype = SOCK_STREAM;
        
        char port_str[6];
        u_short port;
        memcpy(&port, buffer + domain_len, 2);
        sprintf(port_str, "%d", ntohs(port));
        
        if (getaddrinfo(buffer, port_str, &hints, &result) != 0) {
            // Failed to resolve
            char reply[10] = {5, 4, 0, 1, 0, 0, 0, 0, 0, 0}; // Host not found
            send(client_fd, reply, 10, 0);
            goto cleanup;
        }
        
        memcpy(&remote_addr, result->ai_addr, result->ai_addrlen);
        freeaddrinfo(result);
    }
    else {
        // Unsupported address type
        goto cleanup;
    }
    
    // Connect to remote host
    remote_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (remote_fd == INVALID_SOCKET) {
        goto cleanup;
    }
    
    if (connect(remote_fd, (struct sockaddr*)&remote_addr, sizeof(remote_addr)) == SOCKET_ERROR) {
        // Connection failed
        char reply[10] = {5, 5, 0, 1, 0, 0, 0, 0, 0, 0}; // Connection refused
        send(client_fd, reply, 10, 0);
        goto cleanup;
    }
    
    // Send success response
    char reply[10] = {5, 0, 0, 1, 0, 0, 0, 0, 0, 0}; // Success
    send(client_fd, reply, 10, 0);
    
    // 3. Relay data between client and remote host
    while (server_running) {
        fd_set readfds;
        FD_ZERO(&readfds);
        FD_SET(client_fd, &readfds);
        FD_SET(remote_fd, &readfds);
        
        struct timeval tv = {1, 0}; // 1 second timeout
        int ready = select(0, &readfds, NULL, NULL, &tv);
        
        if (ready <= 0) {
            if (ready == 0) continue; // Timeout, try again
            break; // Error
        }
        
        // Check if client sent data
        if (FD_ISSET(client_fd, &readfds)) {
            int bytes = recv(client_fd, buffer, sizeof(buffer), 0);
            if (bytes <= 0) break;
            if (send(remote_fd, buffer, bytes, 0) <= 0) break;
        }
        
        // Check if remote host sent data
        if (FD_ISSET(remote_fd, &readfds)) {
            int bytes = recv(remote_fd, buffer, sizeof(buffer), 0);
            if (bytes <= 0) break;
            if (send(client_fd, buffer, bytes, 0) <= 0) break;
        }
    }

cleanup:
    if (remote_fd != INVALID_SOCKET) closesocket(remote_fd);
    closesocket(client_fd);
    return 0;
}