/*
 * WSSSH Client (wssshc) - C Implementation
 * WebSocket SSH client that registers with wssshd server.
 *
 * 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 <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 <getopt.h>
#include <pthread.h>
#include <errno.h>
#include <signal.h>
#include <time.h>
#include <dirent.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <stdarg.h>

#include "libwsssht/wssshlib.h"
#include "libwsssht/websocket.h"
#include "libwsssht/wssh_ssl.h"
#include "libwsssht/tunnel.h"
#include "libwsssht/control_messages.h"



int global_debug = 0;
time_t start_time = 0;
volatile sig_atomic_t sigint_count = 0;
volatile sig_atomic_t graceful_shutdown = 0;
volatile sig_atomic_t reload_config = 0;

void format_bytes(unsigned long long bytes, char *buffer, size_t buffer_size) {
    if (bytes < 1024) {
        snprintf(buffer, buffer_size, "%llu B", bytes);
    } else if (bytes < 1024 * 1024) {
        snprintf(buffer, buffer_size, "%.2f kB", (double)bytes / 1024.0);
    } else if (bytes < 1024 * 1024 * 1024) {
        snprintf(buffer, buffer_size, "%.2f MB", (double)bytes / (1024.0 * 1024.0));
    } else {
        snprintf(buffer, buffer_size, "%.2f GB", (double)bytes / (1024.0 * 1024.0 * 1024.0));
    }
}

void log_message(const char *level, const char *format, ...) {
    time_t now = time(NULL);
    char time_str[20];
    strftime(time_str, sizeof(time_str), "%Y-%m-%d %H:%M:%S", localtime(&now));
    printf("[%s] [%s] ", time_str, level);
    va_list args;
    va_start(args, format);
    vprintf(format, args);
    va_end(args);
}

void sigint_handler(int sig __attribute__((unused))) {
    sigint_count++;
    if (sigint_count == 1) {
        // Set flag for graceful shutdown - main thread will handle logging
        graceful_shutdown = 1;
    } else if (sigint_count >= 2) {
        // Force immediate exit on second SIGINT
        _exit(1); // Use _exit() instead of exit() for signal safety
    }
}

void sighup_handler(int sig __attribute__((unused))) {
    // Set flag for configuration reload - main thread will handle logging
    reload_config = 1;
}

void print_status() {
    if (global_debug) return;

    time_t current_time = time(NULL);
    time_t uptime = current_time - start_time;

    int days = uptime / 86400;
    int hours = (uptime % 86400) / 3600;
    int minutes = (uptime % 3600) / 60;
    int seconds = uptime % 60;

    // Protect access to active_tunnels_count with mutex
    pthread_mutex_lock(&tunnel_mutex);
    int tunnel_count = active_tunnels_count;

    static time_t last_status_time = 0;

    if (last_status_time == 0) {
        last_status_time = current_time;
    }

    time_t time_since_last_status = current_time - last_status_time;

    if (days > 0) {
        log_message("STATUS", "Uptime: %d-%02d:%02d:%02d | Active tunnels: %d\n",
                    days, hours, minutes, seconds, tunnel_count);
    } else {
        log_message("STATUS", "Uptime: %02d:%02d:%02d | Active tunnels: %d\n",
                    hours, minutes, seconds, tunnel_count);
    }

    // List active tunnels with details
    for (int i = 0; i < active_tunnels_count; i++) {
        tunnel_t *tunnel = active_tunnels[i];
        if (tunnel && tunnel->active) {
            if (tunnel->start_time == 0) tunnel->start_time = current_time;
            time_t tunnel_uptime = current_time - tunnel->start_time;
            unsigned long long total_data = tunnel->total_bytes_sent + tunnel->total_bytes_received;
            double avg_rate = tunnel_uptime > 0 ? (double)total_data / tunnel_uptime : 0.0;

            // Calculate last rate based on data since last status
            static unsigned long long last_total_data[100] = {0}; // Simple array to track last data
            unsigned long long data_since_last = total_data - last_total_data[i];
            double last_rate = time_since_last_status > 0 ? (double)data_since_last / time_since_last_status : 0.0;
            last_total_data[i] = total_data;

            int tunnel_days = tunnel_uptime / 86400;
            int tunnel_hours = (tunnel_uptime % 86400) / 3600;
            int tunnel_minutes = (tunnel_uptime % 3600) / 60;
            int tunnel_seconds = tunnel_uptime % 60;

            char total_data_str[32], sent_data_str[32], recv_data_str[32], avg_rate_str[32], last_rate_str[32];
            format_bytes(total_data, total_data_str, sizeof(total_data_str));
            format_bytes(tunnel->total_bytes_sent, sent_data_str, sizeof(sent_data_str));
            format_bytes(tunnel->total_bytes_received, recv_data_str, sizeof(recv_data_str));
            format_bytes((unsigned long long)avg_rate, avg_rate_str, sizeof(avg_rate_str));
            format_bytes((unsigned long long)last_rate, last_rate_str, sizeof(last_rate_str));

            if (tunnel_days > 0) {
                log_message("STATUS", "  Tunnel ID: %s | Client: %s | Service: %s | Uptime: %d-%02d:%02d:%02d | Data: %s (%s sent, %s recv) | Rate: %s/s avg, %s/s last\n",
                            tunnel->request_id,
                            strlen(tunnel->client_id) > 0 ? tunnel->client_id : "local",
                            strlen(tunnel->service) > 0 ? tunnel->service : "ssh",
                            tunnel_days, tunnel_hours, tunnel_minutes, tunnel_seconds,
                            total_data_str,
                            sent_data_str,
                            recv_data_str,
                            avg_rate_str,
                            last_rate_str);
            } else {
                log_message("STATUS", "  Tunnel ID: %s | Client: %s | Service: %s | Uptime: %02d:%02d:%02d | Data: %s (%s sent, %s recv) | Rate: %s/s avg, %s/s last\n",
                            tunnel->request_id,
                            strlen(tunnel->client_id) > 0 ? tunnel->client_id : "local",
                            strlen(tunnel->service) > 0 ? tunnel->service : "ssh",
                            tunnel_hours, tunnel_minutes, tunnel_seconds,
                            total_data_str,
                            sent_data_str,
                            recv_data_str,
                            avg_rate_str,
                            last_rate_str);
            }
        }
    }

    last_status_time = current_time;

    pthread_mutex_unlock(&tunnel_mutex);
}


typedef struct {
    char *wssshd_server;
    int wssshd_port;
    char *ssh_host;        // Legacy alias for tunnel_host
    int ssh_port;          // Legacy alias for tunnel_port
    char *tunnel_host;
    int tunnel_port;
    char *client_id;
    char *password;
    int interval;
    int debug;
    char *tunnel;
    char *tunnel_control;
    char *wssshd_private_ip;
    char *services_path;
    char *default_service; // Default service name (default: "ssh")
} wssshc_config_t;


void load_config_file(const char *config_path, wssshc_config_t *config) {
    FILE *file = fopen(config_path, "r");
    if (!file) return;

    char line[256];
    int in_section = 0;
    while (fgets(line, sizeof(line), file)) {
        line[strcspn(line, "\n")] = 0;
        if (line[0] == '[' && strstr(line, "wssshc")) {
            in_section = 1;
        } else if (line[0] == '[') {
            in_section = 0;
        } else if (in_section && strchr(line, '=')) {
            char *key = strtok(line, "=");
            char *value = strtok(NULL, "=");
            if (key && value) {
                while (*key == ' ') key++;
                char *end = key + strlen(key) - 1;
                while (end > key && *end == ' ') *end-- = 0;
                while (*value == ' ') value++;
                end = value + strlen(value) - 1;
                while (end > value && *end == ' ') *end-- = 0;
                if (strcmp(key, "password") == 0 && !config->password) {
                    config->password = strdup(value);
                } else if (strcmp(key, "wssshd-server") == 0 && !config->wssshd_server) {
                    config->wssshd_server = strdup(value);
                } else if (strcmp(key, "domain") == 0 && !config->wssshd_server) {
                    config->wssshd_server = strdup(value);
                } else if (strcmp(key, "wssshd-port") == 0) {
                    config->wssshd_port = atoi(value);
                } else if (strcmp(key, "port") == 0 && config->wssshd_port == 9898) {
                    config->wssshd_port = atoi(value);
                } else if (strcmp(key, "ssh-host") == 0 && !config->ssh_host) {
                    config->ssh_host = strdup(value);
                } else if (strcmp(key, "ssh-port") == 0) {
                    config->ssh_port = atoi(value);
                } else if (strcmp(key, "id") == 0 && !config->client_id) {
                    config->client_id = strdup(value);
                } else if (strcmp(key, "interval") == 0) {
                    config->interval = atoi(value);
                } else if (strcmp(key, "services-path") == 0) {
                    // Handle multiple services-path entries
                    if (config->services_path) {
                        char *new_path = malloc(strlen(config->services_path) + strlen(value) + 2);
                        if (new_path) {
                            sprintf(new_path, "%s:%s", config->services_path, value);
                            free(config->services_path);
                            config->services_path = new_path;
                        }
                    } else {
                        config->services_path = strdup(value);
                    }
                } else if (strcmp(key, "tunnel-host") == 0 && !config->tunnel_host) {
                    config->tunnel_host = strdup(value);
                } else if (strcmp(key, "tunnel-port") == 0) {
                    config->tunnel_port = atoi(value);
                } else if (strcmp(key, "service") == 0 && !config->default_service) {
                    config->default_service = strdup(value);
                }
            }
        }
    }
    fclose(file);
}

void load_config(wssshc_config_t *config) {
    // Configuration hierarchy:
    // 1. System config (/etc/wssshc.conf) - lowest priority
    // 2. User config (~/.config/wsssh/wssshc.conf) - medium priority
    // 3. Command line arguments - highest priority (handled in parse_args)

    // Load system config first (lowest priority)
    load_config_file("/etc/wssshc.conf", config);

    // Load user config (medium priority, overrides system)
    char *home = getenv("HOME");
    if (home) {
        char user_config_path[256];
        snprintf(user_config_path, sizeof(user_config_path), "%s/.config/wsssh/wssshc.conf", home);
        load_config_file(user_config_path, config);
    }
}




// Load a single service configuration file
service_config_t *load_service_config(const char *config_path) {
    FILE *file = fopen(config_path, "r");
    if (!file) return NULL;

    service_config_t *service = calloc(1, sizeof(service_config_t));
    if (!service) {
        fclose(file);
        return NULL;
    }

    char line[256];
    char current_section[256] = "";
    int in_service_section = 0;
    int active_parsed = 0; // Track if active was explicitly parsed

    while (fgets(line, sizeof(line), file)) {
        line[strcspn(line, "\n")] = 0;
        if (line[0] == '[' && line[strlen(line) - 1] == ']') {
            // Section header
            strncpy(current_section, line + 1, sizeof(current_section) - 1);
            current_section[strlen(current_section) - 1] = '\0'; // Remove closing bracket
            in_service_section = 1;
            if (!service->name) {
                service->name = strdup(current_section);
            }
        } else if (in_service_section && strchr(line, '=')) {
            char *key = strtok(line, "=");
            char *value = strtok(NULL, "=");
            if (key && value) {
                while (*key == ' ') key++;
                char *end = key + strlen(key) - 1;
                while (end > key && *end == ' ') *end-- = 0;
                while (*value == ' ') value++;
                end = value + strlen(value) - 1;
                while (end > value && *end == ' ') *end-- = 0;

                if (strcmp(key, "tunnel-port") == 0) {
                    service->tunnel_port = atoi(value);
                } else if (strcmp(key, "tunnel-host") == 0 && !service->tunnel_host) {
                    service->tunnel_host = strdup(value);
                } else if (strcmp(key, "command") == 0 && !service->command) {
                    service->command = strdup(value);
                } else if (strcmp(key, "proto") == 0 && !service->proto) {
                    service->proto = strdup(value);
                } else if (strcmp(key, "active") == 0) {
                    service->active = (strcmp(value, "true") == 0 || strcmp(value, "1") == 0 || strcmp(value, "yes") == 0);
                    active_parsed = 1;
                }
            }
        }
    }
    fclose(file);

    // Set defaults
    if (!service->proto) {
        service->proto = strdup("tcp");
    }
    // active defaults to true if not specified
    if (!active_parsed) {
        service->active = 1; // Default to true
    }

    return service;
}

// Free a single service configuration
void free_service_config(service_config_t *service) {
    if (service) {
        free(service->name);
        free(service->tunnel_host);
        free(service->command);
        free(service->proto);
        free(service);
    }
}

// Free service configurations
void free_services_config(service_config_t **services, int num_services) {
    if (!services) return;
    for (int i = 0; i < num_services; i++) {
        if (services[i]) {
            free_service_config(services[i]);
        }
    }
    free(services);
}

// Load all service configurations from multiple directories (colon-separated)
service_config_t **load_services_config(const char *services_path, int *num_services) {
    service_config_t **services = NULL;
    *num_services = 0;

    if (!services_path) return NULL;

    // Handle multiple directories separated by colons
    char *path_copy = strdup(services_path);
    if (!path_copy) return NULL;

    char *dir_token = strtok(path_copy, ":");
    while (dir_token != NULL) {
        DIR *dir = opendir(dir_token);
        if (dir) {
            struct dirent *entry;
            while ((entry = readdir(dir)) != NULL) {
                // Check if file ends with .conf
                size_t len = strlen(entry->d_name);
                if (len > 5 && strcmp(entry->d_name + len - 5, ".conf") == 0) {
                    // Build full path
                    char full_path[1024];
                    snprintf(full_path, sizeof(full_path), "%s/%s", dir_token, entry->d_name);

                    // Check if file is readable
                    if (access(full_path, R_OK) == 0) {
                        service_config_t *service = load_service_config(full_path);
                        if (service && service->name && service->active) {
                            // Check for duplicate service names (later definitions override earlier ones)
                            int duplicate = 0;
                            for (int i = 0; i < *num_services; i++) {
                                if (strcmp(services[i]->name, service->name) == 0) {
                                    // Free the old service and replace it
                                    free_service_config(services[i]);
                                    services[i] = service;
                                    duplicate = 1;
                                    break;
                                }
                            }

                            if (!duplicate) {
                                // Add new service to array
                                services = realloc(services, (*num_services + 1) * sizeof(service_config_t *));
                                if (services) {
                                    services[*num_services] = service;
                                    (*num_services)++;
                                } else {
                                    // Free service if realloc failed
                                    free_service_config(service);
                                }
                            }
                        } else if (service && service->name && !service->active) {
                            // Service is inactive, skip it
                            free_service_config(service);
                        }
                    }
                }
            }
            closedir(dir);
        }
        dir_token = strtok(NULL, ":");
    }

    free(path_copy);
    return services;
}

// Find a service by name
service_config_t *find_service(service_config_t **services, int num_services, const char *service_name) {
    for (int i = 0; i < num_services; i++) {
        if (services[i] && strcmp(services[i]->name, service_name) == 0) {
            return services[i];
        }
    }
    return NULL;
}

// Reload configuration (services and password)
void reload_configuration(wssshc_config_t *config, service_config_t ***services, int *num_services) {
    fprintf(stderr, "Reloading configuration...\n");

    // Reload password from config files
    wssshc_config_t temp_config = *config; // Copy current config
    load_config(&temp_config); // Reload from files

    // Update password if it changed
    if (temp_config.password && (!config->password || strcmp(temp_config.password, config->password) != 0)) {
        fprintf(stderr, "Password updated from configuration\n");
        if (config->password) free(config->password);
        config->password = temp_config.password;
        temp_config.password = NULL; // Don't free it
    }

    // Reload services configuration
    if (config->services_path) {
        // Free existing services
        if (*services) {
            free_services_config(*services, *num_services);
            *services = NULL;
            *num_services = 0;
        }

        // Reload services
        *services = load_services_config(config->services_path, num_services);
        if (*services) {
            fprintf(stderr, "Reloaded %d service configurations\n", *num_services);
        } else {
            fprintf(stderr, "No service configurations found\n");
        }
    }

    // Free temporary config fields that weren't used
    if (temp_config.wssshd_server) free(temp_config.wssshd_server);
    if (temp_config.ssh_host) free(temp_config.ssh_host);
    if (temp_config.tunnel_host) free(temp_config.tunnel_host);
    if (temp_config.client_id) free(temp_config.client_id);
    if (temp_config.tunnel) free(temp_config.tunnel);
    if (temp_config.tunnel_control) free(temp_config.tunnel_control);
    if (temp_config.wssshd_private_ip) free(temp_config.wssshd_private_ip);
    if (temp_config.services_path) free(temp_config.services_path);
    if (temp_config.default_service) free(temp_config.default_service);

    fprintf(stderr, "Configuration reload complete\n");
}


// Execute a command before tunnel connection
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;
    }
}
static void wssshc_print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options]\n", program_name);
    fprintf(stderr, "WSSSH Client - Register with wssshd server\n\n");
    fprintf(stderr, "Protect the dolls!\n\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  --config FILE        Configuration file path (overrides default hierarchy)\n");
    fprintf(stderr, "  --wssshd-server HOST WSSSHD server host (default: mbeted.nexlab.net)\n");
    fprintf(stderr, "  --wssshd-port PORT   WSSSHD server port (default: 9898)\n");
    fprintf(stderr, "  --ssh-host HOST      SSH host to forward tunnel data to (default: 127.0.0.1) [legacy]\n");
    fprintf(stderr, "  --ssh-port PORT      SSH port to forward tunnel data to (default: 22) [legacy]\n");
    fprintf(stderr, "  --tunnel-host HOST   Tunnel host to forward tunnel data to (default: 127.0.0.1)\n");
    fprintf(stderr, "  --tunnel-port PORT   Tunnel port to forward tunnel data to (default: 22)\n");
    fprintf(stderr, "  --id ID              Client identifier\n");
    fprintf(stderr, "  --password PASS      Registration password\n");
    fprintf(stderr, "  --interval SEC       Reconnection interval (default: 30)\n");
    fprintf(stderr, "  --tunnel TYPES       Transport types for data channel (comma-separated or 'any', default: any)\n");
    fprintf(stderr, "  --tunnel-control TYPES Transport types for control channel (comma-separated or 'any', default: any)\n");
    fprintf(stderr, "  --wssshd-private-ip IP Private IP of wssshd server\n");
    fprintf(stderr, "  --services-path DIR   Directory containing service configuration files (default: /etc/wsssh.d/)\n");
    fprintf(stderr, "  --service NAME       Default service name (default: ssh)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "\nConfiguration hierarchy (highest priority first):\n");
    fprintf(stderr, "  1. Command line arguments\n");
    fprintf(stderr, "  2. ~/.config/wsssh/wssshc.conf (user config)\n");
    fprintf(stderr, "  3. /etc/wssshc.conf (system config)\n");
    fprintf(stderr, "\nDonations:\n");
    fprintf(stderr, "  BTC: bc1q3zlkpu95amtcltsk85y0eacyzzk29v68tgc5hx\n");
    fprintf(stderr, "  ETH: 0xdA6dAb526515b5cb556d20269207D43fcc760E51\n");
}

int parse_args(int argc, char *argv[], wssshc_config_t *config) {
    static struct option long_options[] = {
        {"config", required_argument, 0, 'c'},
        {"wssshd-server", required_argument, 0, 's'},
        {"wssshd-port", required_argument, 0, 'p'},
        {"ssh-host", required_argument, 0, 'H'},        // Legacy alias
        {"ssh-port", required_argument, 0, 'P'},        // Legacy alias
        {"tunnel-host", required_argument, 0, 1000},    // New option
        {"tunnel-port", required_argument, 0, 1001},    // New option
        {"id", required_argument, 0, 'i'},
        {"password", required_argument, 0, 'w'},
        {"interval", required_argument, 0, 't'},
        {"tunnel", required_argument, 0, 'T'},
        {"tunnel-control", required_argument, 0, 'C'},
        {"wssshd-private-ip", required_argument, 0, 'I'},
        {"services-path", required_argument, 0, 'S'},
        {"service", required_argument, 0, 1002},        // New option
        {"debug", no_argument, 0, 'd'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int opt;
    char *custom_config = NULL;

    while ((opt = getopt_long(argc, argv, "c:s:p:H:P:i:w:t:T:C:I:S:dh", long_options, NULL)) != -1) {
        switch (opt) {
            case 'c':
                custom_config = optarg;
                break;
            case 's':
                if (config->wssshd_server) free(config->wssshd_server);
                config->wssshd_server = strdup(optarg);
                break;
            case 'p':
                config->wssshd_port = atoi(optarg);
                break;
            case 'H':
                if (config->ssh_host) free(config->ssh_host);
                config->ssh_host = strdup(optarg);
                break;
            case 'P':
                config->ssh_port = atoi(optarg);
                break;
            case 'i':
                if (config->client_id) free(config->client_id);
                config->client_id = strdup(optarg);
                break;
            case 'w':
                if (config->password) free(config->password);
                config->password = strdup(optarg);
                break;
            case 't':
                config->interval = atoi(optarg);
                break;
            case 'T':
                if (config->tunnel) free(config->tunnel);
                config->tunnel = strdup(optarg);
                break;
            case 'C':
                if (config->tunnel_control) free(config->tunnel_control);
                config->tunnel_control = strdup(optarg);
                break;
            case 'I':
                if (config->wssshd_private_ip) free(config->wssshd_private_ip);
                config->wssshd_private_ip = strdup(optarg);
                break;
            case 'S':
                // Handle multiple services-path options
                if (config->services_path) {
                    char *new_path = malloc(strlen(config->services_path) + strlen(optarg) + 2); // +2 for ':' and '\0'
                    if (new_path) {
                        sprintf(new_path, "%s:%s", config->services_path, optarg);
                        free(config->services_path);
                        config->services_path = new_path;
                    }
                } else {
                    config->services_path = strdup(optarg);
                }
                break;
            case 1000: // --tunnel-host
                if (config->tunnel_host) free(config->tunnel_host);
                config->tunnel_host = strdup(optarg);
                break;
            case 1001: // --tunnel-port
                config->tunnel_port = atoi(optarg);
                break;
            case 1002: // --service
                if (config->default_service) free(config->default_service);
                config->default_service = strdup(optarg);
                break;
            case 'd':
                config->debug = 1;
                break;
            case 'h':
            default:
                wssshc_print_usage(argv[0]);
                return 0;
        }
    }

    // Handle --help
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            wssshc_print_usage(argv[0]);
            return 0;
        }
    }

    // If --config is specified, load that file instead of the hierarchy
    if (custom_config) {
        load_config_file(custom_config, config);
    }

    return 1;
}





int connect_to_server(const wssshc_config_t *config, service_config_t ***services, int *num_services) {
    struct sockaddr_in server_addr;
    struct hostent *he;
    int sock;
    SSL_CTX *ssl_ctx;
    SSL *ssl;
    char buffer[BUFFER_SIZE];
    int bytes_read;

    // Reset global shutdown flag for new connection
    global_shutdown = 0;

    // Clean up any leftover tunnel state from previous connection
    cleanup_tunnel(config->debug);

    // Resolve hostname
    if ((he = gethostbyname(config->wssshd_server)) == 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(config->wssshd_port);
    server_addr.sin_addr = *((struct in_addr *)he->h_addr);

    // Connect to server
    if (config->debug) {
        printf("[DEBUG] Attempting to connect to %s:%d...\n", config->wssshd_server, config->wssshd_port);
        fflush(stdout);
    }
    if (connect(sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("Connection failed");
        fprintf(stderr, "Unable to connect to server %s:%d\n", config->wssshd_server, config->wssshd_port);
        fprintf(stderr, "Please check:\n");
        fprintf(stderr, "  1. Server is running\n");
        fprintf(stderr, "  2. Server address and port are correct\n");
        fprintf(stderr, "  3. Network connectivity\n");
        fprintf(stderr, "  4. Firewall settings\n");
        close(sock);
        return 1;
    }
    if (config->debug) {
        printf("[DEBUG] TCP connection established\n");
        fflush(stdout);
    }

    // Initialize SSL
    if (config->debug) {
        printf("[DEBUG] Creating SSL context...\n");
        fflush(stdout);
    }
    ssl_ctx = create_ssl_context();
    if (!ssl_ctx) {
        fprintf(stderr, "Failed to create SSL context\n");
        close(sock);
        return 1;
    }

    if (config->debug) {
        printf("[DEBUG] Creating SSL connection...\n");
        fflush(stdout);
    }
    ssl = SSL_new(ssl_ctx);
    if (!ssl) {
        fprintf(stderr, "Failed to create SSL connection\n");
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }

    SSL_set_fd(ssl, sock);

    if (config->debug) {
        printf("[DEBUG] Performing SSL handshake...\n");
        fflush(stdout);
    }
    int ssl_connect_result = SSL_connect(ssl);
    if (ssl_connect_result <= 0) {
        int ssl_error = SSL_get_error(ssl, ssl_connect_result);
        fprintf(stderr, "SSL connect failed with error %d\n", ssl_error);
        ERR_print_errors_fp(stderr);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }
    if (config->debug) {
        printf("[DEBUG] SSL handshake successful\n");
        fflush(stdout);
    }

    // Perform WebSocket handshake
    if (config->debug) {
        printf("[DEBUG] Performing WebSocket handshake...\n");
        fflush(stdout);
    }
    // websocket_handshake already uses SSL mutex internally
    if (!websocket_handshake(ssl, config->wssshd_server, config->wssshd_port, "/", config->debug)) {
        fprintf(stderr, "WebSocket handshake failed\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }
    if (config->debug) {
        printf("[DEBUG] WebSocket handshake successful\n");
        fflush(stdout);
    }

    // WebSocket handshake and SSL connection are sufficient connectivity tests
    // Skip additional ping test as server may not handle it properly

    // Expand transport lists if "any" is specified
    char *expanded_tunnel = expand_transport_list(config->tunnel, 0); // 0 for data channel
    char *expanded_tunnel_control = expand_transport_list(config->tunnel_control, 1); // 1 for control channel

    // 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);

    // Send registration message with services
    if (config->debug) {
        printf("[DEBUG] Sending registration message with services...\n");
        fflush(stdout);
    }
    int reg_result = send_registration_message_with_services(ssl, config->client_id, config->password, best_tunnel ? best_tunnel : expanded_tunnel, best_tunnel_control ? best_tunnel_control : expanded_tunnel_control, config->wssshd_private_ip, services ? *services : NULL, services ? *num_services : 0);
    if (config->debug && reg_result) {
        // Note: send_registration_message doesn't return the message content, so we can't log the exact JSON here
        // The function internally handles the JSON formatting
        printf("[DEBUG] Registration message sent successfully\n");
        fflush(stdout);
    }

    // Free expanded transport lists and best transport selections
    if (expanded_tunnel != config->tunnel) {
        free(expanded_tunnel);
    }
    if (expanded_tunnel_control != config->tunnel_control) {
        free(expanded_tunnel_control);
    }
    if (best_tunnel) free(best_tunnel);
    if (best_tunnel_control) free(best_tunnel_control);
    if (!reg_result) {
        fprintf(stderr, "Failed to send registration message\n");
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 1;
    }
    if (config->debug) {
        printf("[DEBUG] Registration message sent successfully\n");
        fflush(stdout);
    }

    // Read WebSocket frame with registration response
    if (config->debug) {
        printf("[DEBUG] Waiting for registration response...\n");
        fflush(stdout);
    }

    int sock_fd = SSL_get_fd(ssl);
    fd_set readfds;
    struct timeval tv;

    // Wait for registration response with timeout and SIGINT checking
    time_t start_time_reg = time(NULL);
    int timeout_seconds = 30; // Increased timeout for better reliability

    while (1) {
        // Check for graceful shutdown more frequently
        if (graceful_shutdown) {
            if (config->debug) {
                fprintf(stderr, "[DEBUG] Graceful shutdown requested during registration, exiting...\n");
                fflush(stderr);
            }
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return 1;
        }

        // Check timeout
        if (time(NULL) - start_time_reg > timeout_seconds) {
            fprintf(stderr, "Timeout waiting for registration response after %d seconds\n", timeout_seconds);
            fprintf(stderr, "This may indicate:\n");
            fprintf(stderr, "  1. Server is not running or not responding\n");
            fprintf(stderr, "  2. Network connectivity issues\n");
            fprintf(stderr, "  3. Server rejected the registration\n");
            fprintf(stderr, "  4. Firewall blocking the connection\n");
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return 1;
        }

        // Set up select with longer timeout for better reliability
        FD_ZERO(&readfds);
        FD_SET(sock_fd, &readfds);
        tv.tv_sec = 1;  // 1 second timeout
        tv.tv_usec = 0;

        int select_result = select(sock_fd + 1, &readfds, NULL, NULL, &tv);
        if (select_result == -1) {
            if (errno == EINTR) {
                // Interrupted by signal, continue and check SIGINT at top of loop
                continue;
            }
            fprintf(stderr, "Select error while waiting for registration response: %s\n", strerror(errno));
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return 1;
        } else if (select_result == 0) {
            // Timeout, continue loop
            if (config->debug) {
                printf("[DEBUG] Waiting for registration response... (%ld seconds elapsed)\n",
                       time(NULL) - start_time_reg);
                fflush(stdout);
            }
            // Also print a message every 5 seconds to show we're still waiting
            static time_t last_progress_time = 0;
            time_t current_time = time(NULL);
            if (current_time - last_progress_time >= 5) {
                printf("[DEBUG] Still waiting for server response... (%ld seconds)\n",
                       current_time - start_time_reg);
                fflush(stdout);
                last_progress_time = current_time;
            }
            continue;
        }

        // Data available, try to read
        if (FD_ISSET(sock_fd, &readfds)) {
            // Set socket to non-blocking mode temporarily for SSL_read
            int flags = fcntl(sock_fd, F_GETFL, 0);
            if (flags != -1) {
                fcntl(sock_fd, F_SETFL, flags | O_NONBLOCK);
            }

            bytes_read = SSL_read(ssl, buffer, sizeof(buffer));

            // Restore blocking mode
            if (flags != -1) {
                fcntl(sock_fd, F_SETFL, flags);
            }

            if (bytes_read > 0) {
                if (config->debug) {
                    printf("[DEBUG] SSL_read returned %d bytes (raw WebSocket frame)\n", bytes_read);
                    // Don't print raw frame data as it contains binary headers
                    fflush(stdout);
                }
                break; // Success
            } else if (bytes_read == 0) {
                // Connection closed
                fprintf(stderr, "Connection closed while waiting for registration response\n");
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                return 1;
            } else {
                // Check SSL error
                int ssl_error = SSL_get_error(ssl, bytes_read);
                if (ssl_error == SSL_ERROR_WANT_READ || ssl_error == SSL_ERROR_WANT_WRITE) {
                    // Need more data, continue
                    if (config->debug) {
                        printf("[DEBUG] SSL_read wants more data, continuing...\n");
                        fflush(stdout);
                    }
                    usleep(10000); // Small delay before retry
                    continue;
                } else if (ssl_error == SSL_ERROR_SYSCALL && errno == EAGAIN) {
                    // Non-blocking read would block, continue
                    if (config->debug) {
                        printf("[DEBUG] SSL_read would block, continuing...\n");
                        fflush(stdout);
                    }
                    continue;
                } else {
                    // Real error
                    char error_buf[256];
                    ERR_error_string_n(ssl_error, error_buf, sizeof(error_buf));
                    fprintf(stderr, "SSL error while reading registration response: %d (%s)\n", ssl_error, error_buf);
                    SSL_free(ssl);
                    SSL_CTX_free(ssl_ctx);
                    close(sock);
                    return 1;
                }
            }
        }
    }

    // Parse WebSocket frame with extended length support
    if (bytes_read >= 2 && ((buffer[0] & 0x8F) == 0x81 || (buffer[0] & 0x8F) == 0x82)) { // Text or binary frame
        int masked = buffer[1] & 0x80;
        int len_indicator = buffer[1] & 0x7F;
        int header_len = 2;
        size_t payload_len;

        if (config->debug) {
            printf("[DEBUG] Initial len_indicator = %d, masked = %d\n", len_indicator, masked);
            fflush(stdout);
        }

        if (len_indicator <= 125) {
            payload_len = len_indicator;
        } else if (len_indicator == 126) {
            if (bytes_read < 4) {
                fprintf(stderr, "Incomplete extended length frame\n");
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                return 1;
            }
            payload_len = ((unsigned char)buffer[2] << 8) | (unsigned char)buffer[3];
            header_len = 4;
        } else if (len_indicator == 127) {
            if (bytes_read < 10) {
                fprintf(stderr, "Incomplete extended length frame\n");
                SSL_free(ssl);
                SSL_CTX_free(ssl_ctx);
                close(sock);
                return 1;
            }
            uint64_t len = 0;
            for (int i = 0; i < 8; i++) {
                len = (len << 8) | (unsigned char)buffer[2 + i];
            }
            payload_len = (int)len;
            header_len = 10;
        }

        if (masked) {
            header_len += 4;
        }

        if ((size_t)bytes_read >= (size_t)header_len + payload_len) {
            char *payload = buffer + header_len;
            if (masked) {
                // Unmask the payload
                char *mask_key = buffer + header_len - 4;
                for (size_t i = 0; i < payload_len; i++) {
                    payload[i] ^= mask_key[i % 4];
                }
            }
            // Copy payload to buffer start
            memmove(buffer, payload, payload_len);
            buffer[payload_len] = '\0';
        } else {
            fprintf(stderr, "WebSocket frame incomplete: need %zu bytes, got %d\n", payload_len + header_len, bytes_read);
            SSL_free(ssl);
            SSL_CTX_free(ssl_ctx);
            close(sock);
            return 1;
        }
    } else {
        buffer[bytes_read] = '\0';
        fprintf(stderr, "Unexpected frame: 0x%02x 0x%02x\n", buffer[0], buffer[1]);
    }

    printf("[DEBUG] Registration response received: %s\n", buffer);

    if (strstr(buffer, "registered") == NULL) {
        fprintf(stderr, "Registration failed: %s\n", buffer);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 2;
    }

    if (config->debug) {
        printf("Connected and registered as %s\n", config->client_id);
    } else {
        printf("[EVENT] Connected and registered as %s\n", config->client_id);
    }

    // Keep connection alive and handle tunnel requests
    // active_tunnels = NULL; // Will implement tunnel handling

    // Frame accumulation buffer for handling partial WebSocket frames
    static char frame_buffer[BUFFER_SIZE * 4]; // Quadruple buffer size for accumulation
    static int frame_buffer_used = 0;

    // Keep-alive ping variables
    time_t last_ping_time = time(NULL);
    const int ping_interval = 15; // Send ping every 15 seconds

    // Keep-alive message variables
    time_t last_keepalive_time = time(NULL);
    const int keepalive_interval = 30; // Send keep-alive every 30 seconds

    // Tunnel list debug output variables
    time_t last_tunnel_list_time = time(NULL);
    const int tunnel_list_interval = 60; // Show tunnel list every 60 seconds

    // Status printing variables
    static time_t last_status_time = 0;

    while (1) {
        // Print status every 15 seconds
        if (time(NULL) - last_status_time >= 15) {
            print_status();
            last_status_time = time(NULL);
        }

        // Check for graceful shutdown
        if (graceful_shutdown) {
            fprintf(stderr, "Attempting graceful shutdown of active tunnels...\n");

            // Send tunnel_close messages for all active tunnels
            pthread_mutex_lock(&tunnel_mutex);
            int tunnels_closed = 0;
            for (int i = 0; i < active_tunnels_count; i++) {
                tunnel_t *tunnel = active_tunnels[i];
                if (tunnel && tunnel->active && tunnel->ssl) {
                    if (config->debug) {
                        printf("[DEBUG] Sending tunnel_close for tunnel %s\n", tunnel->request_id);
                        fflush(stdout);
                    }
                    send_tunnel_close(tunnel->ssl, tunnel->request_id, config->debug);
                    tunnels_closed++;
                }
            }
            pthread_mutex_unlock(&tunnel_mutex);

            if (tunnels_closed > 0) {
                fprintf(stderr, "Sent close messages for %d tunnels, waiting up to 3 seconds...\n", tunnels_closed);
                // Wait up to 3 seconds for tunnel close operations to complete
                time_t shutdown_start = time(NULL);
                while (time(NULL) - shutdown_start < 3) {
                    // Check if all tunnels are closed
                    pthread_mutex_lock(&tunnel_mutex);
                    int active_count = 0;
                    for (int i = 0; i < active_tunnels_count; i++) {
                        if (active_tunnels[i] && active_tunnels[i]->active) {
                            active_count++;
                        }
                    }
                    pthread_mutex_unlock(&tunnel_mutex);

                    if (active_count == 0) {
                        fprintf(stderr, "All tunnels closed gracefully\n");
                        break;
                    }

                    usleep(100000); // Sleep 100ms before checking again
                }
            }

            break;
        }

        // Send keep-alive ping if needed
        time_t ping_time = time(NULL);
        if (ping_time - last_ping_time >= ping_interval) {
            if (config->debug) {
                printf("[DEBUG - WebSockets] Sending keep-alive ping\n");
                fflush(stdout);
            }
            // Send ping with empty payload
            if (!send_ping_frame(ssl, "", 0)) {
                if (config->debug) {
                    printf("[DEBUG - WebSockets] Failed to send keep-alive ping\n");
                    fflush(stdout);
                }
            } else {
                last_ping_time = ping_time;
            }
        }

        // Send tunnel keep-alive messages if needed
        time_t current_time = time(NULL);
        if (current_time - last_keepalive_time >= keepalive_interval) {
            // Send keep-alive for all active tunnels
            pthread_mutex_lock(&tunnel_mutex);
            for (int i = 0; i < active_tunnels_count; i++) {
                tunnel_t *tunnel = active_tunnels[i];
                if (tunnel && tunnel->active) {
                    send_tunnel_keepalive(ssl, tunnel, config->debug);
                }
            }
            pthread_mutex_unlock(&tunnel_mutex);
            last_keepalive_time = current_time;
        }

        // Check for keep-alive timeouts
        check_keepalive_timeouts(config->debug);

        // Show tunnel list periodically in debug mode
        if (config->debug && current_time - last_tunnel_list_time >= tunnel_list_interval) {
            printf("[DEBUG - Tunnel] Active tunnels list:\n");
            pthread_mutex_lock(&tunnel_mutex);
            if (active_tunnels_count == 0) {
                printf("[DEBUG - Tunnel]   No active tunnels\n");
            } else {
                for (int i = 0; i < active_tunnels_count; i++) {
                    tunnel_t *tunnel = active_tunnels[i];
                    if (tunnel && tunnel->active) {
                        // 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);
                        }

                        printf("[DEBUG - Tunnel]   Tunnel %s:\n", tunnel->request_id);
                        printf("[DEBUG - Tunnel]     Total sent: %llu bytes\n", tunnel->total_bytes_sent);
                        printf("[DEBUG - Tunnel]     Total received: %llu bytes\n", tunnel->total_bytes_received);
                        printf("[DEBUG - Tunnel]     Current rate: %.2f B/s\n", rate_bps);
                        printf("[DEBUG - Tunnel]     Last keep-alive sent: %ld seconds ago\n", current_time - tunnel->last_keepalive_sent);
                        printf("[DEBUG - Tunnel]     Last keep-alive received: %ld seconds ago\n", current_time - tunnel->last_keepalive_received);
                    }
                }
            }
            pthread_mutex_unlock(&tunnel_mutex);
            printf("[DEBUG - Tunnel] End of tunnel list\n");
            fflush(stdout);
            last_tunnel_list_time = current_time;
        }

        // Always try to read more data if there's space, even if we have a complete frame
        // This ensures we can accumulate data for very large frames
        if ((size_t)frame_buffer_used < sizeof(frame_buffer)) {
            // Validate SSL connection state before reading
            if (SSL_get_shutdown(ssl) & SSL_RECEIVED_SHUTDOWN) {
                if (config->debug) {
                    printf("[DEBUG - WebSockets] SSL connection has received shutdown\n");
                    fflush(stdout);
                }
                // Only clean up if we're actually exiting, not on transient shutdown
                if (!sigint_received) {
                    cleanup_tunnel(config->debug);
                }
                break;
            }

            // Set up timeout for SSL read operation
            fd_set readfds;
            struct timeval tv;
            int sock_fd = SSL_get_fd(ssl);

            FD_ZERO(&readfds);
            FD_SET(sock_fd, &readfds);
            tv.tv_sec = 1;  // 1 second timeout for faster SIGINT response
            tv.tv_usec = 0;

            int select_result = select(sock_fd + 1, &readfds, NULL, NULL, &tv);
            if (select_result == -1) {
                if (errno == EINTR) {
                    // Interrupted by signal, continue and check SIGINT at top of loop
                    continue;
                }
                if (config->debug) {
                    perror("[DEBUG - WebSockets] select failed");
                    fflush(stdout);
                }
                // Don't clean up all tunnels on select failure - this is not a fatal error
                break;
            } else if (select_result == 0) {
                continue; // Timeout, try again
            }

            // Set socket to non-blocking mode temporarily for SSL_read
            int flags = fcntl(sock_fd, F_GETFL, 0);
            if (flags != -1) {
                fcntl(sock_fd, F_SETFL, flags | O_NONBLOCK);
            }

            bytes_read = SSL_read(ssl, frame_buffer + frame_buffer_used, sizeof(frame_buffer) - frame_buffer_used);

            // Restore blocking mode
            if (flags != -1) {
                fcntl(sock_fd, F_SETFL, flags);
            }
            if (bytes_read <= 0) {
                if (bytes_read < 0) {
                    int ssl_error = SSL_get_error(ssl, bytes_read);
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] SSL read error: %d\n", ssl_error);
                        fflush(stdout);
                    }

                    // Handle transient SSL errors with retry
                    if (ssl_error == SSL_ERROR_WANT_READ || ssl_error == SSL_ERROR_WANT_WRITE) {
                        if (config->debug) {
                            printf("[DEBUG - WebSockets] Transient SSL error, retrying...\n");
                            fflush(stdout);
                        }
                        usleep(10000); // Wait 10ms before retry
                        continue; // Retry the read operation
                    } else if (ssl_error == SSL_ERROR_SYSCALL && errno == EAGAIN) {
                        // Non-blocking read would block, continue
                        if (config->debug) {
                            printf("[DEBUG - WebSockets] SSL_read would block, continuing...\n");
                            fflush(stdout);
                        }
                        continue;
                    }

                    // Print detailed SSL error information for non-transient errors
                    char error_buf[256];
                    ERR_error_string_n(ssl_error, error_buf, sizeof(error_buf));
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] SSL error details: %s\n", error_buf);
                        fflush(stdout);
                    }
                    fprintf(stderr, "SSL read error (%d): %s\n", ssl_error, error_buf);
                } else {
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] Connection closed by server (EOF)\n");
                        fflush(stdout);
                    }
                }
                // Don't clean up all tunnels on SSL errors - let individual tunnels handle their own failures
                // Only clean up if we're actually exiting the main loop
                if (!sigint_received) {
                    cleanup_tunnel(config->debug);
                }
                break;
            }

            frame_buffer_used += bytes_read;

            if (config->debug) {
                printf("[DEBUG - WebSockets] Accumulated %d bytes, frame: 0x%02x 0x%02x 0x%02x 0x%02x\n", frame_buffer_used, frame_buffer[0], frame_buffer[1], frame_buffer[2], frame_buffer[3]);
                fflush(stdout);
            }
        }

        // Try to parse WebSocket frame from accumulated buffer
        char *payload;
        int payload_len;
        if (parse_websocket_frame(frame_buffer, frame_buffer_used, &payload, &payload_len)) {
                // Frame is complete, determine frame type
                unsigned char frame_type = frame_buffer[0] & 0x8F;

                if (frame_type == 0x88) { // Close frame
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] Received close frame - treating as tunnel close, keeping WebSocket open\n");
                        fflush(stdout);
                    }
                    // Treat as tunnel close, don't close WebSocket connection
                    // Note: tunnel_close messages are handled in the message processing below
                    // Continue processing, don't return
                } else if (frame_type == 0x89) { // Ping frame
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] Received ping frame\n");
                        fflush(stdout);
                    }
                    // Send pong with same payload
                    // send_pong_frame already uses SSL mutex internally
                    if (!send_pong_frame(ssl, payload, payload_len)) {
                        if (config->debug) {
                            printf("[DEBUG - WebSockets] Failed to send pong frame\n");
                            fflush(stdout);
                        }
                    }
                } else if (frame_type == 0x8A) { // Pong frame
                    if (config->debug) {
                        printf("[DEBUG - WebSockets] Received pong frame\n");
                        fflush(stdout);
                    }
                    // Just acknowledge, no response needed
                } else if (frame_type == 0x81 || frame_type == 0x82) { // Text or binary frame
                    // Copy payload to buffer for processing (null-terminate for string operations)
                    if ((size_t)payload_len < sizeof(buffer)) {
                        memcpy(buffer, payload, payload_len);
                        buffer[payload_len] = '\0';
                    } else {
                        fprintf(stderr, "Payload too large for processing buffer\n");
                        // Skip this frame
                        frame_buffer_used = 0;
                        continue;
                    }

                    // Check if this is a tunnel_data message to suppress verbose logging
                    int is_tunnel_data = (strstr(buffer, "tunnel_data") != NULL);

                    if (config->debug) {
                        if (is_tunnel_data) {
                            // Extract request_id and data size for tunnel_data messages
                            char *request_id_start = strstr(buffer, "\"request_id\"");
                            char request_id[256] = "...";
                            if (request_id_start) {
                                char *colon = strchr(request_id_start, ':');
                                if (colon) {
                                    char *quote = strchr(colon, '"');
                                    if (quote) {
                                        request_id_start = quote + 1;
                                        char *end_quote = strchr(request_id_start, '"');
                                        if (end_quote) {
                                            size_t len = end_quote - request_id_start;
                                            if (len < sizeof(request_id) - 1) {
                                                memcpy(request_id, request_id_start, len);
                                                request_id[len] = '\0';
                                            }
                                        }
                                    }
                                }
                            }
                            char *data_start = strstr(buffer, "\"data\"");
                            size_t data_size = 0;
                            if (data_start) {
                                char *data_colon = strchr(data_start, ':');
                                if (data_colon) {
                                    char *data_quote = strchr(data_colon, '"');
                                    if (data_quote) {
                                        data_start = data_quote + 1;
                                        char *data_end = strchr(data_start, '"');
                                        if (data_end) {
                                            data_size = data_end - data_start;
                                        }
                                    }
                                }
                            }
                            printf("[DEBUG - WebSockets] Received message: {\"type\":\"tunnel_data\", \"request_id\":\"%s\", \"data\":\"<size: %zu bytes>\"}\n", request_id, data_size);
                        } else if (strstr(buffer, "\"type\":\"tunnel_response\"") || strstr(buffer, "\"type\": \"tunnel_response\"")) {
                            // Extract request_id and data size for tunnel_response messages
                            char *request_id_start = strstr(buffer, "\"request_id\"");
                            char request_id[256] = "...";
                            if (request_id_start) {
                                char *colon = strchr(request_id_start, ':');
                                if (colon) {
                                    char *quote = strchr(colon, '"');
                                    if (quote) {
                                        request_id_start = quote + 1;
                                        char *end_quote = strchr(request_id_start, '"');
                                        if (end_quote) {
                                            size_t len = end_quote - request_id_start;
                                            if (len < sizeof(request_id) - 1) {
                                                memcpy(request_id, request_id_start, len);
                                                request_id[len] = '\0';
                                            }
                                        }
                                    }
                                }
                            }
                            char *data_start = strstr(buffer, "\"data\"");
                            size_t data_size = 0;
                            if (data_start) {
                                char *data_colon = strchr(data_start, ':');
                                if (data_colon) {
                                    char *data_quote = strchr(data_colon, '"');
                                    if (data_quote) {
                                        data_start = data_quote + 1;
                                        char *data_end = strchr(data_start, '"');
                                        if (data_end) {
                                            data_size = data_end - data_start;
                                        }
                                    }
                                }
                            }
                            printf("[DEBUG - WebSockets] Received message: {\"type\":\"tunnel_response\", \"request_id\":\"%s\", \"data\":\"<size: %zu bytes>\"}\n", request_id, data_size);
                        } else {
                            printf("[DEBUG - WebSockets] Received message: %.*s\n", payload_len, payload);
                        }
                        fflush(stdout);
                    }

                    // Handle message
                    if (config->debug) {
                        if (is_tunnel_data) {
                            // Extract request_id and data size for tunnel_data messages
                            char *request_id_start = strstr(buffer, "\"request_id\"");
                            char request_id[256] = "...";
                            if (request_id_start) {
                                char *colon = strchr(request_id_start, ':');
                                if (colon) {
                                    char *quote = strchr(colon, '"');
                                    if (quote) {
                                        request_id_start = quote + 1;
                                        char *end_quote = strchr(request_id_start, '"');
                                        if (end_quote) {
                                            size_t len = end_quote - request_id_start;
                                            if (len < sizeof(request_id) - 1) {
                                                memcpy(request_id, request_id_start, len);
                                                request_id[len] = '\0';
                                            }
                                        }
                                    }
                                }
                            }
                            char *data_start = strstr(buffer, "\"data\"");
                            size_t data_size = 0;
                            if (data_start) {
                                char *data_colon = strchr(data_start, ':');
                                if (data_colon) {
                                    char *data_quote = strchr(data_colon, '"');
                                    if (data_quote) {
                                        data_start = data_quote + 1;
                                        char *data_end = strchr(data_start, '"');
                                        if (data_end) {
                                            data_size = data_end - data_start;
                                        }
                                    }
                                }
                            }
                            printf("[DEBUG - WebSockets] Processing message: {\"type\":\"tunnel_data\", \"request_id\":\"%s\", \"data\":\"<size: %zu bytes>\"}\n", request_id, data_size);
                        } else {
                            printf("[DEBUG - WebSockets] Processing message: %s\n", buffer);
                        }
                        fflush(stdout);
                    }

                    // Parse JSON message
                    if (strstr(buffer, "tunnel_request")) {
                        // Extract request_id, service, and enc
                        char *id_start = strstr(buffer, "\"request_id\"");
                        char *service_start = strstr(buffer, "\"service\"");
                        char *enc_start = strstr(buffer, "\"enc\"");
                        char service_name[256] = "";
                        wsssh_encoding_t encoding = ENCODING_HEX; // Default to hex
                        if (id_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        // Extract service name if present
                                        if (service_start) {
                                            char *service_colon = strchr(service_start, ':');
                                            if (service_colon) {
                                                char *service_quote = strchr(service_colon, '"');
                                                if (service_quote) {
                                                    service_start = service_quote + 1;
                                                    char *service_end = strchr(service_start, '"');
                                                    if (service_end) {
                                                        *service_end = '\0';
                                                        strncpy(service_name, service_start, sizeof(service_name) - 1);
                                                    }
                                                }
                                            }
                                        }
                                        // Extract enc mode if present
                                        if (enc_start) {
                                            char *enc_colon = strchr(enc_start, ':');
                                            if (enc_colon) {
                                                char *enc_quote = strchr(enc_colon, '"');
                                                if (enc_quote) {
                                                    enc_start = enc_quote + 1;
                                                    char *enc_end = strchr(enc_start, '"');
                                                    if (enc_end) {
                                                        *enc_end = '\0';
                                                        if (strcmp(enc_start, "hex") == 0) {
                                                            encoding = ENCODING_HEX;
                                                        } else if (strcmp(enc_start, "base64") == 0) {
                                                            encoding = ENCODING_BASE64;
                                                        } else if (strcmp(enc_start, "bin") == 0) {
                                                            encoding = ENCODING_BINARY;
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                        if (config->debug) {
                                            const char *enc_str = "hex";
                                            if (encoding == ENCODING_BASE64) enc_str = "base64";
                                            else if (encoding == ENCODING_BINARY) enc_str = "bin";
                                            printf("[DEBUG - WebSockets] Received tunnel_request for ID: %s, service: %s, enc: %s\n", id_start, service_name[0] ? service_name : "(default)", enc_str);
                                            fflush(stdout);
                                        } else {
                                            printf("[EVENT] New tunnel request: %s\n", id_start);
                                        }
                                        // Services are now loaded in main function

                                        // Find service configuration
                                        service_config_t *service = NULL;
                                        const char *requested_service = service_name[0] ? service_name : config->default_service;

                                        if (requested_service && *services) {
                                            service = find_service(*services, *num_services, requested_service);
                                        }

                                        if (service) {
                                            if (config->debug) {
                                                printf("[DEBUG] Using service configuration: %s (host: %s, port: %d, proto: %s)\n",
                                                       service->name, service->tunnel_host, service->tunnel_port, service->proto);
                                            }
                                            handle_tunnel_request_with_service_and_enc(ssl, id_start, service, config->debug, encoding);
                                        } else {
                                            if (config->debug) {
                                                printf("[DEBUG] No service configuration found for '%s', using default tunnel settings\n",
                                                       requested_service ? requested_service : "(none)");
                                            }
                                            handle_tunnel_request_with_enc(ssl, id_start, config->debug, config->tunnel_host, config->tunnel_port, encoding);
                                        }
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_data")) {
                        if (config->debug) {
                            printf("[DEBUG - WebSockets] Received tunnel_data message\n");
                            fflush(stdout);
                        }
                        // Extract request_id and data
                        char *id_start = strstr(buffer, "\"request_id\"");
                        char *data_start = strstr(buffer, "\"data\"");
                        if (id_start && data_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        char *data_colon = strchr(data_start, ':');
                                        if (data_colon) {
                                            char *data_quote = strchr(data_colon, '"');
                                            if (data_quote) {
                                                data_start = data_quote + 1;
                                                char *data_end = strchr(data_start, '"');
                                                if (data_end) {
                                                    *data_end = '\0';
                                                    handle_tunnel_data(ssl, id_start, data_start, config->debug);
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_close")) {
                        char *id_start = strstr(buffer, "\"request_id\"");
                        if (id_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        if (config->debug) {
                                            printf("[DEBUG - WebSockets] Received tunnel_close for ID: %s\n", id_start);
                                            fflush(stdout);
                                        } else {
                                            printf("[EVENT] Tunnel closed: %s\n", id_start);
                                        }
                                        handle_tunnel_close(ssl, id_start, config->debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_keepalive")) {
                        // Extract request_id, total_bytes, and rate_bps
                        char *id_start = strstr(buffer, "\"request_id\"");
                        char *total_start = strstr(buffer, "\"total_bytes\"");
                        char *rate_start = strstr(buffer, "\"rate_bps\"");

                        if (id_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';

                                        unsigned long long total_bytes = 0;
                                        double rate_bps = 0.0;

                                        // Parse total_bytes
                                        if (total_start) {
                                            char *total_colon = strchr(total_start, ':');
                                            if (total_colon) {
                                                total_bytes = strtoull(total_colon + 1, NULL, 10);
                                            }
                                        }

                                        // Parse rate_bps
                                        if (rate_start) {
                                            char *rate_colon = strchr(rate_start, ':');
                                            if (rate_colon) {
                                                rate_bps = strtod(rate_colon + 1, NULL);
                                            }
                                        }

                                        handle_tunnel_keepalive(ssl, id_start, total_bytes, rate_bps, config->debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_keepalive_ack")) {
                        // Extract request_id
                        char *id_start = strstr(buffer, "\"request_id\"");
                        if (id_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        handle_tunnel_keepalive_ack(ssl, id_start, config->debug);
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_ack")) {
                        // Extract request_id and frame_id
                        char *id_start = strstr(buffer, "\"request_id\"");
                        char *frame_start = strstr(buffer, "\"frame_id\"");
                        if (id_start && frame_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        // Extract frame_id
                                        char *frame_colon = strchr(frame_start, ':');
                                        if (frame_colon) {
                                            uint32_t frame_id = (uint32_t)atoi(frame_colon + 1);
                                            handle_tunnel_ack(ssl, id_start, frame_id, config->debug);
                                        }
                                    }
                                }
                            }
                        }
                    } else if (strstr(buffer, "tunnel_ko")) {
                        // Extract request_id and frame_id
                        char *id_start = strstr(buffer, "\"request_id\"");
                        char *frame_start = strstr(buffer, "\"frame_id\"");
                        if (id_start && frame_start) {
                            char *colon = strchr(id_start, ':');
                            if (colon) {
                                char *open_quote = strchr(colon, '"');
                                if (open_quote) {
                                    id_start = open_quote + 1;
                                    char *close_quote = strchr(id_start, '"');
                                    if (close_quote) {
                                        *close_quote = '\0';
                                        // Extract frame_id
                                        char *frame_colon = strchr(frame_start, ':');
                                        if (frame_colon) {
                                            uint32_t frame_id = (uint32_t)atoi(frame_colon + 1);
                                            handle_tunnel_ko(ssl, id_start, frame_id, config->debug);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                // Remove processed frame from buffer
                int frame_size = (payload - frame_buffer) + payload_len;
                if (frame_size < frame_buffer_used) {
                    // Move remaining data to front
                    memmove(frame_buffer, frame_buffer + frame_size, frame_buffer_used - frame_size);
                    frame_buffer_used -= frame_size;
                } else {
                    // Entire buffer was consumed
                    frame_buffer_used = 0;
                }
            } else {
                // Frame not complete yet, continue reading
                continue;
            }
    }

    // Cleanup
    SSL_free(ssl);
    SSL_CTX_free(ssl_ctx);
    close(sock);
    return 1;
}


int main(int argc, char *argv[]) {
    wssshc_config_t config = {
        .wssshd_server = NULL,
        .wssshd_port = 9898,
        .ssh_host = NULL,
        .ssh_port = 22,
        .tunnel_host = NULL,
        .tunnel_port = 22,
        .client_id = NULL,
        .password = NULL,
        .interval = 30,
        .debug = 0,
        .tunnel = NULL,
        .tunnel_control = NULL,
        .wssshd_private_ip = NULL,
        .services_path = NULL,
        .default_service = NULL
    };

    // Initialize CPU affinity system for multi-core utilization
    init_cpu_affinity();

    pthread_mutex_init(&tunnel_mutex, NULL);

    // Set up signal handlers
    signal(SIGINT, sigint_handler);
    signal(SIGHUP, sighup_handler);

    // Load configuration files first (system and user configs)
    load_config(&config);

    // Parse command line arguments (highest priority, overrides config files)
    if (!parse_args(argc, argv, &config)) {
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    // Set defaults for optional fields
    if (!config.wssshd_server) {
        config.wssshd_server = strdup("mbeted.nexlab.net");
    }

    // Handle tunnel-host/tunnel-port precedence over ssh-host/ssh-port
    if (config.tunnel_host) {
        // tunnel-host has precedence
        if (config.ssh_host) free(config.ssh_host);
        config.ssh_host = strdup(config.tunnel_host);
    } else if (config.ssh_host) {
        // ssh-host provided, use it for tunnel_host
        config.tunnel_host = strdup(config.ssh_host);
    } else {
        // Default
        config.ssh_host = strdup("127.0.0.1");
        config.tunnel_host = strdup("127.0.0.1");
    }

    if (config.tunnel_port != 22) {
        // tunnel-port was explicitly set
        config.ssh_port = config.tunnel_port;
    } else if (config.ssh_port != 22) {
        // ssh-port was explicitly set
        config.tunnel_port = config.ssh_port;
    } else {
        // Default
        config.ssh_port = 22;
        config.tunnel_port = 22;
    }

    if (!config.tunnel) {
        config.tunnel = strdup("any");
    }
    if (!config.tunnel_control) {
        config.tunnel_control = strdup("any");
    }
    if (!config.services_path) {
        config.services_path = strdup("/etc/wsssh.d/");
    }
    if (!config.default_service) {
        config.default_service = strdup("ssh");
    }

    // Validate required arguments
    if (!config.client_id || !config.password) {
        fprintf(stderr, "Error: --id and --password are required\n");
        wssshc_print_usage(argv[0]);
        if (config.wssshd_server) free(config.wssshd_server);
        if (config.ssh_host) free(config.ssh_host);
        if (config.client_id) free(config.client_id);
        if (config.password) free(config.password);
        pthread_mutex_destroy(&tunnel_mutex);
        return 1;
    }

    global_debug = config.debug;
    start_time = time(NULL);

    // Load services configuration
    service_config_t **services = NULL;
    int num_services = 0;
    if (config.services_path) {
        services = load_services_config(config.services_path, &num_services);
        if (config.debug && services) {
            printf("[DEBUG] Loaded %d service configurations from %s\n", num_services, config.services_path);
        }
    }

    // Print configured options
    printf("WSSSH Client starting...\n");
    printf("Configuration:\n");
    printf("  WSSSHD Server: %s\n", config.wssshd_server ? config.wssshd_server : "(null)");
    printf("  WSSSHD Port: %d\n", config.wssshd_port);
    printf("  SSH Host: %s\n", config.ssh_host ? config.ssh_host : "(null)");
    printf("  SSH Port: %d\n", config.ssh_port);
    printf("  Client ID: %s\n", config.client_id ? config.client_id : "(null)");
    printf("  Password: %s\n", config.password ? "***" : "(null)");
    printf("  Reconnection Interval: %d seconds\n", config.interval);
    printf("  Debug Mode: %s\n", config.debug ? "enabled" : "disabled");
    printf("\n");

    while (1) {
        // Check for graceful shutdown before attempting to reconnect
        if (graceful_shutdown) {
            fprintf(stderr, "\nReceived SIGINT, attempting graceful shutdown...\n");
            fflush(stderr);
            printf("Graceful shutdown requested, exiting...\n");
            cleanup_tunnel(config.debug);
            break;
        }

        // Check for configuration reload
        if (reload_config) {
            reload_config = 0; // Reset flag
            fprintf(stderr, "\nReceived SIGHUP, reloading configuration...\n");
            fflush(stderr);
            reload_configuration(&config, &services, &num_services);
        }

        int result = connect_to_server(&config, &services, &num_services);
        if (result == 1) {
            // Error condition - use short retry interval for immediate reconnection
            if (config.debug) {
                printf("Connection lost, retrying in 1 seconds...\n");
            } else {
                log_message("EVENT", "Connection lost, retrying...\n");
            }
            sleep(1);
        } else if (result == 0) {
            // Close frame received - use short delay for immediate reconnection
            if (config.debug) {
                printf("[DEBUG - WebSockets] Server initiated disconnect, reconnecting in 1 seconds...\n");
                fflush(stdout);
            } else {
                log_message("EVENT", "Server initiated disconnect, reconnecting...\n");
            }
            sleep(1);
        }
        // result == 2 (registration failed) also uses shorter interval
    }

    // Cleanup
    if (services) {
        free_services_config(services, num_services);
    }
    free(config.wssshd_server);
    free(config.ssh_host);
    free(config.client_id);
    free(config.password);
    free(config.tunnel);
    free(config.tunnel_control);
    free(config.wssshd_private_ip);
    free(config.services_path);
    free(config.tunnel_host);
    free(config.default_service);
    pthread_mutex_destroy(&tunnel_mutex);

    return 0;
}