#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> // for perror()
#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
    }
}

// Define the decryption function
int stream_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *key, unsigned char *iv, unsigned char *plaintext) {
    /* Declare cipher context */
    EVP_CIPHER_CTX *ctx;

    int len, plaintext_len;

    /* Create and initialise the context */
    ctx = EVP_CIPHER_CTX_new();

    /* Initialise the decryption operation. */
    EVP_DecryptInit_ex(ctx, EVP_aes_256_ctr(), NULL, key, iv);

    /* Provide the message to be decrypted, and obtain the plaintext output. EVP_DecryptUpdate can be called multiple times if necessary */
    EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len);

    /* Finalize the decryption. Further plaintext bytes may be written at this stage. */
    EVP_DecryptFinal_ex(ctx, plaintext + len, &plaintext_len);

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return plaintext_len;
}

int main(void){
    printf("Server starting...\n");
    load_env(".env");

    // Declare variables
    const char *server_ip = std::getenv("SERVER_IP");
    const int server_port = std::atoi(std::getenv("SERVER_PORT"));
    char client_message[2048];
    char server_message[2048];
    const char * custom_message="Server: Hello from server, ";

    // decryption parameters
    const char *inital_vector = std::getenv("INITIAL_VECTOR");
    const char *secret_key = std::getenv("SECRET_KEY");

    // debug
    printf("Server starting at IP: %s, Port: %d\n", server_ip, server_port);

    // Create socket
    const int server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket == -1) {
        perror("Failed to create socket");
        return 1;
    }

    // Bind to the set port and IP
    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 (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
        perror("Failed to bind socket");
        close(server_socket);
        return 1;
    }
    printf("Done with binding with IP: %s, Port: %d\n", server_ip, server_port);

    // Listen for clients:
    const char *client_ip = std::getenv("CLIENT_IP");
    if (listen(server_socket, 1) == -1) {
        perror("Failed to listen on socket");
        close(server_socket);
        return 1;
    }
    printf("Listening for incoming connections...\n");

    // Accept an incoming connection
    struct sockaddr_in client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int client_socket = accept(server_socket, (struct sockaddr*)&client_addr, &client_addr_len);
    if (client_socket == -1) {
        perror("Failed to accept connection");
        close(server_socket);
        return 1;
    }
    printf("Client connected at IP: %s \n", client_ip);

    // clean exising buffer
    memset(client_message, 0, sizeof(client_message));

    // Receive client's message
    while (strcmp(client_message, "\\exit\n") != 0) {
        
        // clean exising buffer
        memset(client_message, 0, sizeof(client_message));

        recv(client_socket, client_message, sizeof(client_message), 0);

        // decrypt client message
        char plaintext[2048];
        int plaintext_len = stream_decrypt((unsigned char*)client_message, strlen(client_message), (unsigned char*)secret_key, (unsigned char*)inital_vector, (unsigned char*)plaintext);

        printf("Msg from client: %.*s", plaintext_len, plaintext);

        // assign client message to plaintext
        memcpy(client_message, plaintext, plaintext_len);

        // Respond to client
        // prepare server message
        memset(server_message, 0, sizeof(server_message));
        ;

        // add my name in the back as response.
        std::snprintf(server_message, sizeof(server_message), "%s%s", custom_message, strtok(client_message, ":"));
        
        if (send(client_socket, server_message, sizeof(server_message), 0) == -1) {
            perror("Failed to send message");
            close(client_socket);
            close(server_socket);
            return 1;
        }
        printf("Response sent to client: %s\n", server_message);

    }
    // Close the socket
    close(client_socket);
    close(server_socket);

    return 0;
}