/**
 * Main server logic 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 <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include <time.h>
#include "config.h"
#include "websocket.h"
#include "web.h"
#include "web_proxy.h"
#include "ssl.h"
#include "plugin.h"

static volatile int shutdown_requested = 0;

// Signal handler for graceful shutdown
static void signal_handler(int signum) {
    if (shutdown_requested) {
        // Force exit on second signal
        fprintf(stderr, "\nReceived second signal, exiting immediately...\n");
        exit(1);
    }

    shutdown_requested = 1;
    printf("\nReceived signal %d, initiating graceful shutdown...\n", signum);
}

// Cleanup thread function
static void *cleanup_thread(void *arg) {
    wssshd_state_t *state = (wssshd_state_t *)arg;

    while (!shutdown_requested) {
        // Run cleanup tasks every 5 seconds
        sleep(5);

        websocket_cleanup_expired_clients(state);
        websocket_check_keepalive_timeouts(state);

        // Print status every 60 seconds (only when not in debug mode)
        static time_t last_status_time = 0;
        time_t current_time = time(NULL);
        if (!state->debug && current_time - last_status_time >= 60) {
            size_t active_clients = 0;
            for (size_t i = 0; i < state->clients_count; i++) {
                if (state->clients[i].active) active_clients++;
            }

            size_t active_tunnels = 0;
            unsigned long long total_transfer_rate = 0;
            for (size_t i = 0; i < state->tunnels_count; i++) {
                if (state->tunnels[i]->status == TUNNEL_STATUS_ACTIVE) {
                    active_tunnels++;
                    total_transfer_rate += state->tunnels[i]->bytes_last_period;
                    // Reset bytes_last_period for next period
                    state->tunnels[i]->bytes_last_period = 0;
                    state->tunnels[i]->last_stats_reset = current_time;
                }
            }

            time_t uptime = current_time - state->start_time;
            int hours = uptime / 3600;
            int minutes = (uptime % 3600) / 60;
            int seconds = uptime % 60;

            printf("[STATUS] Server uptime: %02d:%02d:%02d\n", hours, minutes, seconds);
            printf("[STATUS] Connected wssshc clients (%zu):\n", active_clients);
            for (size_t i = 0; i < state->clients_count; i++) {
                if (state->clients[i].active) {
                    printf("[STATUS]   - %s\n", state->clients[i].client_id);
                }
            }
            printf("[STATUS] Active tunnels (%zu/%zu):\n", active_tunnels, state->tunnels_count);
            for (size_t i = 0; i < state->tunnels_count; i++) {
                if (state->tunnels[i]->status == TUNNEL_STATUS_ACTIVE) {
                    time_t tunnel_uptime = current_time - state->tunnels[i]->created_at;
                    int t_hours = tunnel_uptime / 3600;
                    int t_minutes = (tunnel_uptime % 3600) / 60;
                    int t_seconds = tunnel_uptime % 60;
                    printf("[STATUS]   - %s: uptime %02d:%02d:%02d\n",
                           state->tunnels[i]->request_id, t_hours, t_minutes, t_seconds);
                }
            }
            printf("[STATUS] Transfer rate since last statistic: %llu bytes\n", total_transfer_rate);

            last_status_time = current_time;
        }
    }

    return NULL;
}

int main(int argc, char *argv[]) {
    // Create new process group to avoid receiving SIGINT from terminal
    setpgid(0, 0);

    // Parse configuration
    wssshd_config_t *config = load_config(argc, argv);
    if (!config) {
        fprintf(stderr, "Failed to load configuration\n");
        return 1;
    }

    printf("WSSSH Daemon starting...\n");

    // Print ASCII art banner
    printf("\033[38;5;117m⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠⠖⠢⢄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;117m⠀⠀⠀⠀⠀⠀⠀⢠⠖⠋⠀⠀⠀⣀⣹⡆⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;117m⠀⠀⠀⠀⠀⠀⠀⢸⡄⠀⢠⣶⡞⢁⣸⣇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⠀⠀⠀⠀⠀⠀⠈⢁⣠⠞⢻⡁⢻⣷⣾⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⠀⠀⠀⠀⣠⠤⢄⣘⣿⢬⣛⡷⢋⣈⣉⣀⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⠀⠀⠀⢰⡇⠀⠀⠙⠿⠀⢀⠿⠋⠀⠀⠀⢱⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("⠀⠀⠀⠀⠘⡃⠀⢰⠀⠀⠀⡁⠀⡀⠀⠘⣆⠈⢇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\n");
    printf("\033[38;5;231m⠀⠀⠀⠀⠀⡇⠀⣼⡦⠠⠤⣈⡆⢃⡤⠒⠙⡆⠀⠳⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;231m⠀⠀⠀⠀⢀⠇⠀⡇⢷⣄⣀⡠⠟⠛⠢⠤⠞⡟⠦⡀⠙⠦⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;231m⠀⠀⠀⠀⡞⢀⠀⠀⠀⠳⡄⠀⠀⠠⡀⠀⠀⠘⠉⠙⠦⣀⠀⠉⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;231m⠀⠀⠀⠀⡇⠀⢠⠀⠀⠀⠈⠓⢄⠀⠁⠀⠀⠀⠒⠂⠀⣾⠋⠉⠒⠢⢍⣙⡒⢄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⠀⠀⠀⡇⠀⡜⠀⠀⠀⠀⠀⠀⠱⡀⢀⡀⠁⠀⢀⡼⡇⠀⠀⠀⠀⠀⣏⠙⠯⣆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⠀⠀⠀⡇⢠⠇⠀⠀⠀⠀⠀⠀⠀⢱⣀⣽⠶⠾⠛⠒⠛⠒⠒⠒⠤⠤⣸⡍⠀⠀⠉⠲⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;218m⠀⣀⣀⣀⠇⣼⠀⠀⠀⠀⠀⠀⠀⠀⠀⢫⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠠⣙⠢⣀⠀⠀⠈⠳⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;117m⠉⠙⠛⠥⠰⠟⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⠂⠤⠤⠤⠤⠴⠒⠚⠉⠙⠢⣀⠀⠈⠑⠢⢄⡀⠈⠳⡀⠀⠀⠀⠀⠀⠀⠀\033[0m\n");
    printf("\033[38;5;117m⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠑⠒⠒⠒⠒⠚⠑⠢⢌⡓⠤⠤⠤⠤⣀⠀\033[0m\n");
    printf("\033[38;5;117m⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⠒⠠⠤⣼⠇\033[0m\n");
    printf("\n");

    if (config->debug) {
        print_config(config);
    }

    // Initialize state
    wssshd_state_t *state = websocket_init_state(config->debug, config->password);
    if (!state) {
        fprintf(stderr, "Failed to initialize server state\n");
        free_config(config);
        return 1;
    }

    // Initialize plugin system
    if (!plugin_system_init(state, config)) {
        fprintf(stderr, "Warning: Failed to initialize plugin system\n");
    }

    // Set up signal handlers
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    signal(SIGPIPE, SIG_IGN);  // Ignore SIGPIPE to prevent crashes on broken connections

    // Start web interface if configured
    if (web_start_server(config, state) != 0) {
        fprintf(stderr, "Warning: Failed to start web interface\n");
    }

    // Start WebSocket server
    if (websocket_start_server(config, state) != 0) {
        fprintf(stderr, "Failed to start WebSocket server\n");
        web_stop_server();
        websocket_free_state(state);
        free_config(config);
        return 1;
    }

    // Start web proxy server if enabled
    if (config->web_proxy_enabled) {
        if (web_proxy_start_server(config, state) != 0) {
            fprintf(stderr, "Warning: Failed to start web proxy server\n");
        }
    }

    printf("WSSSH Daemon running on %s:%d\n", config->host, config->port);
    if (config->web_proxy_enabled) {
        printf("Web proxy server running on 127.0.0.1:%d\n", config->web_proxy_port);
    }
    printf("Press Ctrl+C to stop the server\n");

    // Start cleanup thread
    pthread_t cleanup_tid;
    if (pthread_create(&cleanup_tid, NULL, cleanup_thread, state) != 0) {
        fprintf(stderr, "Warning: Failed to start cleanup thread\n");
    } else {
        pthread_detach(cleanup_tid);
    }

    // Main loop - wait for shutdown signal
    while (!shutdown_requested) {
        sleep(1);
    }

    // Shutdown sequence
    printf("\nShutting down WSSSH Daemon...\n");

    // Stop servers
    web_proxy_stop_server();
    web_stop_server();
    websocket_stop_server();

    // Clean up plugin system
    plugin_system_cleanup();

    // Clean up state
    websocket_free_state(state);
    free_config(config);

    // Clean up SSL
    ssl_cleanup();

    printf("WSSSH Daemon stopped cleanly\n");
    return 0;
}