/*
 * Test VNC Client over wsssht
 * Minimal VNC protocol client for testing VNC tunneling via wsssht --pipe
 *
 * Usage: ./test_vnc_client <client_id>
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/wait.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdint.h>

#define BUFFER_SIZE 4096
#define TIMEOUT_SECONDS 20

// Signal handler for timeout
volatile sig_atomic_t timeout_flag = 0;

void timeout_handler(int sig) {
    (void)sig;
    timeout_flag = 1;
}

// Read with timeout from file descriptor
ssize_t read_with_timeout(int fd, char *buffer, size_t size, int timeout_sec) {
    timeout_flag = 0;
    signal(SIGALRM, timeout_handler);
    alarm(timeout_sec);

    ssize_t bytes_read = read(fd, buffer, size);

    alarm(0); // Cancel alarm
    signal(SIGALRM, SIG_DFL);

    if (timeout_flag) {
        return -1; // Timeout
    }

    return bytes_read;
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <client_id>\n", argv[0]);
        return 1;
    }

    const char *client_id = argv[1];

    // Construct the wsssht command
    char command[256];
    snprintf(command, sizeof(command), "wsssht --pipe vnc://%s", client_id);

    printf("Executing command: %s\n", command);

    // Create pipes for bidirectional communication
    int to_child[2];   // Parent writes to child
    int from_child[2]; // Parent reads from child

    if (pipe(to_child) == -1 || pipe(from_child) == -1) {
        perror("Failed to create pipes");
        return 1;
    }

    // Fork the process
    pid_t pid = fork();
    if (pid == -1) {
        perror("Failed to fork");
        return 1;
    }

    if (pid == 0) { // Child process
        // Close unused pipe ends
        close(to_child[1]);
        close(from_child[0]);

        // Redirect stdin and stdout
        dup2(to_child[0], STDIN_FILENO);
        dup2(from_child[1], STDOUT_FILENO);

        // Close the original pipe ends
        close(to_child[0]);
        close(from_child[1]);

        // Execute wsssht
        execlp("wsssht", "wsssht", "--pipe", "vnc://zeiss", NULL);
        perror("Failed to execute wsssht");
        _exit(1);
    } else { // Parent process
        // Close unused pipe ends
        close(to_child[0]);
        close(from_child[1]);

        printf("Connected to wsssht pipe for VNC client_id: %s\n", client_id);
    }

    // Use file descriptors for communication
    int write_fd = to_child[1];
    int read_fd = from_child[0];

    // VNC Protocol Handshake
    // Step 1: Send version message
    const char *version_msg = "RFB 003.008\n";
    size_t version_len = strlen(version_msg);

    printf("Sending VNC version message: %s", version_msg);
    if (write(write_fd, version_msg, version_len) != (ssize_t)version_len) {
        perror("Failed to send VNC version message");
        close(write_fd);
        close(read_fd);
        return 1;
    }

    // Step 2: Read server version response
    char buffer[BUFFER_SIZE];
    printf("Waiting for server version response...\n");
    ssize_t bytes_read = read_with_timeout(read_fd, buffer, sizeof(buffer) - 1, TIMEOUT_SECONDS);
    if (bytes_read > 0) {
        buffer[bytes_read] = '\0';
        printf("Received server version response (%zd bytes): ", bytes_read);
        for (ssize_t i = 0; i < bytes_read; i++) {
            if (buffer[i] >= 32 && buffer[i] <= 126) {
                putchar(buffer[i]);
            } else {
                printf("\\x%02x", (unsigned char)buffer[i]);
            }
        }
        printf("\n");

        // Check if it's a valid VNC version response
        if (bytes_read >= 12 && strncmp(buffer, "RFB ", 4) == 0) {
            printf("✓ Valid VNC version response received\n");

            // Step 3: Read security types
            printf("Reading security types...\n");
            bytes_read = read_with_timeout(read_fd, buffer, sizeof(buffer), TIMEOUT_SECONDS);
            if (bytes_read > 0) {
                printf("Received security types (%zd bytes): ", bytes_read);
                for (ssize_t i = 0; i < bytes_read && i < 10; i++) {  // Show first 10 bytes
                    printf("%02x ", (unsigned char)buffer[i]);
                }
                if (bytes_read > 10) printf("...");
                printf("\n");

                if (bytes_read >= 2) {
                    unsigned char num_types = buffer[0];
                    printf("Number of security types: %u\n", num_types);

                    if (num_types > 0 && bytes_read >= 1 + (ssize_t)num_types) {
                        printf("Available security types: ");
                        for (int i = 0; i < num_types; i++) {
                            printf("%u ", (unsigned char)buffer[1 + i]);
                        }
                        printf("\n");

                        // Select the first security type (usually None=1)
                        unsigned char selected_type = buffer[1];
                        printf("Selecting security type: %u\n", selected_type);

                        // Send security type selection
                        if (write(write_fd, &selected_type, 1) != 1) {
                            perror("Failed to send security type selection");
                            close(write_fd);
                            close(read_fd);
                            return 1;
                        }
                        printf("✓ Security type selection sent\n");

                        // For security type 1 (None), expect security result
                        if (selected_type == 1) {
                            printf("Waiting for security handshake result...\n");
                            bytes_read = read_with_timeout(read_fd, buffer, sizeof(buffer), TIMEOUT_SECONDS);
                            if (bytes_read >= 4) {
                                uint32_t result = ((unsigned char)buffer[0] << 24) |
                                                 ((unsigned char)buffer[1] << 16) |
                                                 ((unsigned char)buffer[2] << 8) |
                                                 (unsigned char)buffer[3];
                                printf("✓ Security handshake result: %u (%s)\n",
                                       result, result == 0 ? "OK" : "Failed");

                                if (result == 0) {
                                    // Send client initialization message (shared flag = 0 for exclusive access)
                                    unsigned char shared_flag = 0;
                                    printf("Sending client initialization (shared=%u)...\n", shared_flag);
                                    if (write(write_fd, &shared_flag, 1) != 1) {
                                        perror("Failed to send client initialization");
                                        close(write_fd);
                                        close(read_fd);
                                        return 1;
                                    }
                                    printf("✓ Client initialization sent\n");

                                    // At this point, VNC handshake is complete
                                    printf("✓ VNC protocol handshake completed successfully!\n");
                                    printf("The connection is now ready for VNC data exchange.\n");
                                }
                            } else {
                                printf("✗ Failed to read security handshake result\n");
                            }
                        }
                    } else {
                        printf("✗ Invalid security types response\n");
                    }
                } else {
                    printf("✗ Security types response too short\n");
                }
            } else {
                printf("✗ Timeout or error reading security types\n");
            }
        } else {
            printf("✗ Unexpected version response format\n");
        }
    } else {
        printf("✗ Timeout or error reading server version response\n");
    }

    // Close pipes
    close(write_fd);
    close(read_fd);

    // Wait for child process
    int status;
    waitpid(pid, &status, 0);
    if (WIFEXITED(status)) {
        printf("wsssht exited with status %d\n", WEXITSTATUS(status));
    } else if (WIFSIGNALED(status)) {
        printf("wsssht terminated by signal %d\n", WTERMSIG(status));
    }

    return 0;
}