/**
 * WebSocket handling implementation for wssshd
 *
 * 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 <time.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <sys/stat.h>
#include <signal.h>
#include <setjmp.h>
#include <ctype.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <stdarg.h>
#include "websocket.h"
#include "websocket_protocol.h"
#include "plugin.h"
#include "ssl.h"

// Note: Removed crash recovery mechanism to prevent resource leaks

// Pre-computed JSON message templates
static const char *REGISTERED_MSG = "{\"type\":\"registered\",\"client_id\":\"%s\"}";
static const char *REGISTRATION_ERROR_MSG = "{\"type\":\"registration_error\",\"error\":\"%s\"}";
static const char *TUNNEL_REQUEST_MSG = "{\"type\":\"tunnel_request\",\"request_id\":\"%s\",\"enc\":\"%s\",\"service\":\"%s\",\"version\":\"%s\"}";
static const char *TUNNEL_ACK_MSG = "{\"type\":\"tunnel_ack\",\"request_id\":\"%s\"}";
static const char *TUNNEL_ERROR_MSG = "{\"type\":\"tunnel_error\",\"request_id\":\"%s\",\"error\":\"%s\"}";

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

wssshd_state_t *websocket_init_state(bool debug, const char *server_password) {
    wssshd_state_t *state = calloc(1, sizeof(wssshd_state_t));
    if (!state) return NULL;

    state->debug = debug;
    state->server_password = server_password ? strdup(server_password) : NULL;
    state->start_time = time(NULL);

    // Initialize dynamic arrays
    state->clients_capacity = 16;
    state->clients = calloc(state->clients_capacity, sizeof(client_t));

    state->tunnels_capacity = 16;
    state->tunnels = calloc(state->tunnels_capacity, sizeof(tunnel_t *));

    state->terminals_capacity = 16;
    state->terminals = calloc(state->terminals_capacity, sizeof(terminal_session_t));

    // Initialize mutexes
    pthread_mutex_init(&state->client_mutex, NULL);
    pthread_mutex_init(&state->tunnel_mutex, NULL);
    pthread_mutex_init(&state->terminal_mutex, NULL);

    return state;
}

void websocket_free_state(wssshd_state_t *state) {
    if (!state) return;

    // Free clients
    free(state->clients);

    // Free tunnels
    for (size_t i = 0; i < state->tunnels_count; i++) {
        tunnel_free(state->tunnels[i]);
    }
    free(state->tunnels);

    // Free terminals
    free(state->terminals);

    // Free password
    free((char *)state->server_password);

    // Destroy mutexes
    pthread_mutex_destroy(&state->client_mutex);
    pthread_mutex_destroy(&state->tunnel_mutex);
    pthread_mutex_destroy(&state->terminal_mutex);

    free(state);
}

// Client management functions
client_t *websocket_find_client(wssshd_state_t *state, const char *client_id) {
    pthread_mutex_lock(&state->client_mutex);
    for (size_t i = 0; i < state->clients_count; i++) {
        if (strcmp(state->clients[i].client_id, client_id) == 0) {
            pthread_mutex_unlock(&state->client_mutex);
            return &state->clients[i];
        }
    }
    pthread_mutex_unlock(&state->client_mutex);
    return NULL;
}

client_t *websocket_add_client(wssshd_state_t *state, const char *client_id, void *websocket, const char *services) {
    pthread_mutex_lock(&state->client_mutex);
    // Check if client already exists
    for (size_t i = 0; i < state->clients_count; i++) {
        if (strcmp(state->clients[i].client_id, client_id) == 0) {
            state->clients[i].active = true;
            state->clients[i].last_seen = time(NULL);
            state->clients[i].websocket = websocket;
            if (services) {
                strncpy(state->clients[i].services, services, sizeof(state->clients[i].services) - 1);
            } else if (strlen(state->clients[i].services) == 0) {
                strcpy(state->clients[i].services, "ssh");
            }
            pthread_mutex_unlock(&state->client_mutex);
            return &state->clients[i];
        }
    }

    // Expand array if needed
    if (state->clients_count >= state->clients_capacity) {
        state->clients_capacity *= 2;
        client_t *new_clients = realloc(state->clients, state->clients_capacity * sizeof(client_t));
        if (!new_clients) {
            pthread_mutex_unlock(&state->client_mutex);
            return NULL;
        }
        state->clients = new_clients;
    }

    // Add new client
    client_t *client = &state->clients[state->clients_count++];
    strncpy(client->client_id, client_id, sizeof(client->client_id) - 1);
    client->websocket = websocket;
    client->last_seen = time(NULL);
    client->active = true;
    strcpy(client->tunnel, "any");
    strcpy(client->tunnel_control, "any");
    strcpy(client->services, "ssh"); // Default service
    if (services) {
        strncpy(client->services, services, sizeof(client->services) - 1);
    } else {
        strcpy(client->services, "ssh"); // Default to ssh if no services specified
    }

    pthread_mutex_unlock(&state->client_mutex);
    return client;
}

void websocket_remove_client(wssshd_state_t *state, const char *client_id) {
    pthread_mutex_lock(&state->client_mutex);
    for (size_t i = 0; i < state->clients_count; i++) {
        if (strcmp(state->clients[i].client_id, client_id) == 0) {
            // Mark as inactive instead of removing
            state->clients[i].active = false;
            state->clients[i].last_seen = time(NULL);
            break;
        }
    }
    pthread_mutex_unlock(&state->client_mutex);
}

void websocket_update_client_activity(wssshd_state_t *state, const char *client_id) {
    pthread_mutex_lock(&state->client_mutex);
    for (size_t i = 0; i < state->clients_count; i++) {
        if (strcmp(state->clients[i].client_id, client_id) == 0) {
            state->clients[i].last_seen = time(NULL);
            state->clients[i].active = true;
            break;
        }
    }
    pthread_mutex_unlock(&state->client_mutex);
}

// Tunnel management functions
tunnel_t *websocket_find_tunnel(wssshd_state_t *state, const char *request_id) {
    pthread_mutex_lock(&state->tunnel_mutex);
    for (size_t i = 0; i < state->tunnels_count; i++) {
        if (strcmp(state->tunnels[i]->request_id, request_id) == 0) {
            pthread_mutex_unlock(&state->tunnel_mutex);
            return state->tunnels[i];
        }
    }
    pthread_mutex_unlock(&state->tunnel_mutex);
    return NULL;
}

tunnel_t *websocket_add_tunnel(wssshd_state_t *state, const char *request_id, const char *client_id) {
    pthread_mutex_lock(&state->tunnel_mutex);
    // Expand array if needed
    if (state->tunnels_count >= state->tunnels_capacity) {
        state->tunnels_capacity *= 2;
        tunnel_t **new_tunnels = realloc(state->tunnels, state->tunnels_capacity * sizeof(tunnel_t *));
        if (!new_tunnels) {
            pthread_mutex_unlock(&state->tunnel_mutex);
            return NULL;
        }
        state->tunnels = new_tunnels;
    }

    // Create new tunnel
    tunnel_t *tunnel = tunnel_create(request_id, client_id);
    if (!tunnel) {
        pthread_mutex_unlock(&state->tunnel_mutex);
        return NULL;
    }

    state->tunnels[state->tunnels_count++] = tunnel;
    pthread_mutex_unlock(&state->tunnel_mutex);
    return tunnel;
}

void websocket_remove_tunnel(wssshd_state_t *state, const char *request_id) {
    pthread_mutex_lock(&state->tunnel_mutex);
    for (size_t i = 0; i < state->tunnels_count; i++) {
        if (strcmp(state->tunnels[i]->request_id, request_id) == 0) {
            tunnel_free(state->tunnels[i]);
            // Shift remaining elements
            memmove(&state->tunnels[i], &state->tunnels[i + 1],
                    (state->tunnels_count - i - 1) * sizeof(tunnel_t *));
            state->tunnels_count--;
            break;
        }
    }
    pthread_mutex_unlock(&state->tunnel_mutex);
}

// Terminal management functions
terminal_session_t *websocket_find_terminal(wssshd_state_t *state, const char *request_id) {
    pthread_mutex_lock(&state->terminal_mutex);
    for (size_t i = 0; i < state->terminals_count; i++) {
        if (strcmp(state->terminals[i].request_id, request_id) == 0) {
            pthread_mutex_unlock(&state->terminal_mutex);
            return &state->terminals[i];
        }
    }
    pthread_mutex_unlock(&state->terminal_mutex);
    return NULL;
}

terminal_session_t *websocket_add_terminal(wssshd_state_t *state, const char *request_id, const char *client_id, const char *username, pid_t proc_pid, int master_fd) {
    pthread_mutex_lock(&state->terminal_mutex);
    // Expand array if needed
    if (state->terminals_count >= state->terminals_capacity) {
        state->terminals_capacity *= 2;
        terminal_session_t *new_terminals = realloc(state->terminals, state->terminals_capacity * sizeof(terminal_session_t));
        if (!new_terminals) {
            pthread_mutex_unlock(&state->terminal_mutex);
            return NULL;
        }
        state->terminals = new_terminals;
    }

    // Add new terminal
    terminal_session_t *terminal = &state->terminals[state->terminals_count++];
    strncpy(terminal->request_id, request_id, sizeof(terminal->request_id) - 1);
    if (client_id) strncpy(terminal->client_id, client_id, sizeof(terminal->client_id) - 1);
    if (username) strncpy(terminal->username, username, sizeof(terminal->username) - 1);
    terminal->proc_pid = proc_pid;
    terminal->master_fd = master_fd;

    pthread_mutex_unlock(&state->terminal_mutex);
    return terminal;
}

void websocket_remove_terminal(wssshd_state_t *state, const char *request_id) {
    pthread_mutex_lock(&state->terminal_mutex);
    for (size_t i = 0; i < state->terminals_count; i++) {
        if (strcmp(state->terminals[i].request_id, request_id) == 0) {
            // Shift remaining elements
            memmove(&state->terminals[i], &state->terminals[i + 1],
                    (state->terminals_count - i - 1) * sizeof(terminal_session_t));
            state->terminals_count--;
            break;
        }
    }
    pthread_mutex_unlock(&state->terminal_mutex);
}

// Message handling with crash protection
int websocket_handle_message(wssshd_state_t *state, ws_connection_t *conn __attribute__((unused)), const char *message, size_t message_len) {
    // Validate input parameters
    if (!state || !message || message_len == 0 || message_len > 50 * 1024 * 1024) { // 50MB safety limit
        if (state && state->debug) {
            printf("[DEBUG - unknown -> wssshd] Invalid message parameters (len=%zu)\n", message_len);
        }
        return -1;
    }

    // Simple string-based JSON parsing for basic functionality
    // This is a simplified implementation - a full JSON parser would be better

    #define MSG_BUFFER_SIZE 2097152  // 2MB buffer for large messages
    char *msg_copy = malloc(MSG_BUFFER_SIZE);
    if (!msg_copy) {
        if (state->debug) {
            printf("[DEBUG - unknown -> wssshd] Failed to allocate message buffer\n");
        }
        return -1;
    }

    // Safe copy with bounds checking
    if (message_len >= MSG_BUFFER_SIZE) {
        if (state->debug) {
            printf("[DEBUG - unknown -> wssshd] Message too long: %zu bytes\n", message_len);
        }
        free(msg_copy);
        return -1;
    }

    // Use memcpy with explicit bounds checking
    memcpy(msg_copy, message, message_len);
    msg_copy[message_len] = '\0';

    // Validate basic JSON structure to prevent crashes from malformed messages
    if (message_len < 2 || message[0] != '{' || message[message_len-1] != '}') {
        if (state->debug) {
            printf("[DEBUG - unknown -> wssshd] Malformed JSON message (invalid braces)\n");
        }
        free(msg_copy);
        return -1;
    }

    // Additional safety checks for message content
    if (strchr(msg_copy, '\0') != msg_copy + message_len) {
        if (state->debug) {
            printf("[DEBUG - unknown -> wssshd] Message contains null bytes\n");
        }
        free(msg_copy);
        return -1;
    }

    // Determine connection direction for debug messages with safety
    const char *direction = "unknown";

    // Safe string operations with null checks
    if (msg_copy && strlen(msg_copy) > 10) { // Minimum length for a valid message
        if ((strstr(msg_copy, "\"type\":\"register\"") != NULL) ||
            (strstr(msg_copy, "\"type\": \"register\"") != NULL)) {
            direction = "wssshc";
        } else if ((strstr(msg_copy, "\"type\":\"tunnel_request\"") != NULL) ||
                   (strstr(msg_copy, "\"type\": \"tunnel_request\"") != NULL)) {
            direction = "wsssht";
        } else {
            // Safe array iteration with bounds checking
            for (size_t i = 0; i < state->clients_count && i < 1000; i++) { // Safety limit
                if (state->clients[i].websocket == conn) {
                    direction = "wssshc";
                    break;
                }
            }
            if (strcmp(direction, "unknown") == 0) {
                for (size_t i = 0; i < state->tunnels_count && i < 1000; i++) { // Safety limit
                    if (state->tunnels[i] && state->tunnels[i]->wsssh_ws == conn) {
                        direction = "wsssht";
                        break;
                    }
                }
            }
        }
    }

    if (state->debug) {
        // Check if this is a tunnel_data or tunnel_response message to avoid printing large data
        if (strstr(msg_copy, "\"type\":\"tunnel_data\"") || strstr(msg_copy, "\"type\": \"tunnel_data\"") ||
            strstr(msg_copy, "\"type\":\"tunnel_response\"") || strstr(msg_copy, "\"type\": \"tunnel_response\"")) {
            // Extract request_id and data size for tunnel messages
            char *request_id_start = strstr(msg_copy, "\"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(msg_copy, "\"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;
                        }
                    }
                }
            }
            const char *msg_type = strstr(msg_copy, "tunnel_data") ? "tunnel_data" : "tunnel_response";
            printf("[DEBUG - %s -> wssshd] Handling message: {\"type\":\"%s\", \"request_id\":\"%s\", \"data\":\"<size: %zu bytes>\"}\n",
                   direction, msg_type, request_id, data_size);
        } else {
            printf("[DEBUG - %s -> wssshd] Handling message: %.*s\n", direction, (int)message_len, message);
        }
    }

    int result = 0;

    // Note: Removed crash recovery mechanism as it can cause resource leaks
    // and interfere with proper error handling. Instead, rely on proper
    // bounds checking and error handling throughout the code.

    // Check for registration message with safe string operations
    if (strstr(msg_copy, "\"type\":\"register\"") || strstr(msg_copy, "\"type\": \"register\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing registration message\n", direction);
            printf("[DEBUG - %s -> wssshd] Full message: %s\n", direction, msg_copy);
        }
        // Extract client_id and password (simplified parsing)
        char *client_id = NULL;
        char *password = NULL;

        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Parsing client_id and password from: %s\n", direction, msg_copy);
        }

        // Extract client_id (make a copy to avoid modifying the original string)
        char *client_id_start = strstr(msg_copy, "\"client_id\":\"");
        const char *key_pattern = "\"client_id\":\"";
        size_t key_len = strlen(key_pattern);

        if (!client_id_start) {
            client_id_start = strstr(msg_copy, "\"id\":\"");
            key_pattern = "\"id\":\"";
            key_len = strlen(key_pattern);
        }

        if (client_id_start && (client_id_start + key_len) < (msg_copy + MSG_BUFFER_SIZE)) {
            client_id_start += key_len;
            char *client_id_end = strchr(client_id_start, '"');
            if (client_id_end && client_id_end > client_id_start &&
                client_id_end < msg_copy + MSG_BUFFER_SIZE &&
                (size_t)(client_id_end - msg_copy) < MSG_BUFFER_SIZE) {
                size_t client_id_len = client_id_end - client_id_start;
                if (client_id_len > 0 && client_id_len < 64) {  // Reasonable limit for client ID
                    char *client_id_copy = malloc(client_id_len + 1);
                    if (client_id_copy) {
                        memcpy(client_id_copy, client_id_start, client_id_len);
                        client_id_copy[client_id_len] = '\0';
                        client_id = client_id_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted client_id: '%s'\n", direction, client_id);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID length invalid: %zu\n", direction, client_id_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID end quote not found or invalid\n", direction);
            }
        } else {
            if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID start not found\n", direction);
        }

        // Extract password (search in original unmodified string)
        char *password_start = strstr(msg_copy, "\"password\":\"");
        if (password_start) {
            password_start += strlen("\"password\":\"");
            char *password_end = strchr(password_start, '"');
            if (password_end && password_end > password_start && password_end < msg_copy + MSG_BUFFER_SIZE) {
                size_t password_len = password_end - password_start;
                if (password_len > 0 && password_len < 256) {  // Reasonable limit for password
                    char *password_copy = malloc(password_len + 1);
                    if (password_copy) {
                        memcpy(password_copy, password_start, password_len);
                        password_copy[password_len] = '\0';
                        password = password_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted password: '%s'\n", direction, password);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Password length invalid: %zu\n", direction, password_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Password end quote not found or invalid\n", direction);
            }
        } else {
            if (state->debug) printf("[DEBUG - %s -> wssshd] Password start not found\n", direction);
        }

        // Extract services (optional, defaults to "ssh")
        char *services = NULL;
        char *services_start = strstr(msg_copy, "\"services\":");
        if (services_start) {
            services_start += strlen("\"services\":");
            // Skip whitespace
            while (*services_start == ' ' || *services_start == '\t' || *services_start == '\n') services_start++;

            if (*services_start == '[') {
                // New format: array of service objects
                char *services_end = strchr(services_start, ']');
                if (services_end && services_end < msg_copy + MSG_BUFFER_SIZE) {
                    services_end++; // Include the closing bracket
                    size_t services_len = services_end - services_start;
                    if (services_len > 0 && services_len < 2048) {  // Reasonable limit for services array
                        char *services_copy = malloc(services_len + 1);
                        if (services_copy) {
                            memcpy(services_copy, services_start, services_len);
                            services_copy[services_len] = '\0';

                            // Parse service names from the array
                            char *service_names = malloc(256); // Buffer for concatenated service names
                            if (service_names) {
                                service_names[0] = '\0';
                                char *service_start = services_copy + 1; // Skip opening bracket

                                while (*service_start && *service_start != ']') {
                                    // Find next service object
                                    if (*service_start == '{') {
                                        char *service_end = strchr(service_start, '}');
                                        if (service_end) {
                                            // Extract service name
                                            char *name_start = strstr(service_start, "\"name\":\"");
                                            if (name_start) {
                                                name_start += strlen("\"name\":\"");
                                                char *name_end = strchr(name_start, '"');
                                                if (name_end) {
                                                    size_t name_len = name_end - name_start;
                                                    if (name_len > 0 && name_len < 64) {
                                                        if (strlen(service_names) > 0) {
                                                            strncat(service_names, ",", 255 - strlen(service_names));
                                                        }
                                                        strncat(service_names, name_start, name_len);
                                                    }
                                                }
                                            }
                                            service_start = service_end + 1;
                                        } else {
                                            break;
                                        }
                                    } else {
                                        service_start++;
                                    }

                                    // Skip commas and whitespace
                                    while (*service_start == ',' || *service_start == ' ' || *service_start == '\t' || *service_start == '\n') {
                                        service_start++;
                                    }
                                }

                                if (strlen(service_names) > 0) {
                                    services = service_names;
                                    // Convert services to lowercase for consistency
                                    for (char *p = services; *p; p++) {
                                        *p = tolower(*p);
                                    }
                                    if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted services from array: '%s'\n", direction, services);
                                } else {
                                    free(service_names);
                                }
                            }
                            free(services_copy);
                        }
                    } else {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Services array length invalid: %zu\n", direction, services_len);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Services array end bracket not found\n", direction);
                }
            } else if (*services_start == '"') {
                // Legacy format: simple string
                services_start++; // Skip opening quote
                char *services_end = strchr(services_start, '"');
                if (services_end && services_end > services_start && services_end < msg_copy + MSG_BUFFER_SIZE) {
                    size_t services_len = services_end - services_start;
                    if (services_len > 0 && services_len < 256) {  // Reasonable limit for services
                        char *services_copy = malloc(services_len + 1);
                        if (services_copy) {
                            memcpy(services_copy, services_start, services_len);
                            services_copy[services_len] = '\0';
                            services = services_copy;
                            // Convert services to lowercase for consistency
                            for (char *p = services; *p; p++) {
                                *p = tolower(*p);
                            }
                            if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted services (legacy format): '%s'\n", direction, services);
                        }
                    } else {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Services length invalid: %zu\n", direction, services_len);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Services end quote not found or invalid\n", direction);
                }
            }
        } else {
            if (state->debug) printf("[DEBUG - %s -> wssshd] Services not specified, using default\n", direction);
        }

        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Password check: received='%s', expected='%s'\n",
                   direction,
                   password ? password : "(null)",
                   state->server_password ? state->server_password : "(null)");
        }

        if (client_id && password && strcmp(password, state->server_password) == 0) {
            client_t *client = websocket_add_client(state, client_id, (void *)conn, services);
            if (client) {
                // Set client IP
                strcpy(client->ip, inet_ntoa(conn->addr.sin_addr));
                // Send registration success
                char response[512];
                snprintf(response, sizeof(response), REGISTERED_MSG, client_id);
                if (state->debug) printf("[DEBUG - wssshd -> wssshc] Sending registration success: %s\n", response);
                ws_send_frame(conn, WS_OPCODE_TEXT, response, strlen(response));
                if (!state->debug) log_message("EVENT", "Client %s registered\n", client_id);
            }
        } else {
            // Send registration error
            char response[512];
            snprintf(response, sizeof(response), REGISTRATION_ERROR_MSG, "Invalid password");
            if (state->debug) printf("[DEBUG - wssshd -> wssshc] Sending registration error: %s\n", response);
            ws_send_frame(conn, WS_OPCODE_TEXT, response, strlen(response));
            if (state->debug) {
                printf("[DEBUG - %s -> wssshd] Client %s registration failed: client_id=%s, password=%s, server_password=%s\n",
                        direction,
                        client_id ? client_id : "unknown",
                        client_id ? client_id : "(null)",
                        password ? password : "(null)",
                        state->server_password ? state->server_password : "(null)");
            }
        }

        // Free allocated strings
        if (client_id) free(client_id);
        if (password) free(password);
        if (services) free(services);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_request\"") || strstr(msg_copy, "\"type\": \"tunnel_request\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel request\n", direction);
            printf("[DEBUG - %s -> wssshd] Full tunnel request: %s\n", direction, msg_copy);
        }
        // Handle tunnel request (simplified)
        char *client_id = NULL;
        char *request_id = NULL;
        char *enc = NULL;
        char *service = NULL;
        char *version = NULL;

        // Extract client_id (make a copy to avoid modifying the original string)
        char *tunnel_client_id_start = strstr(msg_copy, "\"client_id\":\"");
        if (tunnel_client_id_start && (tunnel_client_id_start + strlen("\"client_id\":\"")) < (msg_copy + MSG_BUFFER_SIZE)) {
            tunnel_client_id_start += strlen("\"client_id\":\"");
            char *tunnel_client_id_end = strchr(tunnel_client_id_start, '"');
            if (tunnel_client_id_end && tunnel_client_id_end > tunnel_client_id_start &&
                tunnel_client_id_end < msg_copy + MSG_BUFFER_SIZE &&
                (size_t)(tunnel_client_id_end - msg_copy) < MSG_BUFFER_SIZE) {
                size_t tunnel_client_id_len = tunnel_client_id_end - tunnel_client_id_start;
                if (tunnel_client_id_len > 0 && tunnel_client_id_len < 64) {  // Reasonable limit for client ID
                    char *tunnel_client_id_copy = malloc(tunnel_client_id_len + 1);
                    if (tunnel_client_id_copy) {
                        memcpy(tunnel_client_id_copy, tunnel_client_id_start, tunnel_client_id_len);
                        tunnel_client_id_copy[tunnel_client_id_len] = '\0';
                        client_id = tunnel_client_id_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted tunnel client_id: '%s' (ptr=%p)\n", direction, client_id, client_id);
                    } else {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to allocate memory for client_id\n", direction);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID length invalid: %zu\n", direction, tunnel_client_id_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID end quote not found or invalid\n", direction);
            }
        } else {
            if (state->debug) printf("[DEBUG - %s -> wssshd] Client ID start not found\n", direction);
        }

        // Extract request_id (search in original unmodified string)
        char *tunnel_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (tunnel_request_id_start && (tunnel_request_id_start + strlen("\"request_id\":\"")) < (msg_copy + MSG_BUFFER_SIZE)) {
            tunnel_request_id_start += strlen("\"request_id\":\"");
            char *tunnel_request_id_end = strchr(tunnel_request_id_start, '"');
            if (tunnel_request_id_end && tunnel_request_id_end > tunnel_request_id_start &&
                tunnel_request_id_end < msg_copy + MSG_BUFFER_SIZE &&
                (size_t)(tunnel_request_id_end - msg_copy) < MSG_BUFFER_SIZE) {
                size_t tunnel_request_id_len = tunnel_request_id_end - tunnel_request_id_start;
                if (tunnel_request_id_len > 0 && tunnel_request_id_len < 64) {  // Reasonable limit for request ID
                    char *tunnel_request_id_copy = malloc(tunnel_request_id_len + 1);
                    if (tunnel_request_id_copy) {
                        memcpy(tunnel_request_id_copy, tunnel_request_id_start, tunnel_request_id_len);
                        tunnel_request_id_copy[tunnel_request_id_len] = '\0';
                        request_id = tunnel_request_id_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted request_id: '%s'\n", direction, request_id);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Request ID length invalid: %zu\n", direction, tunnel_request_id_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Request ID end quote not found or invalid\n", direction);
            }
        }

        // Extract enc
        char *enc_start = strstr(msg_copy, "\"enc\":\"");
        if (enc_start) {
            enc_start += strlen("\"enc\":\"");
            char *enc_end = strchr(enc_start, '"');
            if (enc_end && enc_end > enc_start && enc_end < msg_copy + MSG_BUFFER_SIZE) {
                size_t enc_len = enc_end - enc_start;
                if (enc_len > 0 && enc_len < 32) {  // Reasonable limit for encoding type
                    char *enc_copy = malloc(enc_len + 1);
                    if (enc_copy) {
                        memcpy(enc_copy, enc_start, enc_len);
                        enc_copy[enc_len] = '\0';
                        enc = enc_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted enc: '%s'\n", direction, enc);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Enc length invalid: %zu\n", direction, enc_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Enc end quote not found or invalid\n", direction);
            }
        }

        // Extract service
        char *service_start = strstr(msg_copy, "\"service\":\"");
        if (service_start) {
            service_start += strlen("\"service\":\"");
            char *service_end = strchr(service_start, '"');
            if (service_end && service_end > service_start && service_end < msg_copy + MSG_BUFFER_SIZE) {
                size_t service_len = service_end - service_start;
                if (service_len > 0 && service_len < 32) {  // Reasonable limit for service type
                    char *service_copy = malloc(service_len + 1);
                    if (service_copy) {
                        memcpy(service_copy, service_start, service_len);
                        service_copy[service_len] = '\0';
                        service = service_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted service: '%s'\n", direction, service);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Service length invalid: %zu\n", direction, service_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Service end quote not found or invalid\n", direction);
            }
        }

        // Extract version
        char *version_start = strstr(msg_copy, "\"version\":\"");
        if (version_start) {
            version_start += strlen("\"version\":\"");
            char *version_end = strchr(version_start, '"');
            if (version_end && version_end > version_start && version_end < msg_copy + MSG_BUFFER_SIZE) {
                size_t version_len = version_end - version_start;
                if (version_len > 0 && version_len < 32) {  // Reasonable limit for version string
                    char *version_copy = malloc(version_len + 1);
                    if (version_copy) {
                        memcpy(version_copy, version_start, version_len);
                        version_copy[version_len] = '\0';
                        version = version_copy;
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted version: '%s'\n", direction, version);
                    }
                } else {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Version length invalid: %zu\n", direction, version_len);
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Version end quote not found or invalid\n", direction);
            }
        }


        if (client_id && request_id) {
            client_t *client = websocket_find_client(state, client_id);
            if (client && client->active) {
                tunnel_t *tunnel = websocket_add_tunnel(state, request_id, client_id);
                if (tunnel) {
                    // Store the wsssht connection in the tunnel for data forwarding
                    tunnel->wsssh_ws = conn;
                    // Store the wssshc connection in the tunnel
                    tunnel->client_ws = client->websocket;

                    tunnel_update_status(tunnel, TUNNEL_STATUS_ACTIVE, NULL);

                    // Send tunnel request to client (wssshc)
                    char request_msg[1024];
                    // Use defaults if parameters not provided
                    const char *enc_val = enc ? enc : "hex";
                    const char *service_val = service ? service : "ssh";
                    const char *version_val = version ? version : "1.6.5";
                    snprintf(request_msg, sizeof(request_msg), TUNNEL_REQUEST_MSG, request_id, enc_val, service_val, version_val);
                    if (state->debug) printf("[DEBUG - wssshd -> wssshc] Sending tunnel request: %s\n", request_msg);
                    ws_send_frame(client->websocket, WS_OPCODE_TEXT, request_msg, strlen(request_msg));

                    // Send tunnel ack to tool (wsssht/wsscp)
                    char ack_msg[256];
                    snprintf(ack_msg, sizeof(ack_msg), TUNNEL_ACK_MSG, request_id);
                    if (state->debug) printf("[DEBUG - wssshd -> wsssht] Sending tunnel ack: %s\n", ack_msg);
                    ws_send_frame(conn, WS_OPCODE_TEXT, ack_msg, strlen(ack_msg));

                    if (state->debug) printf("[DEBUG - %s -> wssshd] Created tunnel %s for client %s\n", direction, request_id, client_id);
                    if (!state->debug) log_message("EVENT", "New tunnel %s for client %s\n", request_id, client_id);
                }
            } else {
                // Send error to tool
                char error_msg[512];
                snprintf(error_msg, sizeof(error_msg), TUNNEL_ERROR_MSG, request_id, "Client not registered or disconnected");
                if (state->debug) printf("[DEBUG - wssshd -> wsssht] Sending tunnel error: %s\n", error_msg);
                ws_send_frame(conn, WS_OPCODE_TEXT, error_msg, strlen(error_msg));
                if (state->debug) printf("[DEBUG - %s -> wssshd] Tunnel request failed: client %s not found or inactive\n", direction, client_id);

                // Close the connection that sent the invalid tunnel request
                if (state->debug) printf("[DEBUG - %s -> wssshd] Closing connection due to invalid tunnel request from unregistered client %s\n", direction, client_id);
                conn->state = WS_STATE_CLOSED;
            }
        }

        // Free allocated strings
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Freeing tunnel request strings: client_id=%p, request_id=%p, enc=%p, service=%p, version=%p\n", direction, client_id, request_id, enc, service, version);
        }
        if (client_id) free(client_id);
        if (request_id) free(request_id);
        if (enc) free(enc);
        if (service) free(service);
        if (version) free(version);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_ack\"") || strstr(msg_copy, "\"type\": \"tunnel_ack\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel acknowledgment\n", direction);
        }
        // Handle tunnel acknowledgment from wssshc
        char *request_id = NULL;

        char *ack_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (ack_request_id_start) {
            ack_request_id_start += strlen("\"request_id\":\"");
            char *ack_request_id_end = strchr(ack_request_id_start, '"');
            if (ack_request_id_end) {
                size_t ack_request_id_len = ack_request_id_end - ack_request_id_start;
                char *ack_request_id_copy = malloc(ack_request_id_len + 1);
                if (ack_request_id_copy) {
                    memcpy(ack_request_id_copy, ack_request_id_start, ack_request_id_len);
                    ack_request_id_copy[ack_request_id_len] = '\0';
                    request_id = ack_request_id_copy;
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted tunnel_ack request_id: '%s'\n", direction, request_id);
                }
            }
        }

        if (request_id) {
            // Find the tunnel and mark it as acknowledged
            // This is a simple acknowledgment - in a full implementation,
            // you might want to track tunnel state more thoroughly
            if (state->debug) printf("[DEBUG - %s -> wssshd] Tunnel %s acknowledged by client\n", direction, request_id);
        }

        if (request_id) free(request_id);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_data\"") || strstr(msg_copy, "\"type\": \"tunnel_data\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel data\n", direction);
        }
        // Handle tunnel data from wssshc to wsssht
        char *request_id = NULL;
        char *data = NULL;

        // Extract request_id
        char *data_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (data_request_id_start) {
            data_request_id_start += strlen("\"request_id\":\"");
            char *data_request_id_end = strchr(data_request_id_start, '"');
            if (data_request_id_end) {
                size_t data_request_id_len = data_request_id_end - data_request_id_start;
                char *data_request_id_copy = malloc(data_request_id_len + 1);
                if (data_request_id_copy) {
                    memcpy(data_request_id_copy, data_request_id_start, data_request_id_len);
                    data_request_id_copy[data_request_id_len] = '\0';
                    request_id = data_request_id_copy;
                }
            }
        }

        // Extract data
        char *data_start = strstr(msg_copy, "\"data\":\"");
        if (data_start) {
            data_start += strlen("\"data\":\"");
            char *data_end = strchr(data_start, '"');
            if (data_end) {
                size_t data_len = data_end - data_start;
                char *data_copy = malloc(data_len + 1);
                if (data_copy) {
                    memcpy(data_copy, data_start, data_len);
                    data_copy[data_len] = '\0';
                    data = data_copy;
                }
            }
        }

        if (request_id && data) {
            // Find the tunnel
            tunnel_t *tunnel = websocket_find_tunnel(state, request_id);
            if (tunnel) {
                // Determine which side to forward to based on sender
                ws_connection_t *target_conn = NULL;
                const char *target_side = NULL;

                if (strcmp(direction, "wssshc") == 0) {
                    // Message from wssshc, forward to wsssht
                    target_conn = tunnel->wsssh_ws;
                    target_side = "wsssht";
                } else if (strcmp(direction, "wsssht") == 0) {
                    // Message from wsssht, forward to wssshc
                    target_conn = tunnel->client_ws;
                    target_side = "wssshc";
                }

                if (target_conn) {
                    // Check connection health before attempting to send
                    if (!ws_connection_is_healthy(target_conn)) {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Target connection for request %s is not healthy, closing tunnel\n", direction, request_id);
                        // Send tunnel_close to the still-connected end
                        ws_connection_t *other_conn = (strcmp(direction, "wssshc") == 0) ? tunnel->wsssh_ws : tunnel->client_ws;
                        if (other_conn && ((ws_connection_t *)other_conn)->state == WS_STATE_OPEN) {
                            char close_msg[256];
                            snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);
                            ws_send_frame(other_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                            if (state->debug) printf("[DEBUG - wssshd -> %s] Sent tunnel_close for request %s due to unhealthy connection\n", (strcmp(direction, "wssshc") == 0) ? "wsssht" : "wssshc", request_id);
                        }
                        // Mark tunnel as error and remove it
                        tunnel_update_status(tunnel, TUNNEL_STATUS_ERROR, "Target connection unhealthy");
                        websocket_remove_tunnel(state, request_id);
                    } else {
                        // Calculate binary size from hex data length
                        size_t hex_len = strlen(data);
                        size_t binary_size = hex_len / 2;

                        // Update tunnel statistics
                        tunnel->total_bytes_sent += binary_size;
                        tunnel->bytes_last_period += binary_size;

                        // Forward hex data directly without decode/re-encode cycle
                        size_t request_id_len = strlen(request_id);
                        size_t data_len = strlen(data);
                        size_t json_overhead = strlen("{\"type\":\"tunnel_data\",\"request_id\":\"\",\"size\":,\"data\":\"\"}");
                        size_t total_size = json_overhead + request_id_len + 32 + data_len + 1; // Extra 32 for safety margin

                        // Allocate buffer dynamically to handle large messages
                        char *forward_msg = malloc(total_size);
                        if (forward_msg) {
                            int msg_len = snprintf(forward_msg, total_size, "{\"type\":\"tunnel_data\",\"request_id\":\"%s\",\"size\":%zu,\"data\":\"%s\"}", request_id, binary_size, data);
                            if (msg_len > 0 && (size_t)msg_len < total_size) {
                                if (state->debug) printf("[DEBUG - wssshd -> %s] Forwarding tunnel data: request_id=%s, size=%zu bytes\n", target_side, request_id, binary_size);
                                bool send_result = ws_send_frame(target_conn, WS_OPCODE_TEXT, forward_msg, msg_len);
                                if (send_result) {
                                    if (state->debug) printf("[DEBUG - %s -> wssshd] Forwarded tunnel data for request %s to %s, hex length: %zu bytes (binary size: %zu bytes)\n", direction, request_id, target_side, data_len, binary_size);
                                } else {
                                    if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to send tunnel data for request %s to %s (connection may be broken)\n", direction, request_id, target_side);
                                    // Send tunnel_close to the still-connected end when forwarding fails
                                    ws_connection_t *other_conn = (strcmp(direction, "wssshc") == 0) ? tunnel->wsssh_ws : tunnel->client_ws;
                                    if (other_conn && ((ws_connection_t *)other_conn)->state == WS_STATE_OPEN) {
                                        char close_msg[256];
                                        snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);
                                        ws_send_frame(other_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                                        if (state->debug) printf("[DEBUG - wssshd -> %s] Sent tunnel_close for request %s due to forwarding failure\n", (strcmp(direction, "wssshc") == 0) ? "wsssht" : "wssshc", request_id);
                                    }
                                    // Mark tunnel as error and remove it
                                    tunnel_update_status(tunnel, TUNNEL_STATUS_ERROR, "Connection broken during data forwarding");
                                    websocket_remove_tunnel(state, request_id);
                                }
                            } else {
                                if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to format forward message (msg_len=%d, total_size=%zu)\n", direction, msg_len, total_size);
                            }
                            free(forward_msg);
                        } else {
                            if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to allocate buffer for forward message\n", direction);
                        }
                    }
                } else {
                    if (state->debug) {
                        if (!target_conn) {
                            printf("[DEBUG - %s -> wssshd] Could not find target connection for request %s\n", direction, request_id);
                        } else {
                            printf("[DEBUG - %s -> wssshd] Target connection for request %s is not open (state=%d)\n", direction, request_id, target_conn->state);
                        }
                    }
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Could not find tunnel for request %s\n", direction, request_id);
            }
        }

        if (request_id) free(request_id);
        if (data) free(data);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_response\"") || strstr(msg_copy, "\"type\": \"tunnel_response\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel response\n", direction);
        }
        // Handle tunnel response from wssshc to wsssht
        char *request_id = NULL;
        char *data = NULL;

        // Extract request_id
        char *resp_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (resp_request_id_start) {
            resp_request_id_start += strlen("\"request_id\":\"");
            char *resp_request_id_end = strchr(resp_request_id_start, '"');
            if (resp_request_id_end) {
                size_t resp_request_id_len = resp_request_id_end - resp_request_id_start;
                char *resp_request_id_copy = malloc(resp_request_id_len + 1);
                if (resp_request_id_copy) {
                    memcpy(resp_request_id_copy, resp_request_id_start, resp_request_id_len);
                    resp_request_id_copy[resp_request_id_len] = '\0';
                    request_id = resp_request_id_copy;
                }
            }
        }

        // Extract data
        char *resp_data_start = strstr(msg_copy, "\"data\":\"");
        if (resp_data_start) {
            resp_data_start += strlen("\"data\":\"");
            char *resp_data_end = strchr(resp_data_start, '"');
            if (resp_data_end) {
                size_t resp_data_len = resp_data_end - resp_data_start;
                char *resp_data_copy = malloc(resp_data_len + 1);
                if (resp_data_copy) {
                    memcpy(resp_data_copy, resp_data_start, resp_data_len);
                    resp_data_copy[resp_data_len] = '\0';
                    data = resp_data_copy;
                }
            }
        }

        if (request_id && data) {
            // Find the tunnel
            tunnel_t *tunnel = websocket_find_tunnel(state, request_id);
            if (tunnel) {
                // Determine which side to forward to based on sender
                ws_connection_t *target_conn = NULL;
                const char *target_side = NULL;

                if (strcmp(direction, "wssshc") == 0) {
                    // Message from wssshc, forward to wsssht
                    target_conn = tunnel->wsssh_ws;
                    target_side = "wsssht";
                } else if (strcmp(direction, "wsssht") == 0) {
                    // Message from wsssht, forward to wssshc
                    target_conn = tunnel->client_ws;
                    target_side = "wssshc";
                }

                if (target_conn) {
                    // Check connection health before attempting to send
                    if (!ws_connection_is_healthy(target_conn)) {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Target connection for request %s is not healthy, closing tunnel\n", direction, request_id);
                        // Send tunnel_close to the still-connected end
                        ws_connection_t *other_conn = (strcmp(direction, "wssshc") == 0) ? tunnel->wsssh_ws : tunnel->client_ws;
                        if (other_conn && ((ws_connection_t *)other_conn)->state == WS_STATE_OPEN) {
                            char close_msg[256];
                            snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);
                            ws_send_frame(other_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                            if (state->debug) printf("[DEBUG - wssshd -> %s] Sent tunnel_close for request %s due to unhealthy connection\n", (strcmp(direction, "wssshc") == 0) ? "wsssht" : "wssshc", request_id);
                        }
                        // Mark tunnel as error and remove it
                        tunnel_update_status(tunnel, TUNNEL_STATUS_ERROR, "Target connection unhealthy");
                        websocket_remove_tunnel(state, request_id);
                    } else {
                        // Calculate binary size from hex data length
                        size_t hex_len = strlen(data);
                        size_t binary_size = hex_len / 2;

                        // Update tunnel statistics
                        tunnel->total_bytes_sent += binary_size;
                        tunnel->bytes_last_period += binary_size;

                        // Forward hex data directly without decode/re-encode cycle
                        size_t request_id_len = strlen(request_id);
                        size_t data_len = strlen(data);
                        size_t json_overhead = strlen("{\"type\":\"tunnel_response\",\"request_id\":\"\",\"data\":\"\"}");
                        size_t total_size = json_overhead + request_id_len + 32 + data_len + 1; // Extra 32 for safety margin

                        // Allocate buffer dynamically to handle large messages
                        char *forward_msg = malloc(total_size);
                        if (forward_msg) {
                            int msg_len = snprintf(forward_msg, total_size, "{\"type\":\"tunnel_response\",\"request_id\":\"%s\",\"data\":\"%s\"}", request_id, data);
                            if (msg_len > 0 && (size_t)msg_len < total_size) {
                                if (state->debug) printf("[DEBUG - wssshd -> %s] Forwarding tunnel response: %.*s\n", target_side, msg_len, forward_msg);
                                bool send_result = ws_send_frame(target_conn, WS_OPCODE_TEXT, forward_msg, msg_len);
                                if (send_result) {
                                    if (state->debug) printf("[DEBUG - %s -> wssshd] Forwarded tunnel response for request %s to %s, hex length: %zu bytes\n", direction, request_id, target_side, data_len);
                                } else {
                                    if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to send tunnel response for request %s to %s (connection may be broken)\n", direction, request_id, target_side);
                                    // Send tunnel_close to the still-connected end when forwarding fails
                                    ws_connection_t *other_conn = (strcmp(direction, "wssshc") == 0) ? tunnel->wsssh_ws : tunnel->client_ws;
                                    if (other_conn && ((ws_connection_t *)other_conn)->state == WS_STATE_OPEN) {
                                        char close_msg[256];
                                        snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);
                                        ws_send_frame(other_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                                        if (state->debug) printf("[DEBUG - wssshd -> %s] Sent tunnel_close for request %s due to forwarding failure\n", (strcmp(direction, "wssshc") == 0) ? "wsssht" : "wssshc", request_id);
                                    }
                                    // Mark tunnel as error and remove it
                                    tunnel_update_status(tunnel, TUNNEL_STATUS_ERROR, "Connection broken during response forwarding");
                                    websocket_remove_tunnel(state, request_id);
                                }
                            } else {
                                if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to format forward message (msg_len=%d, total_size=%zu)\n", direction, msg_len, total_size);
                            }
                            free(forward_msg);
                        } else {
                            if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to allocate buffer for forward message\n", direction);
                        }
                    }
                } else {
                    if (state->debug) {
                        if (!target_conn) {
                            printf("[DEBUG - %s -> wssshd] Could not find target connection for request %s\n", direction, request_id);
                        } else {
                            printf("[DEBUG - %s -> wssshd] Target connection for request %s is not open (state=%d)\n", direction, request_id, target_conn->state);
                        }
                    }
                }
            } else {
                if (state->debug) printf("[DEBUG - %s -> wssshd] Could not find tunnel for request %s\n", direction, request_id);
            }
        }

        if (request_id) free(request_id);
        if (data) free(data);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_close\"") || strstr(msg_copy, "\"type\": \"tunnel_close\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel close\n", direction);
        }
        // Handle tunnel close from either wssshc or wsssht
        char *request_id = NULL;

        char *close_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (close_request_id_start) {
            close_request_id_start += strlen("\"request_id\":\"");
            char *close_request_id_end = strchr(close_request_id_start, '"');
            if (close_request_id_end) {
                size_t close_request_id_len = close_request_id_end - close_request_id_start;
                char *close_request_id_copy = malloc(close_request_id_len + 1);
                if (close_request_id_copy) {
                    memcpy(close_request_id_copy, close_request_id_start, close_request_id_len);
                    close_request_id_copy[close_request_id_len] = '\0';
                    request_id = close_request_id_copy;
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Extracted tunnel_close request_id: '%s'\n", direction, request_id);
                }
            }
        }

        if (request_id) {
            // Find the tunnel and forward close to the other end
            tunnel_t *tunnel = websocket_find_tunnel(state, request_id);
            if (tunnel) {
                // Determine target based on sender
                ws_connection_t *target_conn = NULL;
                const char *target_side = NULL;

                if (conn == tunnel->client_ws) {
                    // Message from wssshc, forward to wsssht
                    target_conn = tunnel->wsssh_ws;
                    target_side = "wsssht";
                } else if (conn == tunnel->wsssh_ws) {
                    // Message from wsssht, forward to wssshc
                    target_conn = tunnel->client_ws;
                    target_side = "wssshc";
                }

                if (target_conn && ((ws_connection_t *)target_conn)->state == WS_STATE_OPEN) {
                    char close_msg[256];
                    snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id);
                    if (state->debug) printf("[DEBUG - wssshd -> %s] Forwarding tunnel close: %s\n", target_side, close_msg);
                    bool send_result = ws_send_frame(target_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                    if (send_result) {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Forwarded tunnel_close for request %s to %s\n", direction, request_id, target_side);
                    } else {
                        if (state->debug) printf("[DEBUG - %s -> wssshd] Failed to send tunnel_close for request %s to %s (connection may be broken)\n", direction, request_id, target_side);
                    }
                } else if (target_conn && ((ws_connection_t *)target_conn)->state != WS_STATE_OPEN) {
                    if (state->debug) printf("[DEBUG - %s -> wssshd] Not forwarding tunnel_close for request %s - %s connection is not open (state=%d)\n", direction, request_id, target_side ? target_side : "target", target_conn ? (int)((ws_connection_t *)target_conn)->state : -1);
                }
            }

            // Remove the tunnel
            websocket_remove_tunnel(state, request_id);
            if (state->debug) printf("[DEBUG - %s -> wssshd] Closed tunnel %s\n", direction, request_id);
            if (!state->debug) log_message("EVENT", "Tunnel %s closed\n", request_id);
        }

        if (request_id) free(request_id);
    } else if (strstr(msg_copy, "\"type\":\"tunnel_keepalive\"") || strstr(msg_copy, "\"type\": \"tunnel_keepalive\"")) {
        if (state->debug) {
            printf("[DEBUG - %s -> wssshd] Processing tunnel keepalive\n", direction);
        }
        // Handle tunnel keepalive from wssshc or wsssht
        char *request_id = NULL;
        unsigned long long total_bytes = 0;
        double rate_bps = 0.0;

        // Extract request_id
        char *ka_request_id_start = strstr(msg_copy, "\"request_id\":\"");
        if (ka_request_id_start) {
            ka_request_id_start += strlen("\"request_id\":\"");
            char *ka_request_id_end = strchr(ka_request_id_start, '"');
            if (ka_request_id_end) {
                size_t len = ka_request_id_end - ka_request_id_start;
                request_id = malloc(len + 1);
                if (request_id) {
                    memcpy(request_id, ka_request_id_start, len);
                    request_id[len] = '\0';
                }
            }
        }

        // Extract total_bytes
        char *tb_start = strstr(msg_copy, "\"total_bytes\":");
        if (tb_start) {
            tb_start += strlen("\"total_bytes\":");
            total_bytes = strtoull(tb_start, NULL, 10);
        }

        // Extract rate_bps
        char *rb_start = strstr(msg_copy, "\"rate_bps\":");
        if (rb_start) {
            rb_start += strlen("\"rate_bps\":");
            rate_bps = strtod(rb_start, NULL);
        }

        if (request_id) {
            tunnel_t *tunnel = websocket_find_tunnel(state, request_id);
            if (tunnel) {
                // Update keepalive time
                if (conn == tunnel->client_ws) {
                    tunnel->last_keepalive_from_client = time(NULL);
                } else if (conn == tunnel->wsssh_ws) {
                    tunnel->last_keepalive_from_tool = time(NULL);
                }

                // Send ack
                char ack_msg[256];
                snprintf(ack_msg, sizeof(ack_msg), "{\"type\":\"tunnel_keepalive_ack\",\"request_id\":\"%s\"}", request_id);
                if (state->debug) printf("[DEBUG - wssshd -> %s] Sending keepalive ack: %s\n", direction, ack_msg);
                ws_send_frame(conn, WS_OPCODE_TEXT, ack_msg, strlen(ack_msg));

                if (state->debug) {
                    printf("[DEBUG - %s -> wssshd] Tunnel keepalive for %s: total_bytes=%llu, rate_bps=%.2f\n", direction, request_id, total_bytes, rate_bps);
                }
            } else {
                if (state->debug) {
                    printf("[DEBUG - %s -> wssshd] Keepalive for unknown tunnel %s\n", direction, request_id);
                }
            }
            free(request_id);
        }
    } else {
        // Check if any plugin handles this message
        int plugin_result = plugin_handle_message(state, conn, msg_copy, message_len);
        if (plugin_result == 0) {
            if (state->debug) {
                printf("[DEBUG - %s -> wssshd] Plugin handled message: %s\n", direction, msg_copy);
            }
        } else {
            if (state->debug) {
                printf("[DEBUG - %s -> wssshd] Unhandled message type: %s\n", direction, msg_copy);
            }
        }
    }
    // TODO: Handle other message types with similar string parsing

    free(msg_copy);
    return result;
}

// Cleanup functions
void websocket_cleanup_expired_clients(wssshd_state_t *state) {
    time_t current_time = time(NULL);
    for (size_t i = 0; i < state->clients_count; i++) {
        if (!state->clients[i].active && current_time - state->clients[i].last_seen > 30) {
            // Remove expired client
            memmove(&state->clients[i], &state->clients[i + 1],
                    (state->clients_count - i - 1) * sizeof(client_t));
            state->clients_count--;
            i--; // Adjust index after removal
        }
    }
}

void websocket_check_keepalive_timeouts(wssshd_state_t *state) {
    time_t current_time = time(NULL);
    int timeout_seconds = 220; // 4 * 30 + 10 seconds margin

    for (size_t i = 0; i < state->tunnels_count; i++) {
        tunnel_t *tunnel = state->tunnels[i];
        if (tunnel->status == TUNNEL_STATUS_ACTIVE) {
            bool client_timeout = current_time - tunnel->last_keepalive_from_client > timeout_seconds;
            bool tool_timeout = current_time - tunnel->last_keepalive_from_tool > timeout_seconds;

            if (client_timeout || tool_timeout) {
                const char *timeout_reason = client_timeout && tool_timeout ? "both endpoints" :
                                            client_timeout ? "client (wssshc) endpoint" : "tool (wsssht/wsscp) endpoint";

                if (state->debug) {
                    printf("[DEBUG] Keep-alive timeout for tunnel %s from %s\n", tunnel->request_id, timeout_reason);
                }

                tunnel_update_status(tunnel, TUNNEL_STATUS_ERROR, "Keep-alive timeout");
                // TODO: Send close messages and clean up
            }
        }
    }
}

// Thread arguments structure
typedef struct {
    ws_connection_t *conn;
    wssshd_state_t *state;
} client_thread_args_t;

// Global server state
static int server_sock = -1;
static SSL_CTX *ssl_ctx = NULL;
static volatile int server_running = 0;

// Simplified wrapper for WebSocket operations
static int safe_websocket_operation(int (*operation)(void *), void *arg) {
    return operation(arg);
}

// Wrapper function for ws_receive_frame
static int safe_ws_receive_frame(void *arg) {
    struct {
        ws_connection_t *conn;
        uint8_t *opcode;
        void **data;
        size_t *len;
    } *args = arg;

    return ws_receive_frame(args->conn, args->opcode, args->data, args->len) ? 0 : -1;
}

// Client connection thread with crash protection
static void *client_handler_thread(void *arg) {
    client_thread_args_t *args = (client_thread_args_t *)arg;

    ws_connection_t *conn = args->conn;
    wssshd_state_t *state = args->state;
    free(args); // Free the args structure

    // Perform WebSocket handshake
    if (!ws_perform_handshake(conn)) {
        fprintf(stderr, "[ERROR] WebSocket handshake failed\n");
        ws_connection_free(conn);
        return NULL;
    }

    printf("WebSocket connection established\n");

    // Determine connection direction for debug messages
    const char *direction = "unknown";
    for (size_t i = 0; i < state->clients_count; i++) {
        if (state->clients[i].websocket == conn) {
            direction = "wssshc";
            break;
        }
    }
    if (strcmp(direction, "unknown") == 0) {
        for (size_t i = 0; i < state->tunnels_count; i++) {
            if (state->tunnels[i]->wsssh_ws == conn) {
                direction = "wsssht";
                break;
            }
        }
    }

    // Handle WebSocket messages
    while (server_running && conn->state == WS_STATE_OPEN) {
        uint8_t opcode = 0;
        void *data = NULL;
        size_t len = 0;

        // Prepare arguments for safe wrapper
        struct {
            ws_connection_t *conn;
            uint8_t *opcode;
            void **data;
            size_t *len;
        } ws_args = {conn, &opcode, &data, &len};

        // Receive frame
        int receive_result = safe_websocket_operation(safe_ws_receive_frame, &ws_args);

        if (receive_result == 0) {
            if (state->debug) {
                printf("[DEBUG - %s -> wssshd] Received WebSocket frame: opcode=%d, len=%zu\n", direction, opcode, len);
            }

            if (opcode == WS_OPCODE_TEXT && len > 0 && len < 10 * 1024 * 1024) { // 10MB safety limit
                // Handle text message with additional safety
                char *message = (char *)data;
                if (message) {
                    // Update direction if register or tunnel_request message
                    if (strstr(message, "\"type\":\"register\"") || strstr(message, "\"type\": \"register\"")) {
                        direction = "wssshc";
                    } else if (strstr(message, "\"type\":\"tunnel_request\"") || strstr(message, "\"type\": \"tunnel_request\"")) {
                        direction = "wsssht";
                    }
                    if (state->debug) {
                        // Check if this is a tunnel_data or tunnel_response message to avoid printing large data
                        if (strstr(message, "\"type\":\"tunnel_data\"") || strstr(message, "\"type\": \"tunnel_data\"") ||
                            strstr(message, "\"type\":\"tunnel_response\"") || strstr(message, "\"type\": \"tunnel_response\"")) {
                            // Extract request_id
                            char *request_id_start = strstr(message, "\"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';
                                            }
                                        }
                                    }
                                }
                            }
                            // Extract size value
                            size_t size_val = 0;
                            char *size_start = strstr(message, "\"size\"");
                            if (size_start) {
                                char *size_colon = strchr(size_start, ':');
                                if (size_colon) {
                                    size_val = strtoul(size_colon + 1, NULL, 10);
                                }
                            }
                            const char *msg_type = strstr(message, "tunnel_data") ? "tunnel_data" : "tunnel_response";
                            printf("[DEBUG - %s -> wssshd] Received message: {\"type\":\"%s\",\"request_id\":\"%s\",\"size\":%zu,\"data\":\"<size: %zu bytes>\"}\n",
                                    direction, msg_type, request_id, size_val, size_val);
                        } else {
                            printf("[DEBUG - %s -> wssshd] Received message: %s\n", direction, message);
                        }
                    }
                    // Handle message with crash protection
                    websocket_handle_message(state, conn, message, len);
                }
            } else if (opcode == WS_OPCODE_CLOSE) {
                // Handle close frame
                if (state->debug) {
                    printf("[DEBUG - %s -> wssshd] Received close frame\n", direction);
                }
                conn->state = WS_STATE_CLOSED;
            } else if (opcode == WS_OPCODE_PING) {
                // Respond with pong
                if (state->debug) {
                    printf("[DEBUG - %s -> wssshd] Received ping, sending pong\n", direction);
                }
                // Send pong frame with improved retry logic for robustness
                int pong_retries = 0;
                const int max_pong_retries = 5; // Increased retries
                bool pong_sent = false;

                while (!pong_sent && pong_retries < max_pong_retries) {
                    if (ws_send_frame(conn, WS_OPCODE_PONG, data, len)) {
                        pong_sent = true;
                        if (state->debug && pong_retries > 0) {
                            printf("[DEBUG - %s -> wssshd] Pong sent successfully after %d retries\n", direction, pong_retries);
                        }
                    } else {
                        pong_retries++;
                        if (pong_retries < max_pong_retries) {
                            if (state->debug) {
                                printf("[DEBUG - %s -> wssshd] Pong failed, retrying (%d/%d)\n", direction, pong_retries, max_pong_retries);
                            }
                            // Exponential backoff for pong retries
                            usleep(20000 * (1 << pong_retries)); // 20ms, 40ms, 80ms, 160ms
                        }
                    }
                }

                if (!pong_sent) {
                    fprintf(stderr, "[ERROR] Failed to send pong frame after %d retries, connection may be unstable\n", max_pong_retries);
                    // Mark connection as potentially unstable but don't close immediately
                    // The connection health will be monitored through keepalive timeouts
                }
            } else {
                if (state->debug) {
                    printf("[DEBUG - %s -> wssshd] Received unhandled opcode: %d\n", direction, opcode);
                }
            }

            // Safe cleanup
            if (data) {
                free(data);
                data = NULL;
            }
        } else {
            // Connection error
            if (state->debug) {
                printf("[DEBUG - %s -> wssshd] WebSocket frame receive failed\n", direction);
            }
            // Mark connection as closed immediately to prevent further sends
            conn->state = WS_STATE_CLOSED;
            break;
        }
    }

    // Clean up any tunnels that reference this connection before closing
    if (state) {
        // Find and clean up tunnels for this connection
        pthread_mutex_lock(&state->tunnel_mutex);
        for (size_t i = 0; i < state->tunnels_count; ) {
            tunnel_t *tunnel = state->tunnels[i];
            bool tunnel_removed = false;

            if (tunnel && (tunnel->client_ws == conn || tunnel->wsssh_ws == conn)) {
                // This tunnel uses the closing connection
                // Try to send tunnel_close to the other end if possible
                ws_connection_t *other_conn = NULL;
                if (tunnel->client_ws == conn && tunnel->wsssh_ws && ((ws_connection_t *)tunnel->wsssh_ws)->state == WS_STATE_OPEN) {
                    other_conn = tunnel->wsssh_ws;
                } else if (tunnel->wsssh_ws == conn && tunnel->client_ws && ((ws_connection_t *)tunnel->client_ws)->state == WS_STATE_OPEN) {
                    other_conn = tunnel->client_ws;
                }

                // Save request_id before freeing
                char request_id_copy[64];
                strncpy(request_id_copy, tunnel->request_id, sizeof(request_id_copy) - 1);
                request_id_copy[sizeof(request_id_copy) - 1] = '\0';

                if (other_conn) {
                    char close_msg[256];
                    snprintf(close_msg, sizeof(close_msg), "{\"type\":\"tunnel_close\",\"request_id\":\"%s\"}", request_id_copy);
                    ws_send_frame(other_conn, WS_OPCODE_TEXT, close_msg, strlen(close_msg));
                }

                // Remove the tunnel
                tunnel_free(tunnel);
                // Shift remaining elements
                memmove(&state->tunnels[i], &state->tunnels[i + 1],
                        (state->tunnels_count - i - 1) * sizeof(tunnel_t *));
                state->tunnels_count--;
                tunnel_removed = true;

                if (state->debug) {
                    printf("[DEBUG] Cleaned up tunnel %s due to websocket connection closure\n", request_id_copy);
                }
            }

            if (!tunnel_removed) {
                i++;
            }
        }
        pthread_mutex_unlock(&state->tunnel_mutex);

        // Remove client registration if this was a wssshc connection
        pthread_mutex_lock(&state->client_mutex);
        for (size_t i = 0; i < state->clients_count; ) {
            if (state->clients[i].websocket == conn) {
                // Remove this client
                if (!state->debug) {
                    log_message("EVENT", "Client %s disconnected\n", state->clients[i].client_id);
                }
                memmove(&state->clients[i], &state->clients[i + 1],
                        (state->clients_count - i - 1) * sizeof(client_t));
                state->clients_count--;
                if (state->debug) {
                    printf("[DEBUG] Removed client registration due to websocket connection closure\n");
                }
                break;
            } else {
                i++;
            }
        }
        pthread_mutex_unlock(&state->client_mutex);
    }

    printf("WebSocket connection closed\n");
    ws_connection_free(conn);
    return NULL;
}

// Status printing thread
void *status_thread(void *arg) {
    wssshd_state_t *state = (wssshd_state_t *)arg;
    static time_t last_status_time = 0;
    static unsigned long long last_total_data[100] = {0}; // Track last total data for rate calculation

    while (server_running) {
        sleep(15);
        if (!state->debug) {
            time_t current_time = time(NULL);
            time_t time_since_last = last_status_time == 0 ? 15 : current_time - last_status_time;
            time_t uptime = current_time - state->start_time;
            int days = uptime / 86400;
            int hours = (uptime % 86400) / 3600;
            int minutes = (uptime % 3600) / 60;
            int seconds = uptime % 60;

            // Count active clients and build client list with services
            pthread_mutex_lock(&state->client_mutex);
            int client_count = 0;
            char client_list[2048] = "";
            for (size_t i = 0; i < state->clients_count; i++) {
                if (state->clients[i].active) {
                    client_count++;
                    if (strlen(client_list) > 0) strncat(client_list, ", ", sizeof(client_list) - strlen(client_list) - 1);
                    size_t remaining = sizeof(client_list) - strlen(client_list) - 1;
                    if (remaining > 0) {
                        strncat(client_list, state->clients[i].client_id, remaining);
                        remaining = sizeof(client_list) - strlen(client_list) - 1;
                        if (remaining > 2) { // space for (services)
                            strncat(client_list, "(", remaining);
                            remaining = sizeof(client_list) - strlen(client_list) - 1;
                            strncat(client_list, state->clients[i].services, remaining);
                            remaining = sizeof(client_list) - strlen(client_list) - 1;
                            if (remaining > 1) {
                                strncat(client_list, ")", remaining);
                            }
                        }
                    }
                }
            }
            pthread_mutex_unlock(&state->client_mutex);

            pthread_mutex_lock(&state->tunnel_mutex);
            int tunnel_count = state->tunnels_count;

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

            // List active tunnels with details
            for (size_t i = 0; i < state->tunnels_count && i < 100; i++) {
                tunnel_t *tunnel = state->tunnels[i];
                if (tunnel && tunnel->status == TUNNEL_STATUS_ACTIVE) {
                    time_t tunnel_uptime = current_time - tunnel->created_at;
                    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;
                    unsigned long long data_since_last = total_data - last_total_data[i];
                    double last_rate = time_since_last > 0 ? (double)data_since_last / time_since_last : 0.0;
                    last_total_data[i] = total_data;

                    int days = tunnel_uptime / 86400;
                    int hours = (tunnel_uptime % 86400) / 3600;
                    int minutes = (tunnel_uptime % 3600) / 60;
                    int 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 (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 : "unknown",
                                    strlen(tunnel->service) > 0 ? tunnel->service : "ssh",
                                    days, hours, minutes, 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 : "unknown",
                                    strlen(tunnel->service) > 0 ? tunnel->service : "ssh",
                                    hours, minutes, seconds,
                                    total_data_str,
                                    sent_data_str,
                                    recv_data_str,
                                    avg_rate_str,
                                    last_rate_str);
                    }
                }
            }

            last_status_time = current_time;
            pthread_mutex_unlock(&state->tunnel_mutex);
        }
    }
    return NULL;
}

// WebSocket server functions
int websocket_start_server(const wssshd_config_t *config, wssshd_state_t *state) {
    printf("Starting WebSocket server on %s:%d (TLS-only)\n", config->host, config->port);

    // Create TCP socket
    server_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (server_sock < 0) {
        perror("Failed to create socket");
        return -1;
    }

    // Set socket options
    int opt = 1;
    setsockopt(server_sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    // Bind to address
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(config->port);
    if (inet_pton(AF_INET, config->host, &addr.sin_addr) <= 0) {
        perror("Invalid address");
        close(server_sock);
        return -1;
    }

    if (bind(server_sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("Bind failed");
        close(server_sock);
        return -1;
    }

    // Listen
    if (listen(server_sock, 10) < 0) {
        perror("Listen failed");
        close(server_sock);
        return -1;
    }

    // Initialize SSL context (always required for wssshc client connections)
    printf("Creating SSL context...\n");
    ssl_ctx = ssl_create_context();
    if (!ssl_ctx) {
        fprintf(stderr, "Failed to create SSL context\n");
        close(server_sock);
        return -1;
    }
    printf("SSL context created successfully\n");

    // Load or generate certificates
    char cert_file[256], key_file[256];
    bool use_system_location = false;

    // Check if /etc/wssshd directory exists
    if (access("/etc/wssshd", F_OK) == 0) {
        // Try to use system location
        snprintf(cert_file, sizeof(cert_file), "/etc/wssshd/cert.pem");
        snprintf(key_file, sizeof(key_file), "/etc/wssshd/key.pem");
        use_system_location = true;
    } else {
        // Fall back to local directory
        char local_dir[256];
        snprintf(local_dir, sizeof(local_dir), "./certs");
        mkdir(local_dir, 0755); // Create directory if it doesn't exist
        snprintf(cert_file, sizeof(cert_file), "./certs/cert.pem");
        snprintf(key_file, sizeof(key_file), "./certs/key.pem");
        use_system_location = false;
    }

    printf("Using certificates: %s, %s\n", cert_file, key_file);

    if (access(cert_file, F_OK) != 0 || access(key_file, F_OK) != 0) {
        printf("Generating self-signed certificate...\n");
        if (ssl_generate_self_signed_cert(cert_file, key_file) != 0) {
            if (use_system_location) {
                // If system location failed, fall back to local directory
                printf("Failed to generate certificate in system location, falling back to local directory...\n");
                char local_dir[256];
                snprintf(local_dir, sizeof(local_dir), "./certs");
                mkdir(local_dir, 0755); // Create directory if it doesn't exist
                snprintf(cert_file, sizeof(cert_file), "./certs/cert.pem");
                snprintf(key_file, sizeof(key_file), "./certs/key.pem");
                printf("Using certificates: %s, %s\n", cert_file, key_file);

                if (ssl_generate_self_signed_cert(cert_file, key_file) != 0) {
                    fprintf(stderr, "Failed to generate certificate in fallback location\n");
                    SSL_CTX_free(ssl_ctx);
                    close(server_sock);
                    return -1;
                }
                printf("Certificate generated successfully in fallback location\n");
            } else {
                fprintf(stderr, "Failed to generate certificate\n");
                SSL_CTX_free(ssl_ctx);
                close(server_sock);
                return -1;
            }
        } else {
            printf("Certificate generated successfully\n");
        }
    }

    printf("Loading certificates...\n");
    if (ssl_load_certificates(ssl_ctx, cert_file, key_file) != 0) {
        fprintf(stderr, "Failed to load certificates\n");
        SSL_CTX_free(ssl_ctx);
        close(server_sock);
        return -1;
    }
    printf("Certificates loaded successfully\n");

    printf("WebSocket server listening on %s:%d\n", config->host, config->port);
    server_running = 1;

    // Start status printing thread
    pthread_t status_tid;
    pthread_create(&status_tid, NULL, status_thread, state);
    pthread_detach(status_tid);

    printf("Entering accept loop...\n");

    // Accept connections in a loop
    while (server_running) {
        printf("Waiting for connections...\n");
        struct sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);
        int client_sock = accept(server_sock, (struct sockaddr *)&client_addr, &client_len);

        if (client_sock < 0) {
            if (server_running) perror("Accept failed");
            continue;
        }

        // TLS-only mode: always use SSL
        SSL *ssl = SSL_new(ssl_ctx);
        if (!ssl) {
            close(client_sock);
            continue;
        }

        SSL_set_fd(ssl, client_sock);

        if (SSL_accept(ssl) <= 0) {
            ERR_print_errors_fp(stderr);
            SSL_free(ssl);
            close(client_sock);
            continue;
        }

        // Create WebSocket connection with SSL
        ws_connection_t *conn = ws_connection_create(ssl, client_sock);
        if (!conn) {
            // ws_connection_create failed, clean up SSL and socket
            SSL_free(ssl);
            close(client_sock);
            continue;
        }

        // Set the client address
        conn->addr = client_addr;

        // Start client handler thread
        client_thread_args_t *thread_args = malloc(sizeof(client_thread_args_t));

        if (!thread_args) {
            fprintf(stderr, "Failed to allocate thread args\n");
            ws_connection_free(conn);
            continue;
        }

        thread_args->conn = conn;
        thread_args->state = state;

        pthread_t thread;
        if (pthread_create(&thread, NULL, client_handler_thread, thread_args) != 0) {
            fprintf(stderr, "Failed to create client handler thread\n");
            free(thread_args);
            ws_connection_free(conn);
            continue;
        }

        pthread_detach(thread);
    }

    return 0;
}

void websocket_stop_server(void) {
    server_running = 0;

    if (server_sock >= 0) {
        close(server_sock);
        server_sock = -1;
    }

    if (ssl_ctx) {
        SSL_CTX_free(ssl_ctx);
        ssl_ctx = NULL;
    }

    printf("WebSocket server stopped\n");
}