// pq-ftp/server/src/server.cpp
#include <iostream>
#include <string>
#include <vector>
#include <thread>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <fstream>
#include <mutex>
#include <cstring>

// OQS for Post-Quantum Cryptography
#include <oqs/oqs.h>

// Zlib for compression
#include <zlib.h>

#define PORT 2121
#define BUFFER_SIZE 4096
#define PQC_ALGORITHM "Kyber768" // A NIST-selected KEM algorithm

std::mutex log_mutex;

// Simple logger
void log(const std::string& message) {
    std::lock_guard<std::mutex> guard(log_mutex);
    std::ofstream log_file("/logs/server.log", std::ios_base::app);
    time_t now = time(0);
    char* dt = ctime(&now);
    log_file << dt << ": " << message << std::endl;
    std::cout << message << std::endl;
}

// Dummy authentication
bool authenticate(int client_socket) {
    // In a real system, use a secure challenge-response protocol!
    char username[256] = {0};
    char password[256] = {0};
    
    recv(client_socket, username, sizeof(username), 0);
    recv(client_socket, password, sizeof(password), 0);
    
    // For this example, we accept any "user" with password "pass"
    if (std::string(username) == "user" && std::string(password) == "pass") {
        send(client_socket, "OK", 2, 0);
        return true;
    }
    
    send(client_socket, "FAIL", 4, 0);
    return false;
}

// Data compression function
std::vector<uint8_t> compress_data(const std::vector<uint8_t>& data) {
    if (data.empty()) return {};

    uLongf compressed_size = compressBound(data.size());
    std::vector<uint8_t> compressed_data(compressed_size);

    if (compress(compressed_data.data(), &compressed_size, data.data(), data.size()) != Z_OK) {
        log("Compression failed!");
        return {};
    }
    
    compressed_data.resize(compressed_size);
    return compressed_data;
}


void handle_client(int client_socket) {
    log("Client connected. Starting PQC key exchange...");

    // 1. PQC Key Exchange (using liboqs)
    OQS_KEM *kem = NULL;
    uint8_t *public_key = NULL;
    uint8_t *secret_key = NULL;
    uint8_t *ciphertext = NULL;
    uint8_t *shared_secret_server = NULL;
    
    kem = OQS_KEM_new(PQC_ALGORITHM);
    if (kem == NULL) {
        log("Error: OQS_KEM_new failed for " + std::string(PQC_ALGORITHM));
        close(client_socket);
        return;
    }

    public_key = new uint8_t[kem->length_public_key];
    secret_key = new uint8_t[kem->length_secret_key];
    ciphertext = new uint8_t[kem->length_ciphertext];
    shared_secret_server = new uint8_t[kem->length_shared_secret];

    // Generate server's key pair
    if (OQS_KEM_keypair(kem, public_key, secret_key) != OQS_SUCCESS) {
        log("Error: OQS_KEM_keypair failed.");
        // Cleanup would happen here
        close(client_socket);
        return;
    }

    // Send public key to client
    send(client_socket, public_key, kem->length_public_key, 0);

    // Receive ciphertext from client
    recv(client_socket, ciphertext, kem->length_ciphertext, 0);
    
    // Decapsulate to get the shared secret
    if (OQS_KEM_decaps(kem, shared_secret_server, ciphertext, secret_key) != OQS_SUCCESS) {
        log("Error: OQS_KEM_decaps failed.");
        // Cleanup
        close(client_socket);
        return;
    }

    log("PQC key exchange successful. Shared secret established.");
    // NOTE: The shared_secret_server is now our symmetric session key.
    // In a real application, you would derive further keys from this using a KDF.
    // For simplicity, we are not encrypting the file transfer itself, just demonstrating the key exchange.

    // 2. Authentication
    if (!authenticate(client_socket)) {
        log("Authentication failed for client.");
        close(client_socket);
        return;
    }
    log("Client authenticated successfully.");

    // 3. Command Loop (API for scripting)
    char buffer[BUFFER_SIZE] = {0};
    while (recv(client_socket, buffer, BUFFER_SIZE, 0) > 0) {
        std::string command(buffer);
        log("Received command: " + command);
        
        if (command.rfind("put ", 0) == 0) { // Starts with "put "
            std::string filename = "/transfer/" + command.substr(4);
            log("Receiving file: " + filename);
            
            // Get file size
            uint64_t file_size;
            recv(client_socket, &file_size, sizeof(file_size), 0);
            
            std::ofstream outfile(filename, std::ios::binary);
            char file_buffer[BUFFER_SIZE];
            uint64_t bytes_received = 0;
            
            while (bytes_received < file_size) {
                int bytes = recv(client_socket, file_buffer, std::min((uint64_t)BUFFER_SIZE, file_size - bytes_received), 0);
                if (bytes <= 0) break;
                outfile.write(file_buffer, bytes);
                bytes_received += bytes;
            }
            outfile.close();
            log("File received successfully.");
        }
        
        // Clear buffer for next command
        memset(buffer, 0, BUFFER_SIZE);
    }
    
    log("Client disconnected.");
    
    // Cleanup OQS resources
    OQS_KEM_free(kem);
    delete[] public_key;
    delete[] secret_key;
    delete[] ciphertext;
    delete[] shared_secret_server;
    
    close(client_socket);
}

int main() {
    int server_fd;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);

    // Creating socket file descriptor
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0) {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }

    // Forcefully attaching socket to the port
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt))) {
        perror("setsockopt");
        exit(EXIT_FAILURE);
    }
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(PORT);

    // Bind the socket to the network address and port
    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }

    // Start listening for connections
    if (listen(server_fd, 10) < 0) {
        perror("listen");
        exit(EXIT_FAILURE);
    }
    
    log("Server is listening on port " + std::to_string(PORT));
    
    while (true) {
        int client_socket;
        if ((client_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) {
            perror("accept");
            continue; // Continue to next iteration instead of exiting
        }
        
        // Create a new thread to handle the client connection
        std::thread client_thread(handle_client, client_socket);
        client_thread.detach(); // Detach the thread to handle clients concurrently
    }

    return 0;
}