#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>

// open ssl
#include <openssl/conf.h>
#include <openssl/evp.h>
#include <openssl/err.h>

#include <errno.h> // error handling
#include <unistd.h> // for close()

// load environment variables from .env file
#include <fstream>
#include <cstdlib>

void load_env(const std::string& path) {
    std::ifstream f(path);
    std::string line;
    while (std::getline(f, line)) {
        if (line.empty() || line[0] == '#') continue;
        auto pos = line.find('=');
        if (pos == std::string::npos) continue;

        std::string key = line.substr(0, pos);
        std::string val = line.substr(pos + 1);

        #ifdef _WIN32
                _putenv_s(key.c_str(), val.c_str());
        #else
                setenv(key.c_str(), val.c_str(), 1);
        #endif
    }
}

// hex to byte helper function
int hex_to_bytes_upper(const char* hex, unsigned char* out, size_t out_size)
{
    size_t i = 0;
    while (hex[0] && hex[1]) {
        if (i >= out_size) return -1;
        unsigned char h = hex[0];
        unsigned char l = hex[1];
        int hi = (h <= '9') ? (h - '0') : (h - 'A' + 10);
        int lo = (l <= '9') ? (l - '0') : (l - 'A' + 10);
        // minimal sanity check
        if (hi < 0 || hi > 15 || lo < 0 || lo > 15)
            return -1;
        out[i++] = (unsigned char)((hi << 4) | lo);
        hex += 2;
    }
    return (int)i;
}

// block cipher encryption function
int block_encrypt(unsigned char *plaintext, int plaintext_len, unsigned char *key, unsigned char *iv, unsigned char *ciphertext) {
    /* Declare cipher context */
    EVP_CIPHER_CTX *ctx;

    int len, ciphertext_len;

    /* Create and initialise the context */
    ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        ERR_print_errors_fp(stderr);
    }

    /* Initialise the encryption operation. */
    // choice for aes_256 ref: https://stackoverflow.com/questions/1220751/how-to-choose-an-aes-encryption-mode-cbc-ecb-ctr-ocb-cfb
    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv) != 1){
        ERR_print_errors_fp(stderr);
    }

    /* Provide the message to be encrypted, and obtain the encrypted output. EVP_EncryptUpdate can be called multiple times if necessary */
    if (EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len) != 1) {
        ERR_print_errors_fp(stderr);
    }

    /* Finalize the encryption. Further cipher text bytes may be written at this stage. */
    if (EVP_EncryptFinal_ex(ctx, ciphertext + len, &ciphertext_len) != 1) {
        ERR_print_errors_fp(stderr);
    }

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return ciphertext_len;
}

int main(void){
    load_env(".env");
    // Declare variables
    const char *server_ip = std::getenv("SERVER_IP");
    const int server_port = std::atoi(std::getenv("SERVER_PORT"));
    printf("Connecting to server %s:%d\n", server_ip, server_port);
    
    char client_message[1024];
    char server_message[1024];
    const char * custom_message="Zheyuan Wu: ";

    // encryption parameters
    const char *inital_vector = std::getenv("INITIAL_VECTOR");
    const char *secret_key = std::getenv("SECRET_KEY");
    unsigned char key_bytes[32], iv_bytes[16];
    
    int key_len = hex_to_bytes_upper(secret_key, key_bytes, sizeof(key_bytes));
    int iv_len  = hex_to_bytes_upper(inital_vector, iv_bytes, sizeof(iv_bytes));

    if (key_len != 32 || iv_len != 16) {
        fprintf(stderr, "Invalid key/IV size for AES-256\n");
        return 1;
    }

    // Create socket:
    int client_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (client_socket == -1) {
        perror("Failed to create socket");
        return 1;
    }else{
        printf("Socket created successfully\n");
    }

    // Send connection request to server, be sure to set por tand IP the same as server-side
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(server_port);
    inet_pton(AF_INET, server_ip, &server_addr.sin_addr);

    if (connect(client_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
        perror("Failed to connect to server");
        close(client_socket);
        return 1;
    }else{
        printf("Connected to server successfully\n");
    }

    // Get input from the user:
    printf("Enter message sent to the server (type \\quit to exit): ");
    // clean the buffer
    memset(client_message, 0, sizeof(client_message));
    if (fgets(client_message, sizeof(client_message), stdin) == NULL) {
        // EOF or error reading from stdin, exit the loop
        perror("Error reading from stdin");
        return 1;
    }
    
    while (strcmp(client_message, "\\quit\n") != 0) {
        
        // Send the message to server:
        // add my name in the front
        char buffer[2048];
        std::snprintf(buffer, sizeof(buffer), "%s%s", custom_message, client_message);

        unsigned char ciphertext[2048];

        int plaintext_len = (int)strlen(buffer);
        int ciphertext_len = block_encrypt((unsigned char*)buffer, plaintext_len, key_bytes, iv_bytes, ciphertext);
        if (ciphertext_len <= 0) {
            printf("encrypt failed or produced empty ciphertext_len=%d\n", ciphertext_len);
            break;
        }

        ssize_t sent = send(client_socket, ciphertext, (size_t)ciphertext_len, 0);
        if (sent <= 0) { perror("No message sent to server"); break; }

        // Receive the server's response:
        // add terminator for string
        ssize_t recvd = recv(client_socket, server_message, sizeof(server_message) - 1, 0);
        if (recvd <= 0) { perror("No message received from server"); break; }
        server_message[recvd] = '\0';

        printf("Server's response: %s\n", server_message);

        printf("Enter message sent to the server (type \\quit to exit): ");

        // clean the buffer
        memset(client_message, 0, sizeof(client_message));
        memset(server_message, 0, sizeof(server_message));
        
        if (fgets(client_message, sizeof(client_message), stdin) == NULL) {
            // EOF or error reading from stdin, exit the loop
            perror("Error reading from stdin");
            break;
        }
    }

    // Close the socket
    close(client_socket);

    return 0;
}