/**
 * RDP and WebSocket proxy 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 <unistd.h>
#include <fcntl.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <signal.h>
#include <uuid/uuid.h>
#include <pthread.h>
#include "rdp.h"
#include "config.h"

// Thread function to read pipe output
static void *read_pipe_output(void *arg) {
    rdp_session_t *session = (rdp_session_t *)arg;

    char buffer[1024];
    fd_set readfds;
    struct timeval tv;

    while (1) {
        FD_ZERO(&readfds);
        FD_SET(session->read_fd, &readfds);

        tv.tv_sec = 0;
        tv.tv_usec = 100000; // 100ms timeout

        int ret = select(session->read_fd + 1, &readfds, NULL, NULL, &tv);
        if (ret < 0) break;
        if (ret == 0) continue; // timeout

        if (FD_ISSET(session->read_fd, &readfds)) {
            ssize_t bytes_read = read(session->read_fd, buffer, sizeof(buffer));
            if (bytes_read <= 0) break;

            if (session->debug_rdp) {
                printf("RDP: Received %zd bytes from pipe\n", bytes_read);
                // Debug: print first 20 bytes in hex
                printf("RDP received data (first 20 bytes): ");
                for (ssize_t i = 0; i < bytes_read && i < 20; i++) {
                    printf("%02x ", (unsigned char)buffer[i]);
                }
                printf("\n");
            }

            // Append to session output buffer
            pthread_mutex_lock(&session->output_mutex);
            if (session->output_used + bytes_read > session->output_size) {
                size_t new_size = session->output_size * 2;
                while (new_size < session->output_used + bytes_read) {
                    new_size *= 2;
                }
                char *new_buf = realloc(session->output_buffer, new_size);
                if (new_buf) {
                    session->output_buffer = new_buf;
                    session->output_size = new_size;
                } else {
                    // Out of memory, skip this data
                    pthread_mutex_unlock(&session->output_mutex);
                    continue;
                }
            }
            memcpy(session->output_buffer + session->output_used, buffer, bytes_read);
            session->output_used += bytes_read;
            pthread_mutex_unlock(&session->output_mutex);
        }
    }

    return NULL;
}

rdp_session_t *rdp_create_session(const wssshd_config_t *config, const char *client_id, const char *username, bool debug_rdp) {
    rdp_session_t *session = calloc(1, sizeof(rdp_session_t));
    if (!session) return NULL;

    // Generate UUID for request_id
    uuid_t uuid;
    uuid_generate(uuid);
    uuid_unparse(uuid, session->request_id);

    // Copy client_id and username
    if (client_id) {
        strncpy(session->client_id, client_id, sizeof(session->client_id) - 1);
    }
    if (username) {
        strncpy(session->username, username, sizeof(session->username) - 1);
    }

    // Initialize output buffer
    session->output_buffer = malloc(4096);
    if (!session->output_buffer) {
        free(session);
        return NULL;
    }
    session->output_size = 4096;
    session->output_used = 0;
    session->debug_rdp = debug_rdp;
    pthread_mutex_init(&session->output_mutex, NULL);

    // Create pipes for bidirectional communication
    int stdout_pipe[2]; // parent reads from child stdout
    int stdin_pipe[2];  // parent writes to child stdin

    if (pipe(stdout_pipe) < 0 || pipe(stdin_pipe) < 0) {
        if (stdout_pipe[0] >= 0) close(stdout_pipe[0]);
        if (stdout_pipe[1] >= 0) close(stdout_pipe[1]);
        if (stdin_pipe[0] >= 0) close(stdin_pipe[0]);
        if (stdin_pipe[1] >= 0) close(stdin_pipe[1]);
        free(session->output_buffer);
        free(session);
        return NULL;
    }

    session->read_fd = stdout_pipe[0];  // parent reads from child stdout
    session->write_fd = stdin_pipe[1];  // parent writes to child stdin

    // Build command: wsssht --pipe --wssshd-port 9898 rdp://{client_id}@mbetter.nexlab.net
    char command[1024];
    snprintf(command, sizeof(command),
             "wsssht --pipe --wssshd-port %d rdp://%s@%s",
             config->port, client_id, config->domain);

    // Fork and execute
    pid_t pid = fork();
    if (pid < 0) {
        close(stdout_pipe[0]);
        close(stdout_pipe[1]);
        close(stdin_pipe[0]);
        close(stdin_pipe[1]);
        free(session->output_buffer);
        free(session);
        return NULL;
    }

    if (pid == 0) { // Child process
        // Redirect stdout to stdout_pipe
        dup2(stdout_pipe[1], 1);
        // Redirect stdin from stdin_pipe
        dup2(stdin_pipe[0], 0);
        close(stdout_pipe[0]);
        close(stdout_pipe[1]);
        close(stdin_pipe[0]);
        close(stdin_pipe[1]);

        // Execute command
        execl("/bin/sh", "sh", "-c", command, NULL);
        _exit(1); // Should not reach here
    }

    // Parent process
    close(stdout_pipe[1]); // close write end of stdout pipe
    close(stdin_pipe[0]);  // close read end of stdin pipe
    session->proc_pid = pid;

    // Start output reading thread
    pthread_t thread;
    pthread_create(&thread, NULL, read_pipe_output, session);
    pthread_detach(thread);

    return session;
}

void rdp_free_session(rdp_session_t *session) {
    if (!session) return;

    if (session->read_fd >= 0) {
        close(session->read_fd);
        session->read_fd = -1;
    }
    if (session->write_fd >= 0) {
        close(session->write_fd);
        session->write_fd = -1;
    }

    if (session->output_buffer) {
        free(session->output_buffer);
        session->output_buffer = NULL;
    }
    pthread_mutex_destroy(&session->output_mutex);

    free(session);
}

bool rdp_send_data(rdp_session_t *session, const char *data, size_t len) {
    if (!session || session->write_fd < 0) return false;

    // Check if process is still running
    if (waitpid(session->proc_pid, NULL, WNOHANG) > 0) {
        return false;
    }

    if (session->debug_rdp) {
        printf("RDP: Sending %zu bytes to pipe\n", len);
        // Debug: print first 20 bytes in hex
        printf("RDP send data (first 20 bytes): ");
        for (size_t i = 0; i < len && i < 20; i++) {
            printf("%02x ", (unsigned char)data[i]);
        }
        printf("\n");
    }

    ssize_t written = write(session->write_fd, data, len);
    if (session->debug_rdp) {
        printf("RDP: Wrote %zd bytes to pipe\n", written);
    }
    return written > 0;
}

char *rdp_get_output(rdp_session_t *session, size_t *len) {
    if (!session) {
        *len = 0;
        return NULL;
    }

    pthread_mutex_lock(&session->output_mutex);
    if (session->output_used == 0) {
        // For WebSocket connections, don't wait - return immediately
        pthread_mutex_unlock(&session->output_mutex);
        *len = 0;
        return NULL;
    }

    // Return the buffer and length
    *len = session->output_used;
    char *output = session->output_buffer;
    session->output_buffer = malloc(4096);
    if (session->output_buffer) {
        session->output_size = 4096;
    } else {
        // Out of memory
        *len = 0;
        pthread_mutex_unlock(&session->output_mutex);
        return NULL;
    }
    session->output_used = 0;
    pthread_mutex_unlock(&session->output_mutex);

    return output;
}

bool rdp_disconnect(rdp_session_t *session) {
    if (!session) return false;

    if (session->proc_pid > 0) {
        kill(session->proc_pid, SIGTERM);

        // Wait for graceful termination
        for (int i = 0; i < 30; i++) { // Wait up to 3 seconds
            if (waitpid(session->proc_pid, NULL, WNOHANG) > 0) {
                break;
            }
            usleep(100000); // 100ms
        }

        // Force kill if still running
        if (waitpid(session->proc_pid, NULL, WNOHANG) == 0) {
            kill(session->proc_pid, SIGKILL);
            waitpid(session->proc_pid, NULL, 0);
        }
    }

    return true;
}

bool rdp_is_running(rdp_session_t *session) {
    if (!session || session->proc_pid <= 0) return false;

    return waitpid(session->proc_pid, NULL, WNOHANG) == 0;
}