#include "log.h"

#include <dbghelp.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#pragma comment(lib, "DbgHelp.lib")

#include "../../lib/mice/config.h"
#include "../hooks/logging.h"

extern WCHAR exePath[MAX_PATH + 1];
extern DWORD imageOffset;

extern BOOL(WINAPI* TrueWriteFile)(HANDLE hFile, LPCVOID lpBuffer, DWORD nNumberOfBytesToWrite,
                                   LPDWORD lpNumberOfBytesWritten, LPOVERLAPPED lpOverlapped);

BOOL HAS_COLOUR = FALSE;

char _log_prelude[64];
char* log_prelude() {
    time_t rawtime;
    struct tm timeinfo;

    time(&rawtime);
    localtime_s(&timeinfo, &rawtime);

    strftime(_log_prelude, sizeof _log_prelude, "[%Y/%m/%d %H:%M:%S] ", &timeinfo);
    return _log_prelude;
}

HANDLE log_file = NULL;
VOID trace_hook(char* output);
CRITICAL_SECTION logger_lock;

char* log_colours[] = {
    "",          // Always
    "\033[96m",  // Game
    "\033[91m",  // Error
    "\033[33m",  // Warning
    "\033[97m",  // Info
    "\033[90m",  // Misc
    "\033[90m",  // Trace
};
#define LOG_PREFIXES "!GEWIMT"

void force_console_bind() {
    // AttachConsole(ATTACH_PARENT_PROCESS);
    if (GetStdHandle(STD_ERROR_HANDLE) > (HANDLE)0x10) {
        FILE* newstream;
        if (GetStdHandle(STD_INPUT_HANDLE)) {
            freopen_s(&newstream, "CONIN$", "r", stdin);
            setvbuf(stdin, NULL, _IONBF, 0);
        }
        if (GetStdHandle(STD_OUTPUT_HANDLE)) {
            freopen_s(&newstream, "CONOUT$", "w", stdout);
            setvbuf(stdout, NULL, _IONBF, 0);
        }
        if (GetStdHandle(STD_ERROR_HANDLE)) {
            freopen_s(&newstream, "CONOUT$", "w", stderr);
            setvbuf(stderr, NULL, _IONBF, 0);
        }
    }
}

void logcb(LPCSTR param_1) { log_game("amLog", param_1); }
void __stdcall amLogCallback(DWORD level, char* format) {
    if (level == 0)
        log_game("amLog:E", format);
    else if (level == 0)
        log_game("amLog:W", format);
    else
        log_game("amLog:I", format);
}

DWORD pLogcb;
DWORD* ppLogcb;

int _do_log(BYTE log_level, const char* caller, const char* format, va_list args) {
    // TODO: These are all horrible bodges
    if (wcscmp(exePath, L"mxnetwork.exe") == 0) {
        // *((DWORD*)(imageOffset + 0x004438e8)) = (DWORD)(&logcb);
        *((DWORD*)(imageOffset + 0x004438e8)) = 0x00000000;
    }
    if (wcscmp(exePath, L"maimai_dump_.exe") == 0) {
        *((DWORD*)(imageOffset + 0x00c820ec)) = 0x00000001;
        pLogcb = (DWORD)(&amLogCallback);
        ppLogcb = &pLogcb;
        *((DWORD***)(imageOffset + 0x00c820F4)) = &ppLogcb;
        // *((DWORD*)(imageOffset + 0x004438e8)) = (DWORD)(&logcb);
    }

    force_console_bind();

    char prefix = LOG_PREFIXES[log_level];

    EnterCriticalSection(&logger_lock);
    int len = snprintf(NULL, 0, "%s%c:%s:", log_prelude(), prefix, caller) +
              vsnprintf(NULL, 0, format, args);
    char* buf = (char*)malloc(len + 2);
    if (!buf) {
        LeaveCriticalSection(&logger_lock);
        return 0;
    }

    int wrote_a = snprintf(buf, len, "%s%c:%s:", log_prelude(), prefix, caller);
    int wrote_b = vsnprintf(buf + wrote_a, len - wrote_a + 1, format, args);
    buf[len] = '\n';
    buf[len + 1] = '\0';

    // No +1 here to not get the \n
    if (MiceConfig.mice.log_level >= log_level) {
        HANDLE sout = GetStdHandle(STD_OUTPUT_HANDLE);
        if (HAS_COLOUR)
            (TrueWriteFile ? *TrueWriteFile : WriteFile)(
                sout, log_colours[log_level], strlen(log_colours[log_level]), NULL, NULL);
        if (sout != INVALID_HANDLE_VALUE)
            (TrueWriteFile ? TrueWriteFile : WriteFile)(sout, buf, len, NULL, NULL);
        puts(HAS_COLOUR ? "\033[0m" : "");
        FlushFileBuffers(sout);
    }

    if (MiceConfig.mice.log_to_file) {
        if (log_file && log_file != INVALID_HANDLE_VALUE)
            (TrueWriteFile ? TrueWriteFile : WriteFile)(log_file, buf, len + 1, NULL, NULL);
    }

    free(buf);

    LeaveCriticalSection(&logger_lock);
    return wrote_b;
}
int vlog_trace(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_TRACE, caller, format, args);
}
int log_trace(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_trace(caller, format, args);
    va_end(args);
    return ret;
}
int vlog_misc(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_MISC, caller, format, args);
}
int log_misc(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_misc(caller, format, args);
    va_end(args);
    return ret;
}
int vlog_info(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_INFO, caller, format, args);
}
int log_info(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_info(caller, format, args);
    va_end(args);
    return ret;
}
int vlog_warning(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_WARNING, caller, format, args);
}
int log_warning(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_warning(caller, format, args);
    va_end(args);
    return ret;
}
int vlog_error(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_ERROR, caller, format, args);
}
int log_error(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_error(caller, format, args);
    va_end(args);
    return ret;
}
int vlog_game(const char* caller, const char* format, va_list args) {
    return _do_log(LOG_GAME, caller, format, args);
}
int log_game(const char* caller, const char* format, ...) {
    va_list args;
    va_start(args, format);
    int ret = vlog_game(caller, format, args);
    va_end(args);
    return ret;
}

