#include "socks_manager.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "utils.h"
#include <stdint.h>
#include <errno.h>

#ifdef _WIN32
  //#include <winsock2.h>
  //#include <ws2tcpip.h>
  #pragma comment(lib,"Ws2_32.lib")
  typedef SOCKET sock_t;
  #define CLOSESOCK(s) closesocket(s)
#else
  #include <unistd.h>
  #include <pthread.h>
  #include <sys/socket.h>
  #include <sys/types.h>
  #include <netinet/in.h>
  #include <arpa/inet.h>
  #include <netdb.h>
  typedef int sock_t;
  #define INVALID_SOCKET -1
  #define SOCKET_ERROR -1
  #define CLOSESOCK(s) close(s)
#endif


volatile int socks_proxy_active = 0; // Flag global de estado SOCKS

typedef struct conn_entry {
    uint32_t server_id;
    sock_t sock;
    struct conn_entry *next;
    int closed;
} conn_entry_t;

static conn_entry_t *conn_list = NULL;
#ifdef _WIN32
static CRITICAL_SECTION conn_lock;
#else
static pthread_mutex_t conn_lock = PTHREAD_MUTEX_INITIALIZER;
#endif

typedef struct outgoing_node {
    socks_msg_t msg;
    struct outgoing_node *next;
} outgoing_node_t;
static outgoing_node_t *out_head = NULL;
static outgoing_node_t *out_tail = NULL;

#ifdef _WIN32
static CRITICAL_SECTION out_lock;
#else
static pthread_mutex_t out_lock = PTHREAD_MUTEX_INITIALIZER;
#endif

static void lock_init_once(void) {
#ifdef _WIN32
    InitializeCriticalSection(&conn_lock);
    InitializeCriticalSection(&out_lock);
#endif
}

/* === Inicialización / limpieza === */

void socks_manager_init(void) {
_dbg("socks_manager_init()");
    lock_init_once();
#ifdef _WIN32
    WSADATA wsa;
    WSAStartup(MAKEWORD(2,2), &wsa);
#endif
}

void socks_manager_cleanup(void) {
_dbg("socks_manager_cleanup()");
#ifdef _WIN32
    EnterCriticalSection(&conn_lock);
#else
    pthread_mutex_lock(&conn_lock);
#endif
    conn_entry_t *cur = conn_list;
    while (cur) {
        if (cur->sock != INVALID_SOCKET) CLOSESOCK(cur->sock);
        conn_entry_t *n = cur->next;
        free(cur);
        cur = n;
    }
    conn_list = NULL;
#ifdef _WIN32
    LeaveCriticalSection(&conn_lock);
    DeleteCriticalSection(&conn_lock);
    DeleteCriticalSection(&out_lock);
    WSACleanup();
#else
    pthread_mutex_unlock(&conn_lock);
#endif
    socks_proxy_active = 0;
}

/* === Gestión de conexiones === */

static conn_entry_t* find_entry(uint32_t server_id) {
#ifdef _WIN32
    EnterCriticalSection(&conn_lock);
#else
    pthread_mutex_lock(&conn_lock);
#endif
    conn_entry_t *cur = conn_list;
    while (cur) {
        if (cur->server_id == server_id) {
#ifdef _WIN32
            LeaveCriticalSection(&conn_lock);
#else
            pthread_mutex_unlock(&conn_lock);
#endif
            return cur;
        }
        cur = cur->next;
    }
#ifdef _WIN32
    LeaveCriticalSection(&conn_lock);
#else
    pthread_mutex_unlock(&conn_lock);
#endif
    return NULL;
}

static conn_entry_t* create_entry(uint32_t server_id, sock_t s) {
    conn_entry_t *e = malloc(sizeof(conn_entry_t));
    if (!e) return NULL;
    e->server_id = server_id;
    e->sock = s;
    e->next = NULL;
    e->closed = 0;
#ifdef _WIN32
    EnterCriticalSection(&conn_lock);
#else
    pthread_mutex_lock(&conn_lock);
#endif
    e->next = conn_list;
    conn_list = e;
#ifdef _WIN32
    LeaveCriticalSection(&conn_lock);
#else
    pthread_mutex_unlock(&conn_lock);
#endif
_dbg("create_entry sid=%u sock=%d", server_id, (int)s);
    return e;
}

static void remove_entry(uint32_t server_id) {
#ifdef _WIN32
    EnterCriticalSection(&conn_lock);
#else
    pthread_mutex_lock(&conn_lock);
#endif
    conn_entry_t *cur = conn_list, *prev = NULL;
    while (cur) {
        if (cur->server_id == server_id) {
            if (prev) prev->next = cur->next;
            else conn_list = cur->next;
            if (cur->sock != INVALID_SOCKET) CLOSESOCK(cur->sock);
_dbg("remove_entry sid=%u", server_id);
            free(cur);
            break;
        }
        prev = cur;
        cur = cur->next;
    }
#ifdef _WIN32
    LeaveCriticalSection(&conn_lock);
#else
    pthread_mutex_unlock(&conn_lock);
#endif
}

