/**
 * Terminal and PTY 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 <pty.h>
#include <termios.h>
#include <sys/ioctl.h>
#include <sys/select.h>
#include <sys/wait.h>
#include <signal.h>
#include <uuid/uuid.h>
#include <pthread.h>
#include "terminal.h"



// Thread function to read PTY output
static void *read_pty_output(void *arg) {
    terminal_session_t *session = (terminal_session_t *)arg;

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

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

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

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

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

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

terminal_session_t *terminal_create_session(const wssshd_config_t *config, const char *username, const char *client_id) {
    terminal_session_t *session = calloc(1, sizeof(terminal_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);
    }

    // Open PTY
    int master_fd, slave_fd;
    char slave_name[256];

    if (openpty(&master_fd, &slave_fd, slave_name, NULL, NULL) < 0) {
        free(session);
        return NULL;
    }

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

    session->master_fd = master_fd;

    // Set terminal size (default 80x24)
    struct winsize ws = {24, 80, 0, 0};
    ioctl(master_fd, TIOCSWINSZ, &ws);

    // Build command
    snprintf(session->command, sizeof(session->command),
               "sh -c 'stty echo && wsssh -p %d %s@%s.%s'",
               config->port, username, client_id, config->domain);

    // Create pipe to communicate grandchild PID
    int pipefd[2];
    if (pipe(pipefd) < 0) {
        close(master_fd);
        close(slave_fd);
        free(session);
        return NULL;
    }

    // Fork and execute
    pid_t pid = fork();
    if (pid < 0) {
        close(master_fd);
        close(slave_fd);
        free(session);
        return NULL;
    }

    if (pid == 0) { // Child process
        // Create new session
        setsid();

        // Set up controlling terminal
        if (ioctl(slave_fd, TIOCSCTTY, 0) < 0) {
            // Some systems don't support TIOCSCTTY, try alternative
            // Redirect stdio to PTY anyway
        }

        // Set raw mode on the slave PTY
        struct termios term;
        tcgetattr(slave_fd, &term);
        cfmakeraw(&term);
        tcsetattr(slave_fd, TCSANOW, &term);

        // Redirect stdin/stdout/stderr to slave PTY
        dup2(slave_fd, 0);
        dup2(slave_fd, 1);
        dup2(slave_fd, 2);

        close(master_fd);
        close(slave_fd);

        // Set environment
        setenv("TERM", "xterm-256color", 1);
        setenv("COLUMNS", "80", 1);
        setenv("LINES", "24", 1);

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

    // Parent process
    close(slave_fd);
    session->proc_pid = pid;

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

    return session;
}

void terminal_free_session(terminal_session_t *session) {
    if (!session) return;

    if (session->master_fd >= 0) {
        close(session->master_fd);
    }

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

    free(session);
}

bool terminal_send_data(terminal_session_t *session, const char *data) {
    if (!session || !data || session->master_fd < 0) return false;

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

    ssize_t written = write(session->master_fd, data, strlen(data));
    return written > 0;
}

char *terminal_get_output(terminal_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 terminal_disconnect(terminal_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 terminal_resize(terminal_session_t *session, int cols, int rows) {
    if (!session || session->master_fd < 0) return false;

    struct winsize ws = {rows, cols, 0, 0};
    if (ioctl(session->master_fd, TIOCSWINSZ, &ws) < 0) {
        return false;
    }

    // Send SIGWINCH to the process
    if (session->proc_pid > 0) {
        kill(session->proc_pid, SIGWINCH);
    }

    return true;
}

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

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