#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <getopt.h>

#define BUFFER_SIZE 4096
#define DEFAULT_PORT 9898

typedef struct {
    char *server_ip;
    int port;
    char *client_id;
    char *password;
    int interval;
    int debug;
} wssshc_config_t;

void print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options]\n", program_name);
    fprintf(stderr, "WebSocket SSH Client - Register with wssshd server\n\n");
    fprintf(stderr, "Protect the dolls!\n\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  --server-ip IP       Server IP address (required)\n");
    fprintf(stderr, "  --port PORT          Server port (default: %d)\n", DEFAULT_PORT);
    fprintf(stderr, "  --id ID              Client identifier (required)\n");
    fprintf(stderr, "  --password PASS      Registration password (required)\n");
    fprintf(stderr, "  --interval SEC       Reconnection interval (default: 30)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "\nDonations:\n");
    fprintf(stderr, "  BTC: bc1q3zlkpu95amtcltsk85y0eacyzzk29v68tgc5hx\n");
    fprintf(stderr, "  ETH: 0xdA6dAb526515b5cb556d20269207D43fcc760E51\n");
}

int parse_args(int argc, char *argv[], wssshc_config_t *config) {
    static struct option long_options[] = {
        {"server-ip", required_argument, 0, 's'},
        {"port", required_argument, 0, 'p'},
        {"id", required_argument, 0, 'i'},
        {"password", required_argument, 0, 'w'},
        {"interval", required_argument, 0, 't'},
        {"debug", no_argument, 0, 'd'},
        {"help", no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "s:p:i:w:t:dh", long_options, NULL)) != -1) {
        switch (opt) {
            case 's':
                config->server_ip = strdup(optarg);
                break;
            case 'p':
                config->port = atoi(optarg);
                break;
            case 'i':
                config->client_id = strdup(optarg);
                break;
            case 'w':
                config->password = strdup(optarg);
                break;
            case 't':
                config->interval = atoi(optarg);
                break;
            case 'd':
                config->debug = 1;
                break;
            case 'h':
            default:
                print_usage(argv[0]);
                return 0;
        }
    }

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

    // Validate required arguments
    if (!config->server_ip || !config->client_id || !config->password) {
        fprintf(stderr, "Error: --server-ip, --id, and --password are required\n");
        print_usage(argv[0]);
        return 0;
    }

    return 1;
}

SSL_CTX* init_ssl_ctx() {
    SSL_CTX *ctx;

    SSL_library_init();
    OpenSSL_add_all_algorithms();
    SSL_load_error_strings();

    ctx = SSL_CTX_new(TLS_client_method());
    if (!ctx) {
        ERR_print_errors_fp(stderr);
        return NULL;
    }

    // Allow self-signed certificates
    SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, NULL);

    return ctx;
}

int websocket_handshake(int sock, const char *host, int port, const char *path) {
    char request[1024];
    char response[BUFFER_SIZE];
    int bytes_read;

    // Send WebSocket handshake
    snprintf(request, sizeof(request),
             "GET %s HTTP/1.1\r\n"
             "Host: %s:%d\r\n"
             "Upgrade: websocket\r\n"
             "Connection: Upgrade\r\n"
             "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
             "Sec-WebSocket-Version: 13\r\n"
             "\r\n",
             path, host, port);

    if (send(sock, request, strlen(request), 0) < 0) {
        perror("WebSocket handshake send failed");
        return 0;
    }

    // Read response
    bytes_read = recv(sock, response, sizeof(response) - 1, 0);
    if (bytes_read <= 0) {
        perror("WebSocket handshake recv failed");
        return 0;
    }

    response[bytes_read] = '\0';

    // Check for successful handshake
    if (strstr(response, "101 Switching Protocols") == NULL) {
        fprintf(stderr, "WebSocket handshake failed\n");
        return 0;
    }

    return 1;
}

int send_json_message(int sock, const char *type, const char *id, const char *password) {
    char message[1024];

    if (password) {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"id\":\"%s\",\"password\":\"%s\"}",
                 type, id, password);
    } else {
        snprintf(message, sizeof(message),
                 "{\"type\":\"%s\",\"id\":\"%s\"}",
                 type, id);
    }

    if (send(sock, message, strlen(message), 0) < 0) {
        perror("Send failed");
        return 0;
    }

    return 1;
}

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

    // Resolve hostname
    if ((he = gethostbyname(config->server_ip)) == NULL) {
        herror("gethostbyname");
        return 0;
    }

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

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

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

    // Initialize SSL
    ssl_ctx = init_ssl_ctx();
    if (!ssl_ctx) {
        close(sock);
        return 0;
    }

    ssl = SSL_new(ssl_ctx);
    SSL_set_fd(ssl, sock);

    if (SSL_connect(ssl) <= 0) {
        ERR_print_errors_fp(stderr);
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 0;
    }

    // Perform WebSocket handshake
    if (!websocket_handshake(sock, config->server_ip, config->port, "/")) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 0;
    }

    // Send registration message
    if (!send_json_message(sock, "register", config->client_id, config->password)) {
        SSL_free(ssl);
        SSL_CTX_free(ssl_ctx);
        close(sock);
        return 0;
    }

    printf("Connected and registered as %s\n", config->client_id);

    // Keep connection alive and handle messages
    while (1) {
        bytes_read = SSL_read(ssl, buffer, sizeof(buffer) - 1);
        if (bytes_read <= 0) {
            if (config->debug) {
                printf("[DEBUG] Connection lost\n");
            }
            break;
        }

        buffer[bytes_read] = '\0';

        if (config->debug) {
            printf("[DEBUG] Received: %s\n", buffer);
        }

        // Handle different message types if needed
        // For now, just keep the connection alive
    }

    // Cleanup
    SSL_free(ssl);
    SSL_CTX_free(ssl_ctx);
    close(sock);

    return 1;
}

int main(int argc, char *argv[]) {
    wssshc_config_t config = {
        .server_ip = NULL,
        .port = DEFAULT_PORT,
        .client_id = NULL,
        .password = NULL,
        .interval = 30,
        .debug = 0
    };

    if (!parse_args(argc, argv, &config)) {
        return 1;
    }

    printf("WebSocket SSH Client starting...\n");

    while (1) {
        if (connect_to_server(&config)) {
            printf("Connection successful\n");
            // In a real implementation, we might want to keep the connection alive
            // For now, just exit on successful connection
            break;
        } else {
            printf("Connection failed, retrying in %d seconds...\n", config.interval);
            sleep(config.interval);
        }
    }

    // Cleanup
    free(config.server_ip);
    free(config.client_id);
    free(config.password);

    return 0;
}