#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <signal.h>
#include <fcntl.h>
#include <getopt.h>

volatile int running = 1;

void sigint_handler(int sig) {
    running = 0;
}

// Simple base64 encoding function
char *base64_encode(const unsigned char *data, size_t input_length) {
    static const char base64_chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

    size_t output_length = 4 * ((input_length + 2) / 3);
    char *encoded_data = malloc(output_length + 1);
    if (encoded_data == NULL) return NULL;

    for (size_t i = 0, j = 0; i < input_length;) {
        uint32_t octet_a = i < input_length ? data[i++] : 0;
        uint32_t octet_b = i < input_length ? data[i++] : 0;
        uint32_t octet_c = i < input_length ? data[i++] : 0;

        uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;

        encoded_data[j++] = base64_chars[(triple >> 3 * 6) & 0x3F];
        encoded_data[j++] = base64_chars[(triple >> 2 * 6) & 0x3F];
        encoded_data[j++] = base64_chars[(triple >> 1 * 6) & 0x3F];
        encoded_data[j++] = base64_chars[(triple >> 0 * 6) & 0x3F];
    }

    // Add padding
    size_t padding = (3 - (input_length % 3)) % 3;
    for (size_t i = 0; i < padding; i++) {
        encoded_data[output_length - 1 - i] = '=';
    }

    encoded_data[output_length] = '\0';
    return encoded_data;
}

int main(int argc, char *argv[]) {
    char *host = "127.0.0.1";  // Default host
    int port = 8080;           // Default port
    int timeout_seconds = 2;   // Default timeout
    int num_messages = 0;      // Default: infinite loop
    int string_length = 0;     // Default: random length (10-40)

    // Parse command line arguments
    static struct option long_options[] = {
        {"timeout", required_argument, 0, 't'},
        {"n", required_argument, 0, 'n'},
        {"string-length", required_argument, 0, 's'},
        {0, 0, 0, 0}
    };

    int opt;
    int option_index = 0;
    while ((opt = getopt_long(argc, argv, "t:n:s:", long_options, &option_index)) != -1) {
        switch (opt) {
            case 't':
                timeout_seconds = atoi(optarg);
                if (timeout_seconds < 0) {
                    fprintf(stderr, "Invalid timeout value: %s\n", optarg);
                    return 1;
                }
                break;
            case 'n':
                num_messages = atoi(optarg);
                if (num_messages < 0) {
                    fprintf(stderr, "Invalid number of messages: %s\n", optarg);
                    return 1;
                }
                break;
            case 's':
                string_length = atoi(optarg);
                if (string_length < 1 || string_length > 1000) {
                    fprintf(stderr, "Invalid string length: %s (must be 1-1000)\n", optarg);
                    return 1;
                }
                break;
            default:
                fprintf(stderr, "Usage: %s [options] [host] [port]\n", argv[0]);
                fprintf(stderr, "Options:\n");
                fprintf(stderr, "  --timeout, -t <seconds>    Timeout between messages (default: 2)\n");
                fprintf(stderr, "  --n, -n <count>           Number of messages to send (0 for infinite, default: 0)\n");
                fprintf(stderr, "  --string-length, -s <len> String length to generate (1-1000, default: random 10-40)\n");
                fprintf(stderr, "Arguments:\n");
                fprintf(stderr, "  host                      Server host (default: 127.0.0.1)\n");
                fprintf(stderr, "  port                      Server port (default: 8080)\n");
                return 1;
        }
    }

    // Parse positional arguments for host and port
    if (optind < argc) {
        host = argv[optind++];
    }
    if (optind < argc) {
        port = atoi(argv[optind++]);
    }

    if (port <= 0 || port > 65535) {
        fprintf(stderr, "Invalid port: %d\n", port);
        return 1;
    }

    printf("Connecting to %s:%d\n", host, port);
    printf("Timeout: %d seconds\n", timeout_seconds);
    printf("Number of messages: %s\n", num_messages == 0 ? "infinite" : "limited");
    printf("String length: %s\n", string_length > 0 ? "fixed" : "random (10-40)");
    if (string_length > 0) {
        printf("Fixed string length: %d\n", string_length);
    }

    signal(SIGINT, sigint_handler);

    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("socket");
        return 1;
    }

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    if (inet_pton(AF_INET, host, &addr.sin_addr) <= 0) {
        perror("inet_pton");
        close(sock);
        return 1;
    }

    if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("connect");
        close(sock);
        return 1;
    }

    // Make socket non-blocking
    int flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, flags | O_NONBLOCK);

    int message_count = 0;

    while (running) {
        // Check if we've sent the required number of messages
        if (num_messages > 0 && message_count >= num_messages) {
            printf("Sent %d messages as requested. Exiting.\n", num_messages);
            break;
        }

        // Generate string
        srand(time(NULL) + message_count);  // Add message_count to avoid same seed
        int len;
        if (string_length > 0) {
            len = string_length;
        } else {
            len = rand() % 31 + 10; // 10 to 40
        }
        char str[len + 1];
        for (int i = 0; i < len; i++) {
            str[i] = 'A' + rand() % 26;
        }
        str[len] = '\0';

        char *b64_sent = base64_encode((unsigned char *)str, len);
        printf("(SENT) %s <-> %s\n", str, b64_sent ? b64_sent : "ERROR");
        write(sock, str, len);
        message_count++;
        if (b64_sent) free(b64_sent);

        // Wait for response with specified timeout
        fd_set read_set;
        FD_ZERO(&read_set);
        FD_SET(sock, &read_set);
        struct timeval timeout;
        timeout.tv_sec = timeout_seconds;
        timeout.tv_usec = 0;

        int ret = select(sock + 1, &read_set, NULL, NULL, &timeout);
        if (ret > 0 && FD_ISSET(sock, &read_set)) {
            char buffer[1024];
            int n = read(sock, buffer, sizeof(buffer) - 1);
            if (n > 0) {
                buffer[n] = '\0';
                // Remove trailing newline if present
                if (n > 0 && buffer[n-1] == '\n') {
                    buffer[n-1] = '\0';
                    n--;
                }
                char *b64_recv = base64_encode((unsigned char *)buffer, n);
                printf("(RECV) %s <-> %s\n", buffer, b64_recv ? b64_recv : "ERROR");
                if (b64_recv) free(b64_recv);
            } else if (n == 0) {
                // Connection closed
                printf("Connection closed by server\n");
                break;
            } else {
                perror("read");
                break;
            }
        } else if (ret == 0) {
            printf("Timeout waiting for response\n");
        } else {
            perror("select");
            break;
        }

        // Add delay between messages based on timeout parameter
        if (timeout_seconds > 0) {
            sleep(timeout_seconds);
        } else {
            usleep(100000);  // 100ms delay for 0 timeout
        }
    }

    close(sock);
    return 0;
}