/*
 * WSSSH Library - Tunnel management implementation
 *
 * Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "wssshlib.h"
#include "tunnel.h"
#include "websocket.h"
#include "wssh_ssl.h"
#include "control_messages.h"
#include "data_messages.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <fcntl.h>
#include <errno.h>
#include <sched.h>
#include <pthread.h>
#include <sys/wait.h>
#include <ctype.h>

#define INITIAL_FRAME_BUFFER_SIZE 8192

// Global variables
tunnel_t *active_tunnel = NULL;  // For backward compatibility
tunnel_t **active_tunnels = NULL;
int active_tunnels_count = 0;
int active_tunnels_capacity = 0;
pthread_mutex_t tunnel_mutex = PTHREAD_MUTEX_INITIALIZER;

// Global shutdown flag for threads
volatile int global_shutdown = 0;

// CPU affinity management
static int num_cpu_cores = 0;
static int next_cpu_core = 0;
static pthread_mutex_t cpu_affinity_mutex = PTHREAD_MUTEX_INITIALIZER;

// Initialize CPU core detection
void init_cpu_affinity() {
    num_cpu_cores = sysconf(_SC_NPROCESSORS_ONLN);
    if (num_cpu_cores <= 0) {
        num_cpu_cores = 1;  // Fallback to 1 core
    }
    next_cpu_core = 0;
}

// Set CPU affinity for a thread to distribute across cores
void set_thread_cpu_affinity(pthread_t thread) {
    if (num_cpu_cores <= 1) {
        return;  // No need to set affinity for single core
    }

    pthread_mutex_lock(&cpu_affinity_mutex);
    int target_core = next_cpu_core;
    next_cpu_core = (next_cpu_core + 1) % num_cpu_cores;
    pthread_mutex_unlock(&cpu_affinity_mutex);

    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(target_core, &cpuset);

    if (pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuset) != 0) {
        // Silently ignore affinity setting failures - not critical
        // fprintf(stderr, "Warning: Failed to set CPU affinity for thread\n");
    }
}

frame_buffer_t *frame_buffer_init() {
    frame_buffer_t *fb = malloc(sizeof(frame_buffer_t));
    if (!fb) return NULL;

    fb->buffer = malloc(INITIAL_FRAME_BUFFER_SIZE);
    if (!fb->buffer) {
        free(fb);
        return NULL;
    }

    fb->size = INITIAL_FRAME_BUFFER_SIZE;
    fb->used = 0;
    return fb;
}

void frame_buffer_free(frame_buffer_t *fb) {
    if (fb) {
        free(fb->buffer);
        free(fb);
    }
}

int frame_buffer_resize(frame_buffer_t *fb, size_t new_size) {
    if (new_size <= fb->size) return 1;

    char *new_buffer = realloc(fb->buffer, new_size);
    if (!new_buffer) return 0;

    fb->buffer = new_buffer;
    fb->size = new_size;
    return 1;
}

int frame_buffer_append(frame_buffer_t *fb, const char *data, size_t len) {
    if (fb->used + len > fb->size) {
        size_t new_size = fb->size * 2;
        while (new_size < fb->used + len) {
            new_size *= 2;
        }
        if (!frame_buffer_resize(fb, new_size)) {
            return 0;
        }
    }

    memcpy(fb->buffer + fb->used, data, len);
    fb->used += len;
    return 1;
}

int frame_buffer_consume(frame_buffer_t *fb, size_t len) {
    if (len > fb->used) return 0;

    memmove(fb->buffer, fb->buffer + len, fb->used - len);
    fb->used -= len;
    return 1;
}

// Helper functions for managing multiple tunnels
tunnel_t *find_tunnel_by_request_id(const char *request_id) {
    for (int i = 0; i < active_tunnels_count; i++) {
        if (active_tunnels[i] && strcmp(active_tunnels[i]->request_id, request_id) == 0) {
            return active_tunnels[i];
        }
    }
    return NULL;
}

int add_tunnel(tunnel_t *tunnel) {
    if (active_tunnels_count >= active_tunnels_capacity) {
        int new_capacity = active_tunnels_capacity == 0 ? 4 : active_tunnels_capacity * 2;
        tunnel_t **new_tunnels = realloc(active_tunnels, new_capacity * sizeof(tunnel_t *));
        if (!new_tunnels) return 0;
        active_tunnels = new_tunnels;
        active_tunnels_capacity = new_capacity;
    }
    active_tunnels[active_tunnels_count++] = tunnel;
    return 1;
}

void remove_tunnel(const char *request_id) {
    for (int i = 0; i < active_tunnels_count; i++) {
        if (active_tunnels[i] && strcmp(active_tunnels[i]->request_id, request_id) == 0) {
            // Free tunnel resources
            if (active_tunnels[i]->local_sock >= 0) {
                close(active_tunnels[i]->local_sock);
            }
            if (active_tunnels[i]->sock >= 0) {
                close(active_tunnels[i]->sock);
            }
            // Don't free SSL here - it's managed at the connection level
            active_tunnels[i]->ssl = NULL;
            if (active_tunnels[i]->outgoing_buffer) {
                frame_buffer_free(active_tunnels[i]->outgoing_buffer);
            }
            if (active_tunnels[i]->incoming_buffer) {
                frame_buffer_free(active_tunnels[i]->incoming_buffer);
            }
            free(active_tunnels[i]);

            // Shift remaining tunnels
            for (int j = i; j < active_tunnels_count - 1; j++) {
                active_tunnels[j] = active_tunnels[j + 1];
            }
            active_tunnels_count--;
            break;
        }
    }
}

void handle_tunnel_request(SSL *ssl, const char *request_id, int debug, const char *ssh_host, int ssh_port) {
    char actual_request_id[37];
    strcpy(actual_request_id, request_id);
    pthread_mutex_lock(&tunnel_mutex);

    // Check if tunnel with this request_id already exists
    tunnel_t *existing_tunnel = find_tunnel_by_request_id(actual_request_id);
    if (existing_tunnel) {
        if (existing_tunnel->active) {
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s already exists and is active, ignoring duplicate request\n", actual_request_id);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        } else {
            // Tunnel exists but is inactive - generate a new request_id for the new tunnel
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s exists but is inactive, will create new tunnel with new request_id\n", actual_request_id);
            }
            // Generate a new unique request_id for this new tunnel request
            generate_request_id(actual_request_id, sizeof(actual_request_id));
            if (debug) {
                printf("[DEBUG - Tunnel] Generated new request_id %s for tunnel recreation\n", actual_request_id);
            }
        }
    }

    tunnel_t *new_tunnel = malloc(sizeof(tunnel_t));
    if (!new_tunnel) {
        perror("Memory allocation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    memset(new_tunnel, 0, sizeof(tunnel_t));

    // For wssshc: Connect to target TCP endpoint and forward raw TCP data
    struct sockaddr_in target_addr;
    int target_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (target_sock < 0) {
        perror("Target socket creation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    memset(&target_addr, 0, sizeof(target_addr));
    target_addr.sin_family = AF_INET;
    target_addr.sin_port = htons(ssh_port); // Target port

    // Resolve target host
    struct hostent *target_he;
    if ((target_he = gethostbyname(ssh_host)) == NULL) {
        herror("Target host resolution failed");
        close(target_sock);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    target_addr.sin_addr = *((struct in_addr *)target_he->h_addr); // Target host

    if (connect(target_sock, (struct sockaddr *)&target_addr, sizeof(target_addr)) < 0) {
        perror("Connection to target endpoint failed");
        close(target_sock);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    new_tunnel->sock = target_sock; // TCP connection to target
    new_tunnel->local_sock = -1; // Not used in wssshc
    strcpy(new_tunnel->request_id, actual_request_id);
    new_tunnel->active = 1;
    new_tunnel->broken = 0;
    new_tunnel->ssl = ssl;
    new_tunnel->outgoing_buffer = NULL; // wssshc doesn't use buffer
    new_tunnel->incoming_buffer = NULL; // wssshc doesn't need incoming buffer
    new_tunnel->server_version_sent = 0; // Not used for raw TCP

    // Initialize keep-alive statistics
    time_t current_time = time(NULL);
    new_tunnel->last_keepalive_sent = current_time;
    new_tunnel->last_keepalive_received = current_time;
    new_tunnel->total_bytes_sent = 0;
    new_tunnel->total_bytes_received = 0;
    new_tunnel->bytes_last_period = 0;
    new_tunnel->last_stats_reset = current_time;

    // Add the new tunnel to the array
    if (!add_tunnel(new_tunnel)) {
        if (target_sock >= 0) close(target_sock);
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    pthread_mutex_unlock(&tunnel_mutex);

    if (debug) {
        printf("[DEBUG - Tunnel] wssshc connected to target %s:%d\n", ssh_host, ssh_port);
        fflush(stdout);
    }

    // Send tunnel_ack back to server
    char ack_msg[256];
    snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_ack\",\"request_id\":\"%s\"}", actual_request_id);

    if (debug) {
        printf("[DEBUG - WebSockets] Sending tunnel_ack: %s\n", ack_msg);
        fflush(stdout);
    }

    // send_websocket_frame already uses SSL mutex internally
    if (!send_websocket_frame(ssl, ack_msg)) {
        fprintf(stderr, "Send tunnel_ack failed\n");
        return;
    }

    // Start bidirectional forwarding between WebSocket and target TCP endpoint
    thread_args_t *thread_args = malloc(sizeof(thread_args_t));
    if (thread_args) {
        thread_args->ssl = ssl;
        thread_args->tunnel = new_tunnel;
        thread_args->debug = debug;

        pthread_t thread;
        pthread_create(&thread, NULL, forward_ws_to_ssh_server, thread_args);
        set_thread_cpu_affinity(thread);  // Distribute thread across CPU cores
        pthread_detach(thread);
    }
}

// Execute a command before tunnel connection
static int execute_service_command(const char *command, int debug) {
    if (!command || strlen(command) == 0) return 0;

    if (debug) {
        printf("[DEBUG] Executing service command: %s\n", command);
    }

    pid_t pid = fork();
    if (pid == 0) {
        // Child process
        execl("/bin/sh", "sh", "-c", command, NULL);
        _exit(1); // exec failed
    } else if (pid > 0) {
        // Parent process
        int status;
        waitpid(pid, &status, 0);
        if (debug) {
            printf("[DEBUG] Service command exited with status: %d\n", WEXITSTATUS(status));
        }
        return WEXITSTATUS(status) == 0 ? 0 : -1;
    } else {
        // Fork failed
        perror("Failed to fork for service command");
        return -1;
    }
}

void handle_tunnel_request_with_service(SSL *ssl, const char *request_id, service_config_t *service, int debug) {
    handle_tunnel_request_with_service_and_enc(ssl, request_id, service, debug, ENCODING_HEX);
}

void handle_tunnel_request_with_enc(SSL *ssl, const char *request_id, int debug, const char *ssh_host, int ssh_port, wsssh_encoding_t encoding) {
    char actual_request_id[37];
    strcpy(actual_request_id, request_id);
    pthread_mutex_lock(&tunnel_mutex);

    // Check if tunnel with this request_id already exists
    tunnel_t *existing_tunnel = find_tunnel_by_request_id(actual_request_id);
    if (existing_tunnel) {
        if (existing_tunnel->active) {
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s already exists and is active, ignoring duplicate request\n", actual_request_id);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        } else {
            // Tunnel exists but is inactive - generate a new request_id for the new tunnel
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s exists but is inactive, will create new tunnel with new request_id\n", actual_request_id);
            }
            // Generate a new unique request_id for this new tunnel request
            generate_request_id(actual_request_id, sizeof(actual_request_id));
            if (debug) {
                printf("[DEBUG - Tunnel] Generated new request_id %s for tunnel recreation\n", actual_request_id);
            }
        }
    }

    tunnel_t *new_tunnel = malloc(sizeof(tunnel_t));
    if (!new_tunnel) {
        perror("Memory allocation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    memset(new_tunnel, 0, sizeof(tunnel_t));

    // For wssshc: Connect to target TCP endpoint and forward raw TCP data
    struct sockaddr_in target_addr;
    int target_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (target_sock < 0) {
        perror("Target socket creation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    memset(&target_addr, 0, sizeof(target_addr));
    target_addr.sin_family = AF_INET;
    target_addr.sin_port = htons(ssh_port); // Target port

    // Resolve target host
    struct hostent *target_he;
    if ((target_he = gethostbyname(ssh_host)) == NULL) {
        herror("Target host resolution failed");
        close(target_sock);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    target_addr.sin_addr = *((struct in_addr *)target_he->h_addr); // Target host

    if (connect(target_sock, (struct sockaddr *)&target_addr, sizeof(target_addr)) < 0) {
        perror("Connection to target endpoint failed");
        close(target_sock);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    new_tunnel->sock = target_sock; // TCP connection to target
    new_tunnel->local_sock = -1; // Not used in wssshc
    strcpy(new_tunnel->request_id, actual_request_id);
    new_tunnel->active = 1;
    new_tunnel->broken = 0;
    new_tunnel->ssl = ssl;
    new_tunnel->outgoing_buffer = NULL; // wssshc doesn't use buffer
    new_tunnel->incoming_buffer = NULL; // wssshc doesn't need incoming buffer
    new_tunnel->server_version_sent = 0; // Not used for raw TCP
    new_tunnel->bin = (encoding == ENCODING_BINARY); // Set binary mode based on encoding
    new_tunnel->encoding = encoding;

    // Initialize retransmission buffer for reliable transmission
    new_tunnel->retransmission_buffer = retransmission_buffer_init();
    if (!new_tunnel->retransmission_buffer) {
        perror("Failed to initialize retransmission buffer");
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    // Initialize keep-alive statistics
    time_t current_time = time(NULL);
    new_tunnel->last_keepalive_sent = current_time;
    new_tunnel->last_keepalive_received = current_time;
    new_tunnel->total_bytes_sent = 0;
    new_tunnel->total_bytes_received = 0;
    new_tunnel->bytes_last_period = 0;
    new_tunnel->last_stats_reset = current_time;

    // Add the new tunnel to the array
    if (!add_tunnel(new_tunnel)) {
        if (target_sock >= 0) close(target_sock);
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    pthread_mutex_unlock(&tunnel_mutex);

    if (debug) {
        const char *enc_str = "hex";
        if (encoding == ENCODING_BASE64) enc_str = "base64";
        else if (encoding == ENCODING_BINARY) enc_str = "bin";
        printf("[DEBUG - Tunnel] wssshc connected to target %s:%d (encoding: %s)\n", ssh_host, ssh_port, enc_str);
        fflush(stdout);
    }

    // Send tunnel_ack back to server
    char ack_msg[256];
    snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_ack\",\"request_id\":\"%s\"}", actual_request_id);

    if (debug) {
        printf("[DEBUG - WebSockets] Sending tunnel_ack: %s\n", ack_msg);
        fflush(stdout);
    }

    // send_websocket_frame already uses SSL mutex internally
    if (!send_websocket_frame(ssl, ack_msg)) {
        fprintf(stderr, "Send tunnel_ack failed\n");
        return;
    }

    // Start bidirectional forwarding between WebSocket and target TCP endpoint
    thread_args_t *thread_args = malloc(sizeof(thread_args_t));
    if (thread_args) {
        thread_args->ssl = ssl;
        thread_args->tunnel = new_tunnel;
        thread_args->debug = debug;

        pthread_t thread;
        pthread_create(&thread, NULL, forward_ws_to_ssh_server, thread_args);
        set_thread_cpu_affinity(thread);  // Distribute thread across CPU cores
        pthread_detach(thread);
    }
}

void handle_tunnel_request_with_service_and_enc(SSL *ssl, const char *request_id, service_config_t *service, int debug, wsssh_encoding_t encoding) {
    char actual_request_id[37];
    strcpy(actual_request_id, request_id);
    pthread_mutex_lock(&tunnel_mutex);

    // Check if tunnel with this request_id already exists
    tunnel_t *existing_tunnel = find_tunnel_by_request_id(actual_request_id);
    if (existing_tunnel) {
        if (existing_tunnel->active) {
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s already exists and is active, ignoring duplicate request\n", actual_request_id);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        } else {
            // Tunnel exists but is inactive - generate a new request_id for the new tunnel
            if (debug) {
                printf("[DEBUG - Tunnel] Tunnel with request_id %s exists but is inactive, will create new tunnel with new request_id\n", actual_request_id);
            }
            // Generate a new unique request_id for this new tunnel request
            generate_request_id(actual_request_id, sizeof(actual_request_id));
            if (debug) {
                printf("[DEBUG - Tunnel] Generated new request_id %s for tunnel recreation\n", actual_request_id);
            }
        }
    }

    // Execute service command if specified
    if (service && service->command) {
        if (execute_service_command(service->command, debug) != 0) {
            if (debug) {
                printf("[DEBUG - Tunnel] Service command failed, aborting tunnel request\n");
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
    }

    tunnel_t *new_tunnel = malloc(sizeof(tunnel_t));
    if (!new_tunnel) {
        perror("Memory allocation failed");
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    memset(new_tunnel, 0, sizeof(tunnel_t));

    // For wssshc: Connect to target TCP endpoint and forward raw TCP data
    struct sockaddr_in target_addr;
    int target_sock;

    // Determine protocol (TCP or UDP)
    int use_udp = 0;
    if (service && service->proto && strcmp(service->proto, "udp") == 0) {
        use_udp = 1;
        target_sock = socket(AF_INET, SOCK_DGRAM, 0);
        if (debug) {
            printf("[DEBUG - Tunnel] Using UDP protocol for service\n");
        }
    } else {
        target_sock = socket(AF_INET, SOCK_STREAM, 0);
    }

    if (target_sock < 0) {
        perror("Target socket creation failed");
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    memset(&target_addr, 0, sizeof(target_addr));
    target_addr.sin_family = AF_INET;

    // Use service configuration if available
    const char *target_host = service && service->tunnel_host ? service->tunnel_host : "127.0.0.1";
    int target_port = service && service->tunnel_port ? service->tunnel_port : 22;

    target_addr.sin_port = htons(target_port); // Target port

    // Resolve target host
    struct hostent *target_he;
    if ((target_he = gethostbyname(target_host)) == NULL) {
        herror("Target host resolution failed");
        close(target_sock);
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }
    target_addr.sin_addr = *((struct in_addr *)target_he->h_addr); // Target host

    if (!use_udp) {
        // TCP connection
        if (connect(target_sock, (struct sockaddr *)&target_addr, sizeof(target_addr)) < 0) {
            perror("Connection to target endpoint failed");
            close(target_sock);
            free(new_tunnel);
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
    } else {
        // For UDP, we don't connect, just set the target address for sending
        // The connection will be established when we receive data
    }

    new_tunnel->sock = target_sock; // TCP/UDP connection to target
    new_tunnel->local_sock = -1; // Not used in wssshc
    strcpy(new_tunnel->request_id, actual_request_id);
    new_tunnel->active = 1;
    new_tunnel->broken = 0;
    new_tunnel->ssl = ssl;
    new_tunnel->outgoing_buffer = NULL; // wssshc doesn't use buffer
    new_tunnel->incoming_buffer = NULL; // wssshc doesn't need incoming buffer
    new_tunnel->server_version_sent = 0; // Not used for raw TCP/UDP
    new_tunnel->bin = (encoding == ENCODING_BINARY); // Set binary mode based on encoding
    new_tunnel->encoding = encoding;

    // Initialize retransmission buffer for reliable transmission
    new_tunnel->retransmission_buffer = retransmission_buffer_init();
    if (!new_tunnel->retransmission_buffer) {
        perror("Failed to initialize retransmission buffer");
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    // Initialize keep-alive statistics
    time_t current_time = time(NULL);
    new_tunnel->last_keepalive_sent = current_time;
    new_tunnel->last_keepalive_received = current_time;
    new_tunnel->total_bytes_sent = 0;
    new_tunnel->total_bytes_received = 0;
    new_tunnel->bytes_last_period = 0;
    new_tunnel->last_stats_reset = current_time;
    new_tunnel->start_time = current_time;
    strcpy(new_tunnel->service, service ? service->name : "ssh");
    strcpy(new_tunnel->client_id, "local");

    // Add the new tunnel to the array
    if (!add_tunnel(new_tunnel)) {
        if (target_sock >= 0) close(target_sock);
        free(new_tunnel);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    pthread_mutex_unlock(&tunnel_mutex);

    if (debug) {
        const char *enc_str = "hex";
        if (encoding == ENCODING_BASE64) enc_str = "base64";
        else if (encoding == ENCODING_BINARY) enc_str = "bin";
        printf("[DEBUG - Tunnel] wssshc connected to target %s:%d (%s, encoding: %s)\n",
               target_host, target_port, use_udp ? "UDP" : "TCP", enc_str);
    }

    // Send tunnel_ack back to server
    char ack_msg[256];
    snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_ack\",\"request_id\":\"%s\"}", actual_request_id);

    if (debug) {
        printf("[DEBUG - WebSockets] Sending tunnel_ack: %s\n", ack_msg);
        fflush(stdout);
    }

    // send_websocket_frame already uses SSL mutex internally
    if (!send_websocket_frame(ssl, ack_msg)) {
        fprintf(stderr, "Send tunnel_ack failed\n");
        return;
    }

    // Start bidirectional forwarding between WebSocket and target TCP/UDP endpoint
    thread_args_t *thread_args = malloc(sizeof(thread_args_t));
    if (thread_args) {
        thread_args->ssl = ssl;
        thread_args->tunnel = new_tunnel;
        thread_args->debug = debug;

        pthread_t thread;
        if (use_udp) {
            // For UDP, we need a different forwarding function
            // For now, use the TCP version but this should be enhanced for UDP
            pthread_create(&thread, NULL, forward_ws_to_ssh_server, thread_args);
        } else {
            pthread_create(&thread, NULL, forward_ws_to_ssh_server, thread_args);
        }
        set_thread_cpu_affinity(thread);  // Distribute thread across CPU cores
        pthread_detach(thread);
    }
}

void cleanup_tunnel(int debug) {
    pthread_mutex_lock(&tunnel_mutex);

    // Set global shutdown flag to signal all threads to stop
    global_shutdown = 1;

    // First, mark all tunnels as inactive to signal threads to stop
    for (int i = 0; i < active_tunnels_count; i++) {
        if (active_tunnels[i]) {
            active_tunnels[i]->active = 0;
        }
    }

    // Give threads a moment to stop using resources
    pthread_mutex_unlock(&tunnel_mutex);
    usleep(1000000); // 1 second - increased timeout for better thread cleanup
    pthread_mutex_lock(&tunnel_mutex);

    // Now safely clean up - free tunnel structures since threads should have stopped
    for (int i = active_tunnels_count - 1; i >= 0; i--) {
        if (active_tunnels[i]) {
            if (active_tunnels[i]->sock >= 0) {
                // Close the socket directly without validity checks
                close(active_tunnels[i]->sock);
                active_tunnels[i]->sock = -1;
                if (debug) {
                    printf("[DEBUG] [TCP Tunnel] Closed TCP connection for tunnel %s during cleanup\n", active_tunnels[i]->request_id);
                }
            }
            if (active_tunnels[i]->local_sock >= 0) {
                close(active_tunnels[i]->local_sock);
                active_tunnels[i]->local_sock = -1;
            }
            // Don't free SSL here - it's managed at the connection level
            // Just set the pointer to NULL to prevent use-after-free
            active_tunnels[i]->ssl = NULL;

            if (active_tunnels[i]->outgoing_buffer) {
                frame_buffer_free(active_tunnels[i]->outgoing_buffer);
                active_tunnels[i]->outgoing_buffer = NULL;
            }
            if (active_tunnels[i]->incoming_buffer) {
                frame_buffer_free(active_tunnels[i]->incoming_buffer);
                active_tunnels[i]->incoming_buffer = NULL;
            }
            if (active_tunnels[i]->retransmission_buffer) {
                retransmission_buffer_free(active_tunnels[i]->retransmission_buffer);
                active_tunnels[i]->retransmission_buffer = NULL;
            }

            // Clear backward compatibility pointer if it points to this tunnel
            if (active_tunnel == active_tunnels[i]) {
                active_tunnel = NULL;
            }

            // Free the tunnel structure
            free(active_tunnels[i]);
            active_tunnels[i] = NULL;

            // Shift remaining tunnels down
            for (int j = i; j < active_tunnels_count - 1; j++) {
                active_tunnels[j] = active_tunnels[j + 1];
            }
            active_tunnels_count--;
            i++; // Adjust index since we shifted elements
        }
    }

    // Reset global shutdown flag for potential restart
    global_shutdown = 0;

    pthread_mutex_unlock(&tunnel_mutex);
}

void *forward_tcp_to_ws(void *arg) {
    thread_args_t *args = (thread_args_t *)arg;
    SSL *ssl = args->ssl;
    tunnel_t *tunnel = args->tunnel;
    int debug = args->debug;
    char buffer[BUFFER_SIZE];
    int bytes_read;
    fd_set readfds;
    struct timeval tv;

    while (1) {
        pthread_mutex_lock(&tunnel_mutex);
        if (!tunnel || !tunnel->active) {
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }
        int sock = tunnel->local_sock;
        char request_id[37];
        strcpy(request_id, tunnel->request_id);

        // Check if socket is valid
        if (sock < 0) {
            if (debug) {
                printf("[DEBUG - Tunnel] Socket not ready yet, waiting...\n");
                fflush(stdout);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            usleep(10000); // Sleep 10ms and try again
            continue;
        }

        // Check if tunnel is still active after copying values
        tunnel_t *check_tunnel = find_tunnel_by_request_id(request_id);
        if (!check_tunnel || !check_tunnel->active) {
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }

        // For wsscp: The connection should already be established
        int client_sock = sock;
        if (tunnel->sock < 0 && tunnel->outgoing_buffer) {
            // For wsscp, the socket should already be connected when this function starts
            // No need to accept connections here - the main process already did that
            if (debug) {
                printf("[DEBUG - Tunnel] wsscp socket already connected, starting data forwarding\n");
                fflush(stdout);
            }
            // Store the connected socket
            tunnel->sock = sock;
        }

        // Send pending data from outgoing buffer to client socket (wsscp only)
        if (tunnel->outgoing_buffer && tunnel->outgoing_buffer->used > 0) {
            ssize_t sent = send(client_sock, tunnel->outgoing_buffer->buffer, tunnel->outgoing_buffer->used, MSG_DONTWAIT);
            if (sent > 0) {
                frame_buffer_consume(tunnel->outgoing_buffer, sent);
                if (debug) {
                    printf("[DEBUG - TCPConnection] Sent %zd bytes from buffer to local socket\n", sent);
                    fflush(stdout);
                }
            } else if (sent == -1 && errno != EAGAIN && errno != EWOULDBLOCK) {
                if (debug) {
                    perror("[DEBUG] send to local socket failed");
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                break;
            }
        }

        pthread_mutex_unlock(&tunnel_mutex);

        // Use select to wait for data on client socket
        FD_ZERO(&readfds);
        FD_SET(client_sock, &readfds);
        tv.tv_sec = 0;  // 0 seconds
        tv.tv_usec = 200000;  // 200ms timeout (reduced CPU usage)

        int retval = select(client_sock + 1, &readfds, NULL, NULL, &tv);
        if (retval == -1) {
            if (debug) {
                perror("[DEBUG] select failed");
                fflush(stdout);
            }
            break;
        } else if (retval == 0) {
            // Timeout, continue loop
            continue;
        }

        if (FD_ISSET(client_sock, &readfds)) {
            bytes_read = recv(client_sock, buffer, MAX_CHUNK_SIZE, 0);
            if (bytes_read <= 0) {
                if (debug) {
                    if (bytes_read == 0) {
                        printf("[DEBUG - TCPConnection] TCP connection closed by peer (SSH client disconnected)\n");
                    } else {
                        printf("[DEBUG - TCPConnection] TCP connection error: %s\n", strerror(errno));
                    }
                    fflush(stdout);
                }
                // Mark tunnel as inactive since SSH connection is broken
                pthread_mutex_lock(&tunnel_mutex);
                if (tunnel) {
                    tunnel->active = 0;
                    tunnel->broken = 1;
                    // Send tunnel_close notification immediately when local connection breaks
                    if (debug) {
                        printf("[DEBUG - Tunnel] Sending tunnel_close notification from forwarding thread...\n");
                        fflush(stdout);
                    }
                    send_tunnel_close(tunnel->ssl, tunnel->request_id, debug);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                break;
            }

            if (debug) {
                printf("[DEBUG - TCPConnection] Forwarding %d bytes from TCP to WebSocket\n", bytes_read);
                fflush(stdout);
            }

            // Check if binary mode is enabled for this tunnel
            int use_binary = 0;
            pthread_mutex_lock(&tunnel_mutex);
            tunnel_t *bin_tunnel = find_tunnel_by_request_id(request_id);
            if (bin_tunnel) {
                use_binary = bin_tunnel->bin;
            }
            pthread_mutex_unlock(&tunnel_mutex);

            // Track sent bytes for statistics
            size_t data_len = bytes_read;
            pthread_mutex_lock(&tunnel_mutex);
            tunnel_t *stats_tunnel = find_tunnel_by_request_id(request_id);
            if (stats_tunnel) {
                stats_tunnel->total_bytes_sent += data_len;
                stats_tunnel->bytes_last_period += data_len;
            }
            pthread_mutex_unlock(&tunnel_mutex);

            // Use reliable transmission
            if (tunnel->retransmission_buffer) {
                if (!send_tunnel_data_reliable_message(ssl, request_id, (unsigned char *)buffer, bytes_read, tunnel->retransmission_buffer, tunnel->encoding, debug)) {
                    break;
                }
            } else {
                // Fallback to regular transmission if no retransmission buffer
                if (use_binary) {
                    // Send binary data directly
                    if (!send_tunnel_data_binary_message(ssl, request_id, (unsigned char *)buffer, bytes_read, debug)) {
                        break;
                    }
                } else {
                    // Convert to hex with dynamic allocation
                    size_t hex_size = (size_t)bytes_read * 2 + 1;
                    char *hex_data = malloc(hex_size);
                    if (!hex_data) {
                        if (debug) {
                            printf("[DEBUG] Failed to allocate memory for hex data (%zu bytes)\n", hex_size);
                            fflush(stdout);
                        }
                        continue;
                    }
                    size_t actual_hex_len = 0;
                    for (int i = 0; i < bytes_read; i++) {
                        sprintf(hex_data + actual_hex_len, "%02x", (unsigned char)buffer[i]);
                        actual_hex_len += 2;
                    }
                    hex_data[actual_hex_len] = '\0';

                    // Send as tunnel_data
                    if (!send_tunnel_data_message(ssl, request_id, hex_data, debug)) {
                        free(hex_data);
                        break;
                    }
                    free(hex_data);
                }
            }
        }
    }

    if (debug) {
        printf("[DEBUG - TCPConnection] TCP to WebSocket forwarding thread exiting\n");
        fflush(stdout);
    }

    // Mark tunnel as inactive when forwarding thread exits due to broken connection
    if (tunnel) {
        pthread_mutex_lock(&tunnel_mutex);
        if (tunnel->active) {
            tunnel->active = 0;
            if (debug) {
                printf("[DEBUG - TCPConnection] Marked tunnel as inactive due to forwarding thread exit\n");
                fflush(stdout);
            }
            // Send tunnel_close notification
            send_tunnel_close(tunnel->ssl, tunnel->request_id, debug);
        }
        pthread_mutex_unlock(&tunnel_mutex);
    }

    free(args);
    return NULL;
}


void *forward_ws_to_ssh_server(void *arg) {
    thread_args_t *args = (thread_args_t *)arg;
    SSL *ssl = args->ssl;
    tunnel_t *tunnel = args->tunnel;
    int debug = args->debug;
    char buffer[BUFFER_SIZE];
    int bytes_read;
    fd_set readfds;
    struct timeval tv;

    while (1) {
        pthread_mutex_lock(&tunnel_mutex);
        if (!tunnel || !tunnel->active) {
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }
        int target_sock = tunnel->sock; // Target TCP connection
        char request_id[37];
        strcpy(request_id, tunnel->request_id);
        pthread_mutex_unlock(&tunnel_mutex);

        // Check if tunnel is still active after copying values
        pthread_mutex_lock(&tunnel_mutex);
        tunnel_t *check_tunnel = find_tunnel_by_request_id(request_id);
        if (!check_tunnel || !check_tunnel->active) {
            pthread_mutex_unlock(&tunnel_mutex);
            break;
        }
        pthread_mutex_unlock(&tunnel_mutex);

        // Use select to wait for data on target TCP connection
        FD_ZERO(&readfds);
        FD_SET(target_sock, &readfds);
        tv.tv_sec = 0;  // 0 seconds
        tv.tv_usec = 200000;  // 200ms timeout (reduced CPU usage)

        int retval = select(target_sock + 1, &readfds, NULL, NULL, &tv);
        if (retval == -1) {
            if (debug) {
                perror("[DEBUG] select on target socket failed");
                fflush(stdout);
            }
            break;
        } else if (retval == 0) {
            // Timeout, continue loop
            continue;
        }

        if (FD_ISSET(target_sock, &readfds)) {
            bytes_read = recv(target_sock, buffer, MAX_CHUNK_SIZE, 0);
            if (bytes_read <= 0) {
                if (debug) {
                    printf("[DEBUG - TCPConnection] Target connection closed or error\n");
                    fflush(stdout);
                }
                break;
            }

            if (debug) {
                printf("[DEBUG - TCPConnection] Forwarding %d bytes from target to WebSocket\n", bytes_read);
                fflush(stdout);
            }

            // Convert to hex with dynamic allocation
            size_t hex_size = (size_t)bytes_read * 2 + 1;
            char *hex_data = malloc(hex_size);
            if (!hex_data) {
                if (debug) {
                    printf("[DEBUG] Failed to allocate memory for hex data (%zu bytes)\n", hex_size);
                    fflush(stdout);
                }
                continue;
            }
            size_t actual_hex_len = 0;
            for (int i = 0; i < bytes_read; i++) {
                sprintf(hex_data + actual_hex_len, "%02x", (unsigned char)buffer[i]);
                actual_hex_len += 2;
            }
            hex_data[actual_hex_len] = '\0';

            // Track sent bytes for statistics (responses from target)
            size_t data_len = bytes_read;
            pthread_mutex_lock(&tunnel_mutex);
            tunnel_t *stats_tunnel = find_tunnel_by_request_id(request_id);
            if (stats_tunnel) {
                stats_tunnel->total_bytes_sent += data_len;
                stats_tunnel->bytes_last_period += data_len;
            }
            pthread_mutex_unlock(&tunnel_mutex);

            // Send as tunnel_response (from target back to WebSocket)
            if (!send_tunnel_response_message(ssl, request_id, hex_data, debug)) {
                free(hex_data);
                break;
            }
            free(hex_data);
        }
    }

    if (debug) {
        printf("[DEBUG - TCPConnection] Target to WebSocket forwarding thread exiting\n");
        fflush(stdout);
    }

    free(args);
    return NULL;
}

void handle_tunnel_data(SSL *ssl, const char *request_id, const char *data_hex, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (!tunnel) {
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    // Check if tunnel is in binary mode
    int is_binary = tunnel->bin;

    // Check if this is a reliable transmission message (has frame_id and checksum)
    uint32_t frame_id = 0;
    uint32_t expected_checksum = 0;
    int is_reliable = 0;

    // Parse JSON to extract frame_id and checksum if present
    // Simple parsing for frame_id and checksum fields
    const char *frame_id_str = strstr(data_hex, "\"frame_id\":");
    const char *checksum_str = strstr(data_hex, "\"checksum\":");

    if (frame_id_str && checksum_str) {
        // Extract frame_id
        frame_id_str += strlen("\"frame_id\":");
        if (sscanf(frame_id_str, "%u", &frame_id) == 1) {
            // Extract checksum
            checksum_str += strlen("\"checksum\":");
            if (sscanf(checksum_str, "%u", &expected_checksum) == 1) {
                is_reliable = 1;
            }
        }
    }

    // Get data length for processing
    size_t hex_len = strlen(data_hex);

    // For reliable transmission, validate checksum and send ACK/KO
    if (is_reliable) {
        // Extract the actual data part (after checksum field)
        const char *data_start = strstr(data_hex, "\"data\":\"");
        if (!data_start) {
            // Invalid format, send KO
            send_tunnel_ko_message(ssl, request_id, frame_id, debug);
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
        data_start += strlen("\"data\":\"");

        // Find the end of data
        const char *data_end = strstr(data_start, "\"");
        if (!data_end) {
            send_tunnel_ko_message(ssl, request_id, frame_id, debug);
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }

        size_t data_part_len = data_end - data_start;
        char *data_part = malloc(data_part_len + 1);
        if (!data_part) {
            send_tunnel_ko_message(ssl, request_id, frame_id, debug);
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
        memcpy(data_part, data_start, data_part_len);
        data_part[data_part_len] = '\0';

        // Calculate checksum of the data part
        uint32_t calculated_checksum = crc32_checksum((const unsigned char *)data_part, data_part_len);

        if (calculated_checksum == expected_checksum) {
            // Checksum valid, send ACK
            send_tunnel_ack_message(ssl, request_id, frame_id, debug);
            // Use the extracted data for further processing
            data_hex = data_part;
            hex_len = data_part_len;
        } else {
            // Checksum invalid, send KO
            if (debug) {
                printf("[DEBUG] Checksum validation failed: expected %u, got %u\n", expected_checksum, calculated_checksum);
                fflush(stdout);
            }
            send_tunnel_ko_message(ssl, request_id, frame_id, debug);
            free(data_part);
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }

        // Note: data_part will be freed later in the function
    }

    size_t data_len = 0;
    char *data = NULL;

    if (is_binary) {
        // Binary mode: use data as-is (no decoding)
        data_len = hex_len;
        data = malloc(data_len);
        if (!data) {
            if (debug) {
                printf("[DEBUG] Failed to allocate memory for %zu bytes of binary data\n", data_len);
                fflush(stdout);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
        memcpy(data, data_hex, data_len);

        if (debug) {
            printf("[DEBUG] Using %zu bytes of binary data\n", data_len);
            fflush(stdout);
        }

        // Update statistics
        tunnel->total_bytes_received += data_len;
        tunnel->bytes_last_period += data_len;
    } else {
        // Decode based on tunnel encoding
        wsssh_encoding_t encoding = tunnel->encoding;

        if (encoding == ENCODING_BINARY) {
            // Binary mode: use data as-is (no decoding)
            data_len = hex_len;
            data = malloc(data_len);
            if (!data) {
                if (debug) {
                    printf("[DEBUG] Failed to allocate memory for %zu bytes of binary data\n", data_len);
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }
            memcpy(data, data_hex, data_len);

            if (debug) {
                printf("[DEBUG] Using %zu bytes of binary data\n", data_len);
                fflush(stdout);
            }

            // Update statistics
            tunnel->total_bytes_received += data_len;
            tunnel->bytes_last_period += data_len;
        } else if (encoding == ENCODING_BASE64) {
            // Base64 decoding
            data_len = (hex_len * 3) / 4;
            if (hex_len > 0 && data_hex[hex_len - 1] == '=') data_len--;
            if (hex_len > 1 && data_hex[hex_len - 2] == '=') data_len--;

            data = malloc(data_len);
            if (!data) {
                if (debug) {
                    printf("[DEBUG] Failed to allocate memory for %zu bytes of base64 decoded data\n", data_len);
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

            // Simple base64 decoding
            static const char base64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
            size_t i = 0, j = 0;
            int decode_success = 1;
            while (i < hex_len && data_hex[i] != '=') {
                // Find indices in base64 table
                int a = -1, b = -1, c = -1, d = -1;
                for (int k = 0; k < 64; k++) {
                    if (base64_table[k] == data_hex[i]) a = k;
                    if (i + 1 < hex_len && base64_table[k] == data_hex[i + 1]) b = k;
                    if (i + 2 < hex_len && base64_table[k] == data_hex[i + 2]) c = k;
                    if (i + 3 < hex_len && base64_table[k] == data_hex[i + 3]) d = k;
                }

                if (a == -1 || b == -1) {
                    decode_success = 0;
                    break;
                }

                data[j++] = (a << 2) | (b >> 4);
                if (c != -1 && j < data_len) {
                    data[j++] = (b << 4) | (c >> 2);
                }
                if (d != -1 && j < data_len) {
                    data[j++] = (c << 6) | d;
                }
                i += 4;
            }

            if (!decode_success) {
                free(data);
                if (debug) {
                    printf("[DEBUG] Base64 decoding failed\n");
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            } else {
                data_len = j; // Actual decoded length

                if (debug) {
                    printf("[DEBUG] Successfully decoded %zu bytes using base64\n", data_len);
                    fflush(stdout);
                }

                // Update statistics
                tunnel->total_bytes_received += data_len;
                tunnel->bytes_last_period += data_len;
            }
        } else if (encoding == ENCODING_HEX) {
            // Hex decoding
            if (hex_len % 2 != 0) {
                if (debug) {
                    printf("[DEBUG] Invalid hex data length: %zu (must be even)\n", hex_len);
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

            // Additional validation: check for reasonable data size
            size_t expected_data_len = hex_len / 2;
            if (expected_data_len > 64 * 1024) { // 64KB limit for SSH packets
                if (debug) {
                    printf("[DEBUG] Hex data too large: %zu bytes, truncating to 64KB\n", expected_data_len);
                    fflush(stdout);
                }
                // Truncate to reasonable size for SSH
                hex_len = 2 * 64 * 1024; // 128KB hex = 64KB data
                // Make sure we don't truncate in the middle of a hex byte
                if (hex_len % 2 != 0) hex_len--;
                expected_data_len = hex_len / 2;
            }

            data_len = expected_data_len;
            data = malloc(data_len);
            if (!data) {
                if (debug) {
                    printf("[DEBUG] Failed to allocate memory for %zu bytes of hex decoded data\n", data_len);
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

            // Use more efficient hex decoding
            int hex_decode_success = 1;
            for (size_t i = 0; i < data_len; i++) {
                unsigned int byte_val;
                if (i * 2 >= hex_len || sscanf(data_hex + i * 2, "%2x", &byte_val) != 1) {
                    if (debug) {
                        printf("[DEBUG] Failed to decode hex byte at position %zu\n", i * 2);
                        fflush(stdout);
                    }
                    hex_decode_success = 0;
                    break;
                }
                data[i] = (char)byte_val;
            }

            if (!hex_decode_success) {
                free(data);
                if (debug) {
                    printf("[DEBUG] Hex decoding failed\n");
                    fflush(stdout);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }

            if (debug) {
                printf("[DEBUG] Successfully decoded %zu bytes using hex\n", data_len);
                fflush(stdout);
            }

            // Update statistics
            tunnel->total_bytes_received += data_len;
            tunnel->bytes_last_period += data_len;
        } else {
            if (debug) {
                printf("[DEBUG] Unknown encoding type: %d\n", encoding);
                fflush(stdout);
            }
            pthread_mutex_unlock(&tunnel_mutex);
            return;
        }
    }

    int target_sock = -1;
    if (tunnel->outgoing_buffer) {
        // wsscp: Use local_sock (SCP client connection)
        target_sock = tunnel->local_sock;
        if (debug) {
            printf("[DEBUG] Socket selection: wsscp mode, target_sock=%d (local_sock)\n", target_sock);
            fflush(stdout);
        }
    } else if (tunnel->sock >= 0) {
        // wssshc: Use sock (direct SSH server connection)
        target_sock = tunnel->sock;
        if (debug) {
            printf("[DEBUG] Socket selection: wssshc mode, target_sock=%d (sock)\n", target_sock);
            fflush(stdout);
        }
    } else if (tunnel->local_sock >= 0) {
        // wsssh: Use local_sock (accepted SSH client connection)
        target_sock = tunnel->local_sock;
        if (debug) {
            printf("[DEBUG] Socket selection: wsssh mode, target_sock=%d (local_sock)\n", target_sock);
            fflush(stdout);
        }
    } else {
        // No connection established yet - buffer the data
        // Suppress tunnel_data debug messages in debug mode
        // Ensure we have an incoming buffer for wsssh
        if (!tunnel->incoming_buffer) {
            tunnel->incoming_buffer = frame_buffer_init();
            if (!tunnel->incoming_buffer) {
                if (debug) {
                    printf("[DEBUG] Failed to create incoming buffer\n");
                    fflush(stdout);
                }
                free(data);
                pthread_mutex_unlock(&tunnel_mutex);
                return;
            }
        }
        // Buffer the data until connection is established
        if (!frame_buffer_append(tunnel->incoming_buffer, data, data_len)) {
            if (debug) {
                printf("[DEBUG] Failed to buffer incoming data, dropping %zu bytes\n", data_len);
                fflush(stdout);
            }
        }
        free(data);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    // Check if target socket is valid
    if (target_sock < 0) {
        // Suppress tunnel_data debug messages in debug mode
        free(data);
        pthread_mutex_unlock(&tunnel_mutex);
        return;
    }

    // Unlock mutex before sending to avoid blocking other threads
    pthread_mutex_unlock(&tunnel_mutex);

    if (tunnel->outgoing_buffer) {
        // wsscp: Append to outgoing buffer
        pthread_mutex_lock(&tunnel_mutex);
        if (!frame_buffer_append(tunnel->outgoing_buffer, data, data_len)) {
            if (debug) {
                printf("[DEBUG] Failed to append to outgoing buffer, dropping %zu bytes\n", data_len);
                fflush(stdout);
            }
        }
        pthread_mutex_unlock(&tunnel_mutex);
    } else {
        // wsssh/wssshc: Send directly to target socket
        if (debug) {
            printf("[DEBUG] Sending %zu bytes of tunnel data to socket %d\n", data_len, target_sock);
            fflush(stdout);
        }

        size_t total_sent = 0;
        while (total_sent < data_len) {
            ssize_t sent = send(target_sock, data + total_sent, data_len - total_sent, 0);
            if (sent < 0) {
                if (debug) {
                    printf("[DEBUG] Failed to send %zu bytes to target socket %d: %s (errno=%d)\n",
                            data_len - total_sent, target_sock, strerror(errno), errno);
                    fflush(stdout);
                }

                // Mark tunnel inactive for any send error
                pthread_mutex_lock(&tunnel_mutex);
                if (tunnel) {
                    tunnel->active = 0;
                    tunnel->broken = 1;
                    send_tunnel_close(tunnel->ssl, tunnel->request_id, debug);
                }
                pthread_mutex_unlock(&tunnel_mutex);
                free(data);
                return;
            }
            total_sent += sent;
        }

        if (debug) {
            printf("[DEBUG] Successfully sent %zu bytes to target\n", data_len);
            fflush(stdout);
        }
    }
    free(data);
}

void send_tunnel_close(SSL *ssl, const char *request_id, int debug) {
    char close_msg[256];
    snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);

    if (debug) {
        printf("[DEBUG - Tunnel] Sending tunnel_close: %s\n", close_msg);
        fflush(stdout);
    }

    if (!send_websocket_frame(ssl, close_msg)) {
        if (debug) {
            printf("[DEBUG - Tunnel] Failed to send tunnel_close message\n");
            fflush(stdout);
        }
    }
}

void handle_tunnel_close(SSL *ssl __attribute__((unused)), const char *request_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (tunnel) {
        tunnel->active = 0;
        remove_tunnel(request_id);
        if (debug) {
            printf("[DEBUG - Tunnel] Tunnel %s closed\n", request_id);
            fflush(stdout);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void handle_tunnel_keepalive(SSL *ssl, const char *request_id, unsigned long long total_bytes, double rate_bps, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (tunnel) {
        tunnel->last_keepalive_received = time(NULL);
        if (debug) {
            printf("[DEBUG - Tunnel] Keep-alive received for tunnel %s: total=%llu bytes, rate=%.2f B/s\n",
                   request_id, total_bytes, rate_bps);
            fflush(stdout);
        }

        // Send ACK response
        char ack_msg[256];
        snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_keepalive_ack\",\"request_id\":\"%s\"}", request_id);

        if (debug) {
            printf("[DEBUG - WebSockets] Sending tunnel_keepalive_ack: %s\n", ack_msg);
            fflush(stdout);
        }

        if (!send_websocket_frame(ssl, ack_msg)) {
            if (debug) {
                printf("[DEBUG - Tunnel] Failed to send keep-alive ACK for tunnel %s\n", request_id);
                fflush(stdout);
            }
        }
    } else {
        if (debug) {
            printf("[DEBUG - Tunnel] Keep-alive received for unknown tunnel %s\n", request_id);
            fflush(stdout);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void handle_tunnel_keepalive_ack(SSL *ssl __attribute__((unused)), const char *request_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (tunnel) {
        if (debug) {
            printf("[DEBUG - Tunnel] Keep-alive ACK received for tunnel %s\n", request_id);
            fflush(stdout);
        }
        // ACK received, tunnel is alive
    } else {
        if (debug) {
            printf("[DEBUG - Tunnel] Keep-alive ACK received for unknown tunnel %s\n", request_id);
            fflush(stdout);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void handle_tunnel_ack(SSL *ssl __attribute__((unused)), const char *request_id, uint32_t frame_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (tunnel && tunnel->retransmission_buffer) {
        retransmission_buffer_ack(tunnel->retransmission_buffer, frame_id);
        if (debug) {
            printf("[DEBUG - Tunnel] ACK received for frame %u in tunnel %s\n", frame_id, request_id);
            fflush(stdout);
        }
    } else {
        if (debug) {
            printf("[DEBUG - Tunnel] ACK received for unknown tunnel %s or no retransmission buffer\n", request_id);
            fflush(stdout);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void handle_tunnel_ko(SSL *ssl __attribute__((unused)), const char *request_id, uint32_t frame_id, int debug) {
    pthread_mutex_lock(&tunnel_mutex);
    tunnel_t *tunnel = find_tunnel_by_request_id(request_id);
    if (tunnel && tunnel->retransmission_buffer) {
        retransmission_buffer_ko(tunnel->retransmission_buffer, frame_id);
        if (debug) {
            printf("[DEBUG - Tunnel] KO received for frame %u in tunnel %s\n", frame_id, request_id);
            fflush(stdout);
        }
    } else {
        if (debug) {
            printf("[DEBUG - Tunnel] KO received for unknown tunnel %s or no retransmission buffer\n", request_id);
            fflush(stdout);
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

void send_tunnel_keepalive(SSL *ssl, tunnel_t *tunnel, int debug) {
    if (!tunnel || !tunnel->active) return;

    time_t current_time = time(NULL);

    // Reset stats every 30 seconds
    if (current_time - tunnel->last_stats_reset >= 30) {
        tunnel->bytes_last_period = 0;
        tunnel->last_stats_reset = current_time;
    }

    // Calculate rate (bytes per second over last 30 seconds)
    double rate_bps = 0.0;
    if (current_time > tunnel->last_stats_reset) {
        rate_bps = (double)tunnel->bytes_last_period / (current_time - tunnel->last_stats_reset);
    }

    // Send keep-alive message
    char keepalive_msg[512];
    unsigned long long total_bytes = tunnel->total_bytes_sent + tunnel->total_bytes_received;
    snprintf(keepalive_msg, sizeof(keepalive_msg),
             "{\"type\":\"tunnel_keepalive\",\"request_id\":\"%s\",\"total_bytes\":%llu,\"rate_bps\":%.2f}",
             tunnel->request_id, total_bytes, rate_bps);

    if (debug) {
        printf("[DEBUG - Tunnel] Sending keep-alive for tunnel %s: %s\n", tunnel->request_id, keepalive_msg);
        fflush(stdout);
    }

    if (send_websocket_frame(ssl, keepalive_msg)) {
        tunnel->last_keepalive_sent = current_time;
    } else if (debug) {
        printf("[DEBUG - Tunnel] Failed to send keep-alive for tunnel %s\n", tunnel->request_id);
        fflush(stdout);
    }
}

void check_keepalive_timeouts(int debug) {
    time_t current_time = time(NULL);
    const int timeout_seconds = 190; // 3 * 30 + 10 seconds margin

    pthread_mutex_lock(&tunnel_mutex);
    for (int i = 0; i < active_tunnels_count; i++) {
        tunnel_t *tunnel = active_tunnels[i];
        if (tunnel && tunnel->active) {
            // Check if we haven't received a keep-alive from the other side for too long
            if (current_time - tunnel->last_keepalive_received > timeout_seconds) {
                if (debug) {
                    printf("[DEBUG - Tunnel] Keep-alive timeout for tunnel %s (%ld seconds since last keep-alive)\n",
                           tunnel->request_id, current_time - tunnel->last_keepalive_received);
                    fflush(stdout);
                }

                // Send tunnel close and mark as broken
                tunnel->active = 0;
                tunnel->broken = 1;
                send_tunnel_close(tunnel->ssl, tunnel->request_id, debug);
            }

            // Garbage collect retransmission buffer
            if (tunnel->retransmission_buffer) {
                retransmission_buffer_gc(tunnel->retransmission_buffer);
            }
        }
    }
    pthread_mutex_unlock(&tunnel_mutex);
}

int reconnect_websocket(tunnel_t *tunnel, const char *wssshd_host, int wssshd_port, const char *client_id, const char *request_id, int debug) {
    struct sockaddr_in server_addr;
    struct hostent *he;
    int sock;
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    if (debug) {
        printf("[DEBUG] Attempting to reconnect WebSocket connection...\n");
        fflush(stdout);
    }

    // Resolve hostname
    if ((he = gethostbyname(wssshd_host)) == NULL) {
        herror("gethostbyname");
        return -1;
    }

    // Create socket
    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("Socket creation failed");
        return -1;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(wssshd_port);
    server_addr.sin_addr = *((struct in_addr *)he->h_addr);

    // Connect to server
    if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Connection failed");
        close(sock);
        return -1;
    }

    // Create SSL context and connection
    ssl_ctx = create_ssl_context();
    if (!ssl_ctx) {
        close(sock);
        return -1;
    }

    ssl = create_ssl_connection(ssl_ctx, sock, debug);
    if (!ssl) {
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Perform WebSocket handshake
    if (!websocket_handshake(ssl, wssshd_host, wssshd_port, "/", debug)) {
        fprintf(stderr, "WebSocket handshake failed\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Send detailed tunnel request with transport information
    char *expanded_tunnel = expand_transport_list("any", 0); // Data channel transports
    char *expanded_tunnel_control = expand_transport_list("any", 1); // Control channel transports

    // Select best transport based on weight (lowest weight = highest priority)
    char *best_tunnel = select_best_transport(expanded_tunnel);
    char *best_tunnel_control = select_best_transport(expanded_tunnel_control);

    if (!send_tunnel_request_message_with_enc(ssl, client_id, request_id, best_tunnel ? best_tunnel : expanded_tunnel, best_tunnel_control ? best_tunnel_control : expanded_tunnel_control, "ssh", tunnel->bin ? ENCODING_BINARY : ENCODING_HEX)) {
        free(expanded_tunnel);
        free(expanded_tunnel_control);
        if (best_tunnel) free(best_tunnel);
        if (best_tunnel_control) free(best_tunnel_control);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    free(expanded_tunnel);
    free(expanded_tunnel_control);
    if (best_tunnel) free(best_tunnel);
    if (best_tunnel_control) free(best_tunnel_control);

    // Read acknowledgment
    bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
    if (bytes_read <= 0) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Check if it's a close frame
    if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) {
        if (debug) {
            printf("[DEBUG] Server sent close frame during reconnection\n");
            fflush(stdout);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Parse WebSocket frame
    char *payload;
    int payload_len;
    if (!parse_websocket_frame(buffer, bytes_read, &payload, &payload_len)) {
        fprintf(stderr, "Failed to parse WebSocket frame during reconnection\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Null terminate payload
    payload[payload_len] = '\0';

    // Check for tunnel acknowledgment
    if (strstr(payload, "tunnel_ack") == NULL) {
        fprintf(stderr, "Tunnel request denied during reconnection: %s\n", payload);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return -1;
    }

    // Success - update the tunnel's SSL connection
    if (tunnel->ssl) {
        SSL_free(tunnel->ssl);
    }
    tunnel->ssl = ssl;

    // Clean up old SSL context
    SSL_CTX_free(ssl_ctx);

    if (debug) {
        printf("[DEBUG] WebSocket reconnection successful\n");
        fflush(stdout);
    }

    return 0;
}

tunnel_setup_result_t setup_tunnel(const char *wssshd_host, int wssshd_port, const char *client_id, int local_port, int debug, int use_buffer, const char *tunnel_host, wsssh_encoding_t encoding, int send_tunnel_request_immediately) {
    struct sockaddr_in server_addr;
    struct hostent *he;
    int sock;
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // Generate request ID
    char request_id[37];
    generate_request_id(request_id, sizeof(request_id));

    // Resolve hostname
    if ((he = gethostbyname(wssshd_host)) == NULL) {
        herror("gethostbyname");
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    // Create socket
    if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("Socket creation failed");
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(wssshd_port);
    server_addr.sin_addr = *((struct in_addr *)he->h_addr);

    // Connect to server
    if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Connection failed");
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    // Create SSL context and connection
    ssl_ctx = create_ssl_context();
    if (!ssl_ctx) {
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    ssl = create_ssl_connection(ssl_ctx, sock, debug);
    if (!ssl) {
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    // Perform WebSocket handshake
    if (debug) {
        printf("[DEBUG] Performing WebSocket handshake to %s:%d\n", wssshd_host, wssshd_port);
        fflush(stdout);
    }
    if (!websocket_handshake(ssl, wssshd_host, wssshd_port, "/", debug)) {
        fprintf(stderr, "WebSocket handshake failed\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }
    if (debug) {
        printf("[DEBUG] WebSocket handshake successful\n");
        fflush(stdout);
    }

    if (send_tunnel_request_immediately) {
        // Send detailed tunnel request with transport information
        char *expanded_tunnel = expand_transport_list("any", 0); // Data channel transports
        char *expanded_tunnel_control = expand_transport_list("any", 1); // Control channel transports

        // Select best transport based on weight (lowest weight = highest priority)
        char *best_tunnel = select_best_transport(expanded_tunnel);
        char *best_tunnel_control = select_best_transport(expanded_tunnel_control);

        if (!send_tunnel_request_message_with_enc(ssl, client_id, request_id, best_tunnel ? best_tunnel : expanded_tunnel, best_tunnel_control ? best_tunnel_control : expanded_tunnel_control, "ssh", encoding)) {
            free(expanded_tunnel);
            free(expanded_tunnel_control);
            if (best_tunnel) free(best_tunnel);
            if (best_tunnel_control) free(best_tunnel_control);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        free(expanded_tunnel);
        free(expanded_tunnel_control);
        if (best_tunnel) free(best_tunnel);
        if (best_tunnel_control) free(best_tunnel_control);

        if (debug) {
            printf("[DEBUG] Tunnel request sent for client: %s, request_id: %s\n", client_id, request_id);
        }

        // Read acknowledgment
        bytes_read = SSL_read(ssl, buffer, sizeof(buffer));
        if (bytes_read <= 0) {
            if (debug) {
                printf("[DEBUG] No acknowledgment received\n");
            }
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        if (debug) {
            printf("[DEBUG] Read %d bytes acknowledgment\n", bytes_read);
            printf("[DEBUG] Raw acknowledgment: ");
            for (int i = 0; i < bytes_read && i < 50; i++) {
                if (buffer[i] >= 32 && buffer[i] < 127) {
                    printf("%c", buffer[i]);
                } else {
                    printf("\\x%02x", (unsigned char)buffer[i]);
                }
            }
            printf("\n");
            fflush(stdout);
        }

        // Check if it's a close frame
        if (bytes_read >= 2 && (buffer[0] & 0x8F) == 0x88) {
            if (debug) {
                printf("[DEBUG] Server sent close frame\n");
                if (bytes_read >= 6) {
                    int close_code = (buffer[2] << 8) | buffer[3];
                    printf("[DEBUG] Close code: %d\n", close_code);
                    if (bytes_read > 6) {
                        printf("[DEBUG] Close reason: %.*s\n", bytes_read - 6, buffer + 6);
                    }
                }
                fflush(stdout);
            }
            fprintf(stderr, "Tunnel request rejected by server\n");
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        // Parse WebSocket frame
        char *payload;
        int payload_len;
        if (!parse_websocket_frame(buffer, bytes_read, &payload, &payload_len)) {
            fprintf(stderr, "Failed to parse WebSocket frame\n");
            if (debug) {
                printf("[DEBUG] Frame parsing failed. First few bytes: ");
                for (int i = 0; i < bytes_read && i < 10; i++) {
                    printf("%02x ", (unsigned char)buffer[i]);
                }
                printf("\n");
                fflush(stdout);
            }
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        // Null terminate payload
        payload[payload_len] = '\0';

        if (debug) {
            printf("[DEBUG] Received payload: '%s' (length: %d)\n", payload, payload_len);
            printf("[DEBUG] Looking for tunnel_ack in: ");
            for (int i = 0; i < payload_len && i < 50; i++) {
                if (payload[i] >= 32 && payload[i] < 127) {
                    printf("%c", payload[i]);
                } else {
                    printf("\\x%02x", (unsigned char)payload[i]);
                }
            }
            printf("\n");
            fflush(stdout);
        }

        // Check for tunnel acknowledgment
        if (strstr(payload, "tunnel_ack") == NULL) {
            fprintf(stderr, "Tunnel request denied or failed: %s\n", payload);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        if (debug) {
            printf("[DEBUG] Tunnel established, local port: %d\n", local_port);
        }
    }

    tunnel_t *new_tunnel = NULL;
    if (send_tunnel_request_immediately) {
        // Create tunnel structure
        new_tunnel = malloc(sizeof(tunnel_t));
        if (!new_tunnel) {
            perror("Memory allocation failed");
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }
        memset(new_tunnel, 0, sizeof(tunnel_t));

        if (use_buffer) {
            new_tunnel->outgoing_buffer = frame_buffer_init();
            if (!new_tunnel->outgoing_buffer) {
                perror("Failed to initialize outgoing buffer");
                free(new_tunnel);
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                tunnel_setup_result_t result = {-1, NULL, NULL, ""};
                return result;
            }
        } else {
            new_tunnel->outgoing_buffer = NULL;
        }

        // Initialize incoming buffer for buffering data before connection is established
        new_tunnel->incoming_buffer = frame_buffer_init();
        if (!new_tunnel->incoming_buffer) {
            perror("Failed to initialize incoming buffer");
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            free(new_tunnel);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        strcpy(new_tunnel->request_id, request_id);
        new_tunnel->sock = -1;  // wsssh doesn't connect to remote server
        new_tunnel->local_sock = -1;
        new_tunnel->active = 1;
        new_tunnel->broken = 0;
        new_tunnel->ssl = ssl;
        new_tunnel->server_version_sent = 0;
        new_tunnel->bin = (encoding == ENCODING_BINARY);

        // Initialize retransmission buffer for reliable transmission
        new_tunnel->retransmission_buffer = retransmission_buffer_init();
        if (!new_tunnel->retransmission_buffer) {
            perror("Failed to initialize retransmission buffer");
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
            pthread_mutex_unlock(&tunnel_mutex);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        // Add the new tunnel to the array for multiple tunnel support
        pthread_mutex_lock(&tunnel_mutex);
        if (!add_tunnel(new_tunnel)) {
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
            pthread_mutex_unlock(&tunnel_mutex);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            tunnel_setup_result_t result = {-1, NULL, NULL, ""};
            return result;
        }

        // For backward compatibility with wsssh/wsscp that use active_tunnel global
        active_tunnel = new_tunnel;

        pthread_mutex_unlock(&tunnel_mutex);
    }

    // Start listening on local port
    int listen_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_sock < 0) {
        perror("Local socket creation failed");
        if (send_tunnel_request_immediately && new_tunnel) {
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    // Set SO_REUSEADDR to allow immediate reuse of the port
    int opt = 1;
    if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("setsockopt on listen_sock failed");
        close(listen_sock);
        if (send_tunnel_request_immediately && new_tunnel) {
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    struct sockaddr_in local_addr;
    memset(&local_addr, 0, sizeof(local_addr));
    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(local_port);

    // Use specified tunnel_host or default to 127.0.0.1
    if (tunnel_host && strcmp(tunnel_host, "127.0.0.1") != 0) {
        // Resolve the tunnel_host
        struct hostent *tunnel_he;
        if ((tunnel_he = gethostbyname(tunnel_host)) == NULL) {
            if (debug) {
                fprintf(stderr, "[DEBUG] Failed to resolve tunnel_host '%s', using 127.0.0.1\n", tunnel_host);
            }
            local_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        } else {
            local_addr.sin_addr = *((struct in_addr *)tunnel_he->h_addr);
            if (debug) {
                printf("[DEBUG] Binding tunnel to %s:%d\n", tunnel_host, local_port);
                fflush(stdout);
            }
        }
    } else {
        local_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        if (debug) {
            printf("[DEBUG] Binding tunnel to 127.0.0.1:%d\n", local_port);
            fflush(stdout);
        }
    }

    if (bind(listen_sock, (struct sockaddr *)&local_addr, sizeof(local_addr)) < 0) {
        perror("Local bind failed");
        close(listen_sock);
        if (send_tunnel_request_immediately && new_tunnel) {
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    if (listen(listen_sock, 1) < 0) {
        perror("Local listen failed");
        close(listen_sock);
        if (send_tunnel_request_immediately && new_tunnel) {
            if (use_buffer) frame_buffer_free(new_tunnel->outgoing_buffer);
            frame_buffer_free(new_tunnel->incoming_buffer);
            free(new_tunnel);
        }
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        tunnel_setup_result_t result = {-1, NULL, NULL, ""};
        return result;
    }

    if (debug) {
        printf("[DEBUG] Listening on localhost:%d\n", local_port);
        fflush(stdout);
    }

    if (send_tunnel_request_immediately) {
        // Clean up SSL context
        SSL_CTX_free(ssl_ctx);
    }

    // Return success - tunnel is set up and listening
    tunnel_setup_result_t result = {listen_sock, ssl, ssl_ctx, ""};
    strcpy(result.request_id, request_id);
    return result;
}