VOID trace_hook(char* output) {
    output[strcspn(output, "\n")] = 0;
    log_error("trace", output);
}

void setup_logging() {
    // Force stdio even for GUI applications
    // TODO: Is there a more robust way to check if we have a proper stdio?
    if (false) {
        force_console_bind();
    }

    // Enable colour in CMD
    HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
    DWORD dwMode = 0;
    if (GetConsoleMode(hConsole, &dwMode))
        HAS_COLOUR = SetConsoleMode(hConsole, dwMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);

    InitializeCriticalSection(&logger_lock);

    if (MiceConfig.mice.log_to_file) {
        if (log_file == NULL && MiceConfig.mice.log_file[0] != '\0')
            log_file = CreateFileA(MiceConfig.mice.log_file, GENERIC_WRITE, FILE_SHARE_READ, NULL,
                                   CREATE_ALWAYS, 0, NULL);
    }
}

void log_stack(const char* caller) {
    char name[MAX_PATH * sizeof(TCHAR)];
    char Storage[sizeof(IMAGEHLP_SYMBOL64) + sizeof(name)];

    IMAGEHLP_SYMBOL64* symbol = (IMAGEHLP_SYMBOL64*)Storage;
    CONTEXT context;
    RtlCaptureContext(&context);

    STACKFRAME64 stack = { 0 };
    stack.AddrPC.Offset = context.Eip;
    stack.AddrPC.Mode = AddrModeFlat;
    stack.AddrStack.Offset = context.Esp;
    stack.AddrStack.Mode = AddrModeFlat;
    stack.AddrFrame.Offset = context.Ebp;
    stack.AddrFrame.Mode = AddrModeFlat;

    HANDLE process = GetCurrentProcess();
    HANDLE thread = GetCurrentThread();
    BOOL initres = SymInitialize(process, NULL, true);
    for (ULONG frame = 0;; frame++) {
        BOOL result = StackWalk64(IMAGE_FILE_MACHINE_I386, process, thread, &stack, &context, NULL,
                                  SymFunctionTableAccess64, SymGetModuleBase64, NULL);

        symbol->SizeOfStruct = sizeof(Storage);
        symbol->MaxNameLength = sizeof(name);

        DWORD64 displacement;
        SymGetSymFromAddr64(process, (ULONG64)stack.AddrPC.Offset, &displacement, symbol);
        UnDecorateSymbolName(symbol->Name, (PSTR)name, sizeof(name), UNDNAME_COMPLETE);

        log_error(caller, "%02u called from 0x%08X STACK=0x%08X FRAME=0x%08X  %s\n", frame,
                  (ULONG64)stack.AddrPC.Offset, (ULONG64)stack.AddrStack.Offset,
                  (ULONG64)stack.AddrFrame.Offset, symbol->Name);

        if (result == FALSE) {
            DWORD frameError = GetLastError();
            break;
        }
    }
}
