#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()

#include <string>
#include <cstring>
#include <cstdlib> // for atoi()
#include <vector>

#include "helper.h"

// GCM encryption
int gcm_encrypt(unsigned char *plaintext, int plaintext_len,
                unsigned char *aad, int aad_len,
                unsigned char *key,
                unsigned char *iv, int iv_len,
                unsigned char *ciphertext,
                unsigned char *tag)
{
    EVP_CIPHER_CTX *ctx;
    int len, ciphertext_len;
    /* Create and initialize the context */
    ctx = EVP_CIPHER_CTX_new();
    if (!ctx)
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }
    /* Initialize the encryption operation. */
    if (!EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, key, iv))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }
    /*
     * Set IV length if default 12 bytes (96 bits) is not appropriate
     */
    if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, iv_len, NULL))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    };
    /* Initialize key and IV */
    if (!EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }

    /*
     * Provide any AAD data. This can be called zero or more times as
     * required
     */
    if (!EVP_EncryptUpdate(ctx, NULL, &len, aad, aad_len))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }

    /*
    * 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))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }
    ciphertext_len = len;
    /*
     * Finalize the encryption. Normally ciphertext bytes may be written at
     * this stage, but this does not occur in GCM mode
     */
    if (!EVP_EncryptFinal_ex(ctx, ciphertext + len, &len))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }
    ciphertext_len += len;
    /* Get the tag */
    if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag))
    {
        ERR_print_errors_fp(stderr);
        EVP_CIPHER_CTX_free(ctx);
        return 0;
    }
    /* 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");
    const char *add = std::getenv("ADD");
    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-GCM\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);
        printf("Message sent to server: %s, length: %d \n", buffer, (int)strlen(buffer));

        unsigned char tag[16] = {0};
        unsigned char cipher_text[2048] = {0};
        int cipher_text_len = gcm_encrypt((unsigned char *)buffer,
                    (int)strlen(buffer),
                    (unsigned char *)add,
                    (int)strlen(add),
                    key_bytes,
                    iv_bytes,
                    iv_len,
                    cipher_text,
                    tag);

        printf("Ciphertext sent to server: %s\n", byte_to_hex(cipher_text, cipher_text_len));
        ssize_t sent = send(client_socket, cipher_text, cipher_text_len, 0);
        if (sent <= 0)
        {
            perror("No message sent to server");
            break;
        }
        sleep(1);
        // send tag of the message
        printf("TAG sent to server: %s\n", byte_to_hex(tag, 16));

        ssize_t sent_tag = send(client_socket, tag, 16, 0);
        if (sent_tag <= 0)
        {
            perror("No TAG 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;
}