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

#include "transport.h"

#define SYSINFO_METHOD_REGISTRY TRUE // TRUE = use registry method, FALSE = use systeminfo cmd

void CollectSystemInfo(char *osVersion, char *architecture, char *hostname, char *username, char *localIP, char* procname, unsigned long* pid) {

#if SYSINFO_METHOD_REGISTRY == TRUE
    // OS Version - registry method
    HKEY hKey;
    DWORD bufferSize = 256;
    char temp_OS_name[256] = "Unknown OS Name";
    char temp_OS_version[256] = "Unknown OS Version";
    char temp_current_build[256] = "Unknown Current Build";

    // Open the registry key
    if (RegOpenKeyExA(HKEY_LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion", 0, KEY_READ, &hKey) == ERROR_SUCCESS) {

        // Get ProductName (OS Name)
        bufferSize = sizeof(temp_OS_name);
        if (RegQueryValueExA(hKey, "ProductName", NULL, NULL, (LPBYTE)temp_OS_name, &bufferSize) != ERROR_SUCCESS) {
            strcpy(temp_OS_name, "Unknown OS Name");
        }
        
        // Get DisplayVersion (OS Version)
        bufferSize = sizeof(temp_OS_version);
        if (RegQueryValueExA(hKey, "DisplayVersion", NULL, NULL, (LPBYTE)temp_OS_version, &bufferSize) != ERROR_SUCCESS) {
            strcpy(temp_OS_version, "Unknown OS Version");
        }

        // Get CurrentBuild (sometimes different from BuildNumber)
        bufferSize = sizeof(temp_current_build);
        if (RegQueryValueExA(hKey, "CurrentBuild", NULL, NULL, (LPBYTE)temp_current_build, &bufferSize) != ERROR_SUCCESS) {
            strcpy(temp_current_build, "Unknown Current Build");
        }

        RegCloseKey(hKey);

        // Format the complete version string
        snprintf(osVersion, 256, "%s (Version %s) Build %s", temp_OS_name, temp_OS_version, temp_current_build);
    } 
    else {
        strcpy(osVersion, "Unknown OS");
    }
#else
    // OS Version - CMD method
    FILE *fp;
    char buffer[512];
    char command[] = "systeminfo | findstr /B /C:\"OS Name\" /C:\"OS Version\"";
    char temp_OS_name[128] = "Unknown OS Name";
    char temp_OS_version[128] = "Unknown OS Version";

    fp = _popen(command, "r");
    if (fp == NULL) {
        LOG_ERROR("Failed to run command.\n");
        strcpy(osVersion, "Unknown OS");
    } else {
        while (fgets(buffer, sizeof(buffer), fp) != NULL) {
            if (strstr(buffer, "OS Name:")) {
                sscanf(buffer, "OS Name: %[^\n]", temp_OS_name);
            } else if (strstr(buffer, "OS Version:")) {
                sscanf(buffer, "OS Version: %[^\n]", temp_OS_version);
            }
        }
        _pclose(fp);
        snprintf(osVersion, 256, "%s %s", temp_OS_name, temp_OS_version);
        LOG("temp_OS_name: %s\n", temp_OS_name);
        LOG("temp_OS_version: %s\n", temp_OS_version);
    }
#endif
    // Get PID and name
    strncpy(procname, _pgmptr, 128);
    procname[127] = '\0'; // Ensure null termination

    *pid = GetCurrentProcessId();
    
    // LOG("Proc name: %s\n", procname);
    // LOG("PID: %lu\n", *pid);

    // Architecture
    SYSTEM_INFO sysInfo;
    GetSystemInfo(&sysInfo);
    switch(sysInfo.wProcessorArchitecture) {
        case PROCESSOR_ARCHITECTURE_AMD64: 
            strcpy(architecture, "x64"); 
            break;
        case PROCESSOR_ARCHITECTURE_INTEL: 
            strcpy(architecture, "x86"); 
            break;
        case PROCESSOR_ARCHITECTURE_ARM64: 
            strcpy(architecture, "ARM64"); 
            break;
        default: 
            strcpy(architecture, "Unknown");
    }

    // Hostname
    DWORD hostname_len = 128;
    GetComputerNameA(hostname, &hostname_len);

    // Username
    DWORD username_len = 128;
    GetUserNameA(username, &username_len);

    // Local IP (first non-loopback)
    ULONG outBufLen = 0;
    GetAdaptersInfo(NULL, &outBufLen);
    PIP_ADAPTER_INFO pAdapterInfo = malloc(outBufLen);
    GetAdaptersInfo(pAdapterInfo, &outBufLen);
    
    strcpy(localIP, "Unknown");
    for (PIP_ADAPTER_INFO pAdapter = pAdapterInfo; pAdapter; pAdapter = pAdapter->Next) {
        if (strcmp(pAdapter->IpAddressList.IpAddress.String, "127.0.0.1") != 0) {
            strcpy(localIP, pAdapter->IpAddressList.IpAddress.String);
            break;
        }
    }
    free(pAdapterInfo);
}