/*
 * WSSSH Tunnel (wsssht) - Utility Functions Implementation
 * Utility functions for wsssht
 *
 * 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 <getopt.h>

#include "wssshlib.h"
#include "utils.h"

int parse_connection_string(const char *conn_str, char **service, char **client_id, char **wssshd_host, int *wssshd_port) {
    char *str = strdup(conn_str);
    if (!str) return 0;

    char *working_str = str; // Keep track of what to free

    // Check for service prefix (e.g., "ssh://client@host:port")
    char *service_end = strstr(str, "://");
    if (service_end) {
        *service_end = '\0';
        *service = strdup(str);
        working_str = service_end + 3; // Point to part after "://"
    }

    // Find @ separator for client@host
    char *at_pos = strchr(working_str, '@');
    if (!at_pos) {
        // No @ found, treat whole string as client_id
        *client_id = strdup(working_str);
        free(str);
        return 1;
    }

    // Split client_id and host part
    *at_pos = '\0';
    *client_id = strdup(working_str);

    // Parse host and port
    char *host_part = at_pos + 1;
    char *colon_pos = strchr(host_part, ':');

    if (colon_pos) {
        // Has port
        *colon_pos = '\0';
        *wssshd_host = strdup(host_part);
        *wssshd_port = atoi(colon_pos + 1);
    } else {
        // No port
        *wssshd_host = strdup(host_part);
    }

    free(str);
    return 1;
}

void print_usage(const char *program_name) {
    fprintf(stderr, "Usage: %s [options] [service://]clientid[@wssshd-host][:wssshd-port]\n", program_name);
    fprintf(stderr, "WSSSH Tunnel - Setup WebSocket tunnels for manual connections\n\n");
    fprintf(stderr, "Protect the dolls!\n\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  --config FILE        Use custom config file (takes precedence over default)\n");
    fprintf(stderr, "  --clientid ID        Client ID for the tunnel (default: from config)\n");
    fprintf(stderr, "  --tunnel-port PORT   Local tunnel port (default: auto)\n");
    fprintf(stderr, "  --tunnel-host HOST   Local IP address to bind tunnel to (default: 127.0.0.1)\n");
    fprintf(stderr, "  --wssshd-host HOST   wssshd server hostname (required if not in config)\n");
    fprintf(stderr, "  --wssshd-port PORT   wssshd server port (default: 9898)\n");
    fprintf(stderr, "  --interval SEC       Connection retry interval in seconds (default: 5)\n");
    fprintf(stderr, "  --debug              Enable debug output\n");
    fprintf(stderr, "  --tunnel TRANSPORT      Transport for data channel (comma-separated or 'any', or 'websocket'  default: any)\n");
    fprintf(stderr, "  --tunnel-control TYPES Transport types for control channel (comma-separated or 'any', default: any)\n");
    fprintf(stderr, "  --service SERVICE    Service type (default: ssh)\n");
    fprintf(stderr, "  --enc ENCODING       Data encoding: hex, base64, or bin (default: hex)\n");
    fprintf(stderr, "  --mode MODE          Operating mode: interactive, silent, bridge, script, pipe (default: interactive)\n");
    fprintf(stderr, "  --silent             Shortcut for --mode silent\n");
    fprintf(stderr, "  --bridge             Shortcut for --mode bridge\n");
    fprintf(stderr, "  --script             Shortcut for --mode script\n");
    fprintf(stderr, "  --pipe               Shortcut for --mode pipe\n");
    fprintf(stderr, "  --daemon             Enable daemon mode for lazy initialization\n");
    fprintf(stderr, "  --help               Show this help\n");
    fprintf(stderr, "\nExamples:\n");
    fprintf(stderr, "  %s --clientid myclient --wssshd-host mbetter.nexlab.net\n", program_name);
    fprintf(stderr, "  %s myclient@mbetter.nexlab.net:9898\n", program_name);
    fprintf(stderr, "  %s ssh://myclient@mbetter.nexlab.net\n", program_name);
    fprintf(stderr, "  %s --tunnel websocket myclient@mbetter.nexlab.net\n", program_name);
    fprintf(stderr, "  %s --silent myclient@mbetter.nexlab.net\n", program_name);
    fprintf(stderr, "  %s --daemon --clientid myclient --wssshd-host mbetter.nexlab.net\n", program_name);
    fprintf(stderr, "\nDonations:\n");
    fprintf(stderr, "  BTC: bc1q3zlkpu95amtcltsk85y0eacyzzk29v68tgc5hx\n");
    fprintf(stderr, "  ETH: 0xdA6dAb526515b5cb556d20269207D43fcc760E51\n");
}

int wsssht_parse_args(int argc, char *argv[], wsssh_config_t *config, int *remaining_argc, char ***remaining_argv) {
    // Parse wsssht options and optional connection string
    int target_start = -1;
    char *custom_config = NULL;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--config") == 0 && i + 1 < argc) {
            custom_config = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--clientid") == 0 && i + 1 < argc) {
            if (config->client_id) free(config->client_id);
            config->client_id = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--tunnel-port") == 0 && i + 1 < argc) {
            config->local_port = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--tunnel-host") == 0 && i + 1 < argc) {
            if (config->tunnel_host) free(config->tunnel_host);
            config->tunnel_host = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--wssshd-host") == 0 && i + 1 < argc) {
            if (config->wssshd_host) free(config->wssshd_host);
            config->wssshd_host = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--wssshd-port") == 0 && i + 1 < argc) {
            config->wssshd_port = atoi(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--interval") == 0 && i + 1 < argc) {
            config->interval = atoi(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--tunnel") == 0 && i + 1 < argc) {
            if (config->tunnel) free(config->tunnel);
            config->tunnel = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--tunnel-control") == 0 && i + 1 < argc) {
            if (config->tunnel_control) free(config->tunnel_control);
            config->tunnel_control = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--service") == 0 && i + 1 < argc) {
            if (config->service) free(config->service);
            config->service = strdup(argv[i + 1]);
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--enc") == 0 && i + 1 < argc) {
            if (strcmp(argv[i + 1], "hex") == 0) {
                config->encoding = ENCODING_HEX;
            } else if (strcmp(argv[i + 1], "base64") == 0) {
                config->encoding = ENCODING_BASE64;
            } else if (strcmp(argv[i + 1], "bin") == 0) {
                config->encoding = ENCODING_BINARY;
            } else {
                fprintf(stderr, "Error: Invalid encoding: %s (must be hex, base64, or bin)\n", argv[i + 1]);
                return 0;
            }
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--debug") == 0) {
            config->debug = 1;
        } else if (strcmp(argv[i], "--mode") == 0 && i + 1 < argc) {
            if (strcmp(argv[i + 1], "interactive") == 0) {
                config->mode = MODE_INTERACTIVE;
            } else if (strcmp(argv[i + 1], "silent") == 0) {
                config->mode = MODE_SILENT;
            } else if (strcmp(argv[i + 1], "bridge") == 0) {
                config->mode = MODE_BRIDGE;
            } else if (strcmp(argv[i + 1], "script") == 0) {
                config->mode = MODE_SCRIPT;
            } else if (strcmp(argv[i + 1], "pipe") == 0) {
                config->mode = MODE_PIPE;
            } else {
                fprintf(stderr, "Error: Invalid mode: %s\n", argv[i + 1]);
                return 0;
            }
            i++; // Skip the argument
        } else if (strcmp(argv[i], "--daemon") == 0) {
            config->daemon = 1;
        } else if (strcmp(argv[i], "--silent") == 0) {
            config->mode = MODE_SILENT;
        } else if (strcmp(argv[i], "--bridge") == 0) {
            config->mode = MODE_BRIDGE;
        } else if (strcmp(argv[i], "--script") == 0) {
            config->mode = MODE_SCRIPT;
        } else if (strcmp(argv[i], "--pipe") == 0) {
            config->mode = MODE_PIPE;
        } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
            print_usage(argv[0]);
            return 0;
        } else if (argv[i][0] == '-') {
            // Unknown option
            fprintf(stderr, "Error: Unknown option: %s\n", argv[i]);
            print_usage(argv[0]);
            return 0;
        } else {
            // Non-option argument - should be connection string
            if (target_start == -1) {
                target_start = i;
            } else {
                // Multiple non-option arguments not allowed
                fprintf(stderr, "Error: Multiple connection strings not allowed\n");
                print_usage(argv[0]);
                return 0;
            }
        }
    }

    // Handle custom config file if specified
    if (custom_config) {
        // TODO: Implement custom config file loading
        // For now, just validate the file exists
        FILE *test = fopen(custom_config, "r");
        if (!test) {
            fprintf(stderr, "Error: Config file not found: %s\n", custom_config);
            free(custom_config);
            return 0;
        }
        fclose(test);
        fprintf(stderr, "Warning: Custom config file support not yet implemented\n");
        free(custom_config);
    }

    // Handle connection string if present
    if (target_start != -1) {
        char *service = NULL;
        char *client_id = NULL;
        char *wssshd_host = NULL;
        int wssshd_port = 9898;

        if (!parse_connection_string(argv[target_start], &service, &client_id, &wssshd_host, &wssshd_port)) {
            fprintf(stderr, "Error: Invalid connection string format\n");
            free(service);
            free(client_id);
            free(wssshd_host);
            return 0;
        }

        // Set values from connection string (these take precedence over config file)
        if (service) {
            if (config->service) free(config->service);
            config->service = service;
        }

        if (client_id) {
            if (config->client_id) free(config->client_id);
            config->client_id = client_id;
        }

        if (wssshd_host) {
            if (config->wssshd_host) free(config->wssshd_host);
            config->wssshd_host = wssshd_host;
        }

        if (wssshd_port != 9898) {
            config->wssshd_port = wssshd_port;
        }
    }

    // Return remaining arguments (should be none)
    if (target_start != -1) {
        *remaining_argc = argc - target_start - 1;
        *remaining_argv = &argv[target_start + 1];
    } else {
        *remaining_argc = 0;
        *remaining_argv = NULL;
    }

    return 1;
}
