/**
 * Web Proxy server 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 <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <errno.h>
#include <time.h>
#include <ctype.h>
#include <signal.h>
#include "web_proxy.h"
#include "websocket.h"
#include "websocket_protocol.h"

// Global state
static wssshd_state_t *global_state = NULL;
static const wssshd_config_t *global_config = NULL;
static int server_socket = -1;
static volatile int server_running = 0;

// Active proxy connections
#define MAX_PROXY_CONNECTIONS 100
typedef struct {
    int client_fd;           // Connection from web proxy client
    void *tunnel_ws;         // WebSocket tunnel to the client
    char client_id[256];
    char hostname[256];
    bool active;
    pthread_t thread;
} proxy_connection_t;

static proxy_connection_t proxy_connections[MAX_PROXY_CONNECTIONS];
static int proxy_connections_count = 0;
static pthread_mutex_t proxy_mutex = PTHREAD_MUTEX_INITIALIZER;

// Forward declaration for proxy connection handler
static void *proxy_connection_handler(void *arg);

// Find client by hostname (matching subdomain or client_id)
static client_t *find_client_by_hostname(const char *hostname) {
    if (!hostname || !global_state) return NULL;

    pthread_mutex_lock(&global_state->client_mutex);

    for (size_t i = 0; i < global_state->clients_count; i++) {
        client_t *client = &global_state->clients[i];
        if (!client->active) continue;

        // Check if client has web service
        if (strstr(client->services, "web") == NULL &&
            strstr(client->services, "http") == NULL) {
            continue;
        }

        // Try exact match with client_id
        if (strcasecmp(hostname, client->client_id) == 0) {
            pthread_mutex_unlock(&global_state->client_mutex);
            return client;
        }

        // Try subdomain match: hostname.client_id.domain
        if (global_config && global_config->domain) {
            char expected_hostname[512];
            snprintf(expected_hostname, sizeof(expected_hostname), "%s.%s",
                     client->client_id, global_config->domain);
            if (strcasecmp(hostname, expected_hostname) == 0) {
                pthread_mutex_unlock(&global_state->client_mutex);
                return client;
            }
        }
    }

    pthread_mutex_unlock(&global_state->client_mutex);
    return NULL;
}

// Parse Host header from HTTP request
static int parse_host_header(const char *request, char *hostname, size_t hostname_size) {
    if (!request || !hostname) return -1;

    const char *host_header = strstr(request, "Host:");
    if (!host_header) return -1;

    host_header += 5; // Skip "Host:"
    while (*host_header == ' ' || *host_header == '\t') host_header++;

    // Extract hostname (up to port or end of line)
    const char *end = host_header;
    while (*end && *end != '\r' && *end != '\n' && *end != ' ') {
        end++;
    }

    size_t len = end - host_header;
    if (len >= hostname_size) len = hostname_size - 1;

    strncpy(hostname, host_header, len);
    hostname[len] = '\0';

    // Remove port if present
    char *colon = strchr(hostname, ':');
    if (colon) {
        *colon = '\0';
    }

    return 0;
}

// Send HTTP error response
static void send_http_error(int client_fd, int status_code, const char *status_text) {
    const char *body_format = "<html><body><h1>%d %s</h1></body></html>";
    char body[512];
    snprintf(body, sizeof(body), body_format, status_code, status_text);

    char response[1024];
    int len = snprintf(response, sizeof(response),
        "HTTP/1.1 %d %s\r\n"
        "Content-Type: text/html\r\n"
        "Content-Length: %zu\r\n"
        "Connection: close\r\n"
        "\r\n"
        "%s",
        status_code, status_text, strlen(body), body);

    send(client_fd, response, len, 0);
}

// Proxy data between client socket and tunnel WebSocket
static void *proxy_data_forward(void *arg) {
    proxy_connection_t *conn = (proxy_connection_t *)arg;
    if (!conn) return NULL;

    char buffer[8192];
    fd_set read_fds;
    int max_fd = conn->client_fd;
    struct timeval tv;

    printf("[WEB-PROXY] Started data forwarding for %s\n", conn->hostname);

    while (conn->active && server_running) {
        FD_ZERO(&read_fds);
        FD_SET(conn->client_fd, &read_fds);

        tv.tv_sec = 1;
        tv.tv_usec = 0;

        int activity = select(max_fd + 1, &read_fds, NULL, NULL, &tv);
        if (activity < 0) {
            if (errno == EINTR) continue;
            break;
        }

        if (activity == 0) {
            continue; // Timeout, check again
        }

        // Read from client socket
        if (FD_ISSET(conn->client_fd, &read_fds)) {
            ssize_t bytes_read = recv(conn->client_fd, buffer, sizeof(buffer), 0);
            if (bytes_read <= 0) {
                // Connection closed or error
                break;
            }

            // Forward to tunnel WebSocket
            if (conn->tunnel_ws) {
                ws_send_binary_frame((ws_connection_t *)conn->tunnel_ws, buffer, bytes_read, false);
            }
        }
    }

    printf("[WEB-PROXY] Stopped data forwarding for %s\n", conn->hostname);
    return NULL;
}

// Handle incoming HTTP request and establish tunnel
static int handle_proxy_request(int client_fd) {
    char request[8192];
    ssize_t bytes_read = recv(client_fd, request, sizeof(request) - 1, 0);
    if (bytes_read <= 0) {
        return -1;
    }
    request[bytes_read] = '\0';

    // Parse Host header
    char hostname[256];
    if (parse_host_header(request, hostname, sizeof(hostname)) != 0) {
        send_http_error(client_fd, 400, "Bad Request");
        return -1;
    }

    printf("[WEB-PROXY] Received request for hostname: %s\n", hostname);

    // Find client by hostname
    client_t *client = find_client_by_hostname(hostname);
    if (!client) {
        printf("[WEB-PROXY] No client found for hostname: %s\n", hostname);
        send_http_error(client_fd, 404, "Not Found");
        return -1;
    }

    printf("[WEB-PROXY] Found client: %s\n", client->client_id);

    // TODO: Establish tunnel to the client's web service
    // For now, we'll return a simple response indicating the proxy is working
    const char *response =
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/plain\r\n"
        "Content-Length: 50\r\n"
        "Connection: close\r\n"
        "\r\n"
        "Web proxy tunnel established for client: ";

    char response_body[512];
    snprintf(response_body, sizeof(response_body), "%s%s", response, client->client_id);

    send(client_fd, response_body, strlen(response_body), 0);

    // Close the connection for now (full tunnel implementation requires more work)
    return 0;
}

// Proxy connection handler thread
static void *proxy_connection_handler(void *arg) {
    int client_fd = *(int *)arg;
    free(arg);

    printf("[WEB-PROXY] New connection received\n");

    handle_proxy_request(client_fd);

    close(client_fd);
    printf("[WEB-PROXY] Connection closed\n");

    return NULL;
}

// Main HTTP server thread
static void *http_proxy_thread(void *arg __attribute__((unused))) {
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);

    server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket < 0) {
        perror("[WEB-PROXY] Failed to create server socket");
        return NULL;
    }

    // Set socket options
    int opt = 1;
    setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr("127.0.0.1");  // Listen only on localhost
    server_addr.sin_port = htons(global_config->web_proxy_port);

    if (bind(server_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("[WEB-PROXY] Failed to bind server socket");
        close(server_socket);
        return NULL;
    }

    if (listen(server_socket, 10) < 0) {
        perror("[WEB-PROXY] Failed to listen on server socket");
        close(server_socket);
        return NULL;
    }

    printf("[WEB-PROXY] Web proxy server starting on 127.0.0.1:%d\n", global_config->web_proxy_port);
    server_running = 1;

    // Ignore SIGPIPE to prevent crashes on broken connections
    signal(SIGPIPE, SIG_IGN);

    while (server_running) {
        int client_fd = accept(server_socket, (struct sockaddr *)&client_addr, &client_len);
        if (client_fd < 0) {
            if (server_running) perror("[WEB-PROXY] Failed to accept client connection");
            continue;
        }

        // Create a new thread to handle the connection
        int *client_fd_ptr = malloc(sizeof(int));
        if (!client_fd_ptr) {
            close(client_fd);
            continue;
        }
        *client_fd_ptr = client_fd;

        pthread_t thread;
        if (pthread_create(&thread, NULL, proxy_connection_handler, client_fd_ptr) != 0) {
            perror("[WEB-PROXY] Failed to create connection handler thread");
            free(client_fd_ptr);
            close(client_fd);
            continue;
        }

        pthread_detach(thread);
    }

    close(server_socket);
    server_socket = -1;
    printf("[WEB-PROXY] Web proxy server stopped\n");
    return NULL;
}

int web_proxy_start_server(const wssshd_config_t *config, wssshd_state_t *state) {
    if (!config->web_proxy_enabled || config->web_proxy_port == 0) {
        return 0; // Web proxy not enabled
    }

    global_config = config;
    global_state = state;

    // Set default port if not specified
    if (config->web_proxy_port == 0) {
        // This won't work directly since config is const, but the CLI parsing sets it
    }

    // Start HTTP proxy server thread
    pthread_t thread;
    if (pthread_create(&thread, NULL, http_proxy_thread, NULL) != 0) {
        perror("[WEB-PROXY] Failed to create HTTP proxy server thread");
        return -1;
    }

    pthread_detach(thread);
    return 0;
}

void web_proxy_stop_server(void) {
    server_running = 0;

    // Close server socket to unblock accept()
    if (server_socket >= 0) {
        close(server_socket);
        server_socket = -1;
    }

    // Close all active proxy connections
    pthread_mutex_lock(&proxy_mutex);
    for (int i = 0; i < proxy_connections_count; i++) {
        if (proxy_connections[i].active) {
            proxy_connections[i].active = false;
            if (proxy_connections[i].client_fd >= 0) {
                close(proxy_connections[i].client_fd);
            }
        }
    }
    proxy_connections_count = 0;
    pthread_mutex_unlock(&proxy_mutex);

    printf("[WEB-PROXY] Web proxy server stopping\n");
}
