#include "win-platform.h"
#include "utils.h"
#include "spoofer.h"
#include "log-sink.h"
#include "proxy-configurator.h"
#include "cache-cleaner.h"

#include <seallib/log.h>
#include <tinymitm/proxy.h>

#include <memory>

#include <processenv.h>
#include <cassert>
#include <processthreadsapi.h>
#include <handleapi.h>

#ifndef CREATE_NO_WINDOW
    // from: WinBase.h
    #define CREATE_NO_WINDOW 0x08000000
#endif

bool running = true;
TinyMITMProxy* proxy = nullptr;
ProxyConfigurator* conf = nullptr;
CacheCleaner* cleaner = nullptr;

BOOL WINAPI consoleHandler(DWORD dwType)
{
    if (dwType == CTRL_C_EVENT || dwType == CTRL_CLOSE_EVENT || dwType == CTRL_LOGOFF_EVENT ||
        dwType == CTRL_SHUTDOWN_EVENT)
    {
        if (conf) conf->clearProxy();
        running = false;
        return TRUE;
    }
    return FALSE;
}

void StartWatchdog()
{
    DWORD pid = GetCurrentProcessId();

    /*
        this probably looks so sketchy for some antiviruses
        should test, probably
    */
    std::string psCommand =
        std::format("powershell.exe -WindowStyle Hidden -Command \""
                    "$proc = Get-Process -Id {} -ErrorAction SilentlyContinue;"
                    "if ($proc) {{ $proc.WaitForExit(); }}"
                    "$reg = 'HKCU:\\Software\\Microsoft\\Windows\\CurrentVersion\\Internet Settings';"
                    "Set-ItemProperty -Path $reg -Name ProxyEnable -Value 0;"
                    "Set-ItemProperty -Path $reg -Name ProxyServer -Value '';"
                    "rundll32.exe user32.dll,UpdatePerUserSystemParameters;\"",
                    pid);

    STARTUPINFOA si;
    PROCESS_INFORMATION pi;
    RtlZeroMemory(&si, sizeof(si));
    RtlZeroMemory(&pi, sizeof(pi));
    si.cb = sizeof(si);

    CreateProcessA(NULL, const_cast<char*>(psCommand.c_str()), NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &si,
                   &pi);

    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
}

bool run()
{
    seallib::Logger mainLog("Main");
    mainLog.addSink(std::make_shared<ConOutSink>());

    mainLog.info("Init");

    mainLog.info("Setting up cache cleaner");
    cleaner = new CacheCleaner();
    cleaner->init();

    mainLog.info("Creating proxy");
    TinyMITMConfig mitmConfig;
    mitmConfig.port = 44444;
    mitmConfig.threadCount = 255;
    mitmConfig.caName = utils::randomizeString(32);
    proxy = new TinyMITMProxy(mitmConfig);

    mainLog.info("Creating configurator and setting proxy addr");
    conf = new ProxyConfigurator();
    conf->setProxy("127.0.0.1", 44444);

    mainLog.info("Instantiating spoofer and registering proxy listeners");
    Spoofer* spoofer = new Spoofer();
    spoofer->registerListeners(proxy);

    mainLog.info("Starting proxy");
    if (!proxy->init())
    {
        mainLog.error("Failed to initialize proxy, shutting down");
        proxy->shutdown();
        cleaner->shutdown();
        return false;
    }

    mainLog.info("Proxy running, Ctrl+C to stop.");
    while (running)
        std::this_thread::sleep_for(std::chrono::milliseconds(100));

    mainLog.info("Shutting down...");
    proxy->shutdown();
    cleaner->shutdown();

    delete proxy;
    proxy = nullptr;
    delete conf;
    conf = nullptr;

    return true;
}

int main()
{

    /*
        handlers for cleaning proxy conf on exit / crash / taskkill / whatever
    */
    SetConsoleCtrlHandler(consoleHandler, TRUE);
    StartWatchdog();

    // __try
    // {
    if (!run()) return 1;
    // }
    // __except (EXCEPTION_EXECUTE_HANDLER)
    // {
    //     if (conf) conf->clearProxy();
    // }

    return 0;
}