/* === Cola de salida === */

static void enqueue_outgoing(socks_msg_t *m) {
_dbg("enqueue_outgoing: sid=%u exit=%d len=%zu", m->server_id, m->exit_flag, m->data_len);
    outgoing_node_t *n = malloc(sizeof(outgoing_node_t));
    memset(n,0,sizeof(*n));
    n->msg.server_id = m->server_id;
    n->msg.exit_flag = m->exit_flag;
    n->msg.data_len = m->data_len;
    if (m->data_len) {
        n->msg.data = malloc(m->data_len);
        memcpy(n->msg.data, m->data, m->data_len);
    } else n->msg.data = NULL;
    n->next = NULL;
#ifdef _WIN32
    EnterCriticalSection(&out_lock);
#else
    pthread_mutex_lock(&out_lock);
#endif
    if (!out_tail) out_head = out_tail = n;
    else { out_tail->next = n; out_tail = n; }
#ifdef _WIN32
    LeaveCriticalSection(&out_lock);
#else
    pthread_mutex_unlock(&out_lock);
#endif
_dbg("Mensaje encolado OK (sid=%u)", m->server_id);
}

/* === Volcado de cola de salida === */
void socks_manager_flush_outgoing(socks_send_cb_t cb, void *ctx) {
_dbg("flush_outgoing: iniciando");
    if (!cb) {
        _err("ERROR: callback es NULL");
        return;
    }
#ifdef _WIN32
    EnterCriticalSection(&out_lock);
#else
    pthread_mutex_lock(&out_lock);
#endif
    outgoing_node_t *n = out_head;
    out_head = out_tail = NULL;
#ifdef _WIN32
    LeaveCriticalSection(&out_lock);
#else
    pthread_mutex_unlock(&out_lock);
#endif

    int count = 0;
    while (n) {
        count++;
_dbg("flush_outgoing: procesando mensaje #%d (sid=%u exit=%d len=%zu)", 
             count, n->msg.server_id, n->msg.exit_flag, n->msg.data_len);
        outgoing_node_t *next = n->next;
        cb(n->msg.server_id, n->msg.exit_flag, n->msg.data, n->msg.data_len, ctx);
        if (n->msg.data) free(n->msg.data);
        free(n);
        n = next;
    }
_dbg("flush_outgoing: completado (%d mensajes procesados)", count);
}

/* === Base64 === */
extern int base64_decode(const char *b64, unsigned char **out, size_t *out_len);
extern int base64_encode(const unsigned char *data, size_t dlen, char **out_b64);

/* === Parseo SOCKS5 === */
static int parse_socks5_connect(const unsigned char *data, size_t data_len, char *out_ip, size_t ip_sz, uint16_t *out_port) {
    if (data_len < 7) return -1;
    if (data[0] == 0x05) {
        if (data_len >= 10 && data[1] == 0x01 && data[2] == 0x00) {
            unsigned char atyp = data[3];
            if (atyp == 0x01 && data_len >= 10) {
                snprintf(out_ip, ip_sz, "%u.%u.%u.%u", data[4], data[5], data[6], data[7]);
                *out_port = (uint16_t)((data[8] << 8) | data[9]);
                return 0;
            } else if (atyp == 0x03) {
                unsigned int len = data[4];
                if (data_len >= 7 + len) {
                    memcpy(out_ip, &data[5], len);
                    out_ip[len] = 0;
                    *out_port = (uint16_t)((data[5+len] << 8) | data[6+len]);
                    return 0;
                }
            }
        }
    }
    return -1;
}

/* === Thread lector === */
#ifdef _WIN32
DWORD WINAPI remote_reader_thread(LPVOID arg);
#else
static void* remote_reader_thread(void* arg);
#endif

typedef struct reader_args {
    uint32_t server_id;
    sock_t s;
} reader_args_t;

#ifdef _WIN32
DWORD WINAPI remote_reader_thread(LPVOID arg) {
    reader_args_t *ra = (reader_args_t*)arg;
    uint32_t sid = ra->server_id;
    sock_t s = ra->s;
    free(ra);
_dbg("Thread iniciado para sid=%u socket=%d", sid, (int)s);
    unsigned char buf[4096];
    int loop_count = 0;
    while (1) {
        loop_count++;
_dbg("sid=%u loop #%d, llamando recv()...", sid, loop_count);
        int r = recv(s, (char*)buf, sizeof(buf), 0);
_dbg("sid=%u recv() retornó r=%d", sid, r);
        if (r < 0) {
            int err = WSAGetLastError();
_dbg(" WSAError=%d", err);
        }
        if (r <= 0) break;
_dbg("sid=%u recibió %d bytes, encolando...", sid, r);
        socks_msg_t m = {sid, 0, buf, (size_t)r};
        enqueue_outgoing(&m);
_dbg("sid=%u datos encolados OK", sid);
    }
_dbg("sid=%u salió del loop, cerrando socket", sid);
_dbg("reader sid=%u closing", sid);
    socks_msg_t m2 = {sid, 1, NULL, 0};
    enqueue_outgoing(&m2);
    CLOSESOCK(s);
_dbg("Thread terminado para sid=%u", sid);
    return 0;
}
#else
static void* remote_reader_thread(void* arg) {
    reader_args_t *ra = (reader_args_t*)arg;
    uint32_t sid = ra->server_id;
    sock_t s = ra->s;
    free(ra);
    unsigned char buf[4096];
    while (1) {
        ssize_t r = recv(s, buf, sizeof(buf), 0);
        if (r <= 0) break;
_dbg("reader sid=%u recv=%zd", sid, r);
        socks_msg_t m = {sid, 0, buf, (size_t)r};
        enqueue_outgoing(&m);
    }
_dbg("reader sid=%u closing", sid);
    socks_msg_t m2 = {sid, 1, NULL, 0};
    enqueue_outgoing(&m2);
    CLOSESOCK(s);
    return NULL;
}
#endif

/* === Conexión remota === */
typedef struct {
    uint32_t ipv4;   // en network byte order
    uint16_t port;   // en network byte order
} socks_bound_addr_t;

static int connect_and_spawn_reader(const char *dest, uint16_t port, uint32_t server_id, socks_bound_addr_t *bound) {
_dbg("connect_and_spawn_reader: dest=%s port=%u sid=%u", dest, port, server_id);
    if (server_id == 0) {
        _err("ERROR: invalid server_id=%u", server_id);
_dbg("invalid server_id=%u (ignored)", server_id);
        return -4;
    }
_dbg("connect_and_spawn_reader dest=%s port=%u sid=%u", dest, port, server_id);
    struct sockaddr_in sa;
    sock_t s = socket(AF_INET, SOCK_STREAM, 0);
    if (s == INVALID_SOCKET) {
        _err("ERROR: socket() failed");
        return -1;
    }
_dbg("Socket creado OK");
    memset(&sa,0,sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port = htons(port);
    
    /* Intentar parsear como dirección IP usando inet_addr (compatible con MinGW) */
    unsigned long addr = inet_addr(dest);
    if (addr == INADDR_NONE) {
        /* No es una IP válida, intentar resolver como hostname */
_dbg("Resolviendo hostname: %s", dest);
        struct hostent *h = gethostbyname(dest);
        if (!h) {
            _err("ERROR: gethostbyname failed");
            CLOSESOCK(s);
            return -2;
        }
        sa.sin_addr = *(struct in_addr*)h->h_addr;
_dbg("Hostname resuelto OK");
    } else {
        /* Es una IP válida */
        sa.sin_addr.s_addr = addr;
    }
_dbg("Conectando a %s:%u...", dest, port);
    if (connect(s, (struct sockaddr*)&sa, sizeof(sa)) != 0) {
        _err("ERROR: connect() failed errno=%d", errno);
_dbg("connect() failed errno=%d", errno);
        CLOSESOCK(s);
        return -3;
    }
_dbg("Conexión exitosa!");

    /* Guardar la dirección bound (la IP/puerto del destino conectado) */
    if (bound) {
        bound->ipv4 = sa.sin_addr.s_addr;  // ya en network byte order
        bound->port = sa.sin_port;         // ya en network byte order
_dbg("Bound addr: %08X:%u", ntohl(bound->ipv4), ntohs(bound->port));
    }

    /* timeout para evitar bloqueo indefinido */
    int timeout_ms = 15000; // 15s
    setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout_ms, sizeof(timeout_ms));
_dbg("socket timeout set %d ms", timeout_ms);

_dbg("Creando entry para sid=%u", server_id);
    create_entry(server_id, s);
#ifdef _WIN32
    reader_args_t *ra = malloc(sizeof(reader_args_t));
    ra->server_id = server_id; ra->s = s;
_dbg("Creando thread lector para sid=%u", server_id);
    CreateThread(NULL,0, remote_reader_thread, ra, 0, NULL);
#else
    pthread_t tid;
    reader_args_t *ra = malloc(sizeof(reader_args_t));
    ra->server_id = server_id; ra->s = s;
    pthread_create(&tid, NULL, remote_reader_thread, ra);
    pthread_detach(tid);
#endif
_dbg("connect_and_spawn_reader completado OK para sid=%u", server_id);
    return 0;
}

/* === Procesamiento principal === */
int socks_manager_process_incoming(uint32_t server_id, int exit_flag, const unsigned char *data, size_t data_len) {
_dbg("process_incoming: sid=%u exit=%d len=%zu", server_id, exit_flag, data_len);
_dbg("incoming sid=%u exit=%d len=%zu", server_id, exit_flag, data_len);
    if (server_id == 0) {
        _err("ERROR: invalid server_id=%u", server_id);
_dbg("invalid server_id=%u", server_id);
        return -4;
    }

    if (exit_flag) {
_dbg("Exit flag set, removiendo entry sid=%u", server_id);
        remove_entry(server_id);
        socks_msg_t m = {server_id, 1, NULL, 0};
        enqueue_outgoing(&m);
        return 0;
    }

    conn_entry_t *e = find_entry(server_id);
    if (!e) {
        char dest[256]; uint16_t dport=0;
        int r = parse_socks5_connect(data, data_len, dest, sizeof(dest), &dport);
_dbg("Nueva conexión: sid=%u parse_rc=%d dest=%s port=%u", server_id, r, dest, dport);
_dbg("sid=%u new conn parse_rc=%d dest=%s port=%u", server_id, r, dest, dport);
        if (r != 0) return -1;
        
        socks_bound_addr_t bound;
        memset(&bound, 0, sizeof(bound));
        int rc = connect_and_spawn_reader(dest, dport, server_id, &bound);
_dbg("connect_and_spawn_reader rc=%d", rc);
_dbg("connect_and_spawn_reader rc=%d", rc);
        if (rc != 0) {
            _err("ERROR: connect_and_spawn_reader falló con rc=%d", rc);
            return -2;
        }
        
        /* Construir respuesta SOCKS5 con bound address real */
        unsigned char reply[10];
        reply[0] = 0x05;  // SOCKS5
        reply[1] = 0x00;  // Success
        reply[2] = 0x00;  // Reserved
        reply[3] = 0x01;  // IPv4
        /* Copiar IP (4 bytes en network byte order) */
        memcpy(&reply[4], &bound.ipv4, 4);
        /* Copiar puerto (2 bytes en network byte order) */
        memcpy(&reply[8], &bound.port, 2);
        
_dbg("Enviando respuesta SOCKS (10 bytes) a sid=%u con bound %u.%u.%u.%u:%u", 
             server_id, reply[4], reply[5], reply[6], reply[7], ntohs(bound.port));
        
        socks_msg_t m = {server_id, 0, reply, sizeof(reply)};
        enqueue_outgoing(&m);
_dbg("Respuesta SOCKS encolada para sid=%u", server_id);
        return 0;
    } else {
_dbg("Entry encontrada para sid=%u, enviando %zu bytes al socket remoto", server_id, data_len);
        if (e->sock != INVALID_SOCKET) {
            size_t total = 0;
            while (total < data_len) {
_dbg("sid=%u enviando %zu bytes (total=%zu de %zu)", server_id, data_len - total, total, data_len);
                int w = send(e->sock, (const char*)(data + total), (int)(data_len - total), 0);
_dbg("sid=%u send() retornó w=%d", server_id, w);
                if (w <= 0) {
                    _err("send() falló para sid=%u, cerrando conexión", server_id);
_dbg("send error sid=%u w=%d", server_id, w);
                    remove_entry(server_id);
                    socks_msg_t m = {server_id, 1, NULL, 0};
                    enqueue_outgoing(&m);
                    return -1;
                }
                total += w;
            }
_dbg("sid=%u enviados %zu bytes OK", server_id, data_len);
_dbg("sent %zu bytes sid=%u", data_len, server_id);
            return 0;
        } else {
            _err("socket inválido para sid=%u", server_id);
            return -2;
        }
    }
}

/* === Stubs start/stop === */
int socks_manager_start_proxy(uint16_t local_port) {
    socks_proxy_active = 1;
_dbg("socks_manager_start_proxy port=%u (active)", local_port);
    (void)local_port;
    return 0;
}

int socks_manager_stop_proxy(void) {
_dbg("socks_manager_stop_proxy()");
    socks_proxy_active = 0;
#ifdef _WIN32
    EnterCriticalSection(&conn_lock);
#else
    pthread_mutex_lock(&conn_lock);
#endif
    conn_entry_t *cur = conn_list;
    while (cur) {
        if (cur->sock != INVALID_SOCKET) CLOSESOCK(cur->sock);
        cur = cur->next;
    }
#ifdef _WIN32
    LeaveCriticalSection(&conn_lock);
#else
    pthread_mutex_unlock(&conn_lock);
#endif
    return 0;
}
