#include <windows.h>
#include <winnls.h>
#include <shobjidl.h>
#include <objbase.h>
#include <Shlwapi.h>
#include <objidl.h>
#include <shlguid.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <thread>
#include "hid_handler.h"
#include "process_watch.h"
#include "utils.h"
#include "ui.h"
#include "localization.h"

// Is this process running with elevated rights?
BOOL is_elevated()
{
    BOOL fRet = FALSE;
    HANDLE hToken = NULL;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
    {
        TOKEN_ELEVATION Elevation;
        DWORD cbSize = sizeof(TOKEN_ELEVATION);
        if (GetTokenInformation(hToken, TokenElevation, &Elevation, sizeof(Elevation), &cbSize))
        {
            fRet = Elevation.TokenIsElevated;
        }
    }
    if (hToken)
    {
        CloseHandle(hToken);
    }
    return fRet;
}

/// Create a shortcut.
static HRESULT create_link(LPCWSTR working_dir,
    LPCWSTR exe_path,
    LPCWSTR shorcut_path,
    LPCWSTR description)
{
    HRESULT hres;
    IShellLink* psl;
    hres = CoCreateInstance(CLSID_ShellLink, NULL, CLSCTX_INPROC_SERVER, IID_IShellLink, (LPVOID*)&psl);
    if (SUCCEEDED(hres))
    {
        IPersistFile* ppf;
        psl->SetPath(exe_path);
        psl->SetDescription(description);
        psl->SetWorkingDirectory(working_dir);
        psl->SetArguments(L"-startupLaunch");
        hres = psl->QueryInterface(IID_IPersistFile, (LPVOID*)&ppf);
        if (SUCCEEDED(hres))
        {
            hres = ppf->Save(shorcut_path, TRUE);
            ppf->Release();
        }
        psl->Release();
    }
    return hres;
}

bool UriSchemeExists(const std::wstring& scheme)
{
    HKEY hKey;
    std::wstring keyPath = L"Software\\Classes\\" + scheme;

    LONG result = RegOpenKeyExW(HKEY_CURRENT_USER, keyPath.c_str(), 0, KEY_READ, &hKey);
    if (result == ERROR_SUCCESS)
    {
        RegCloseKey(hKey);
        return true;
    }
    return false;
}

void RegisterUri(const std::wstring& scheme, const std::wstring& exePath)
{
    if (UriSchemeExists(scheme))
    {
        return;
    }

    HKEY hKey;
    std::wstring baseKeyPath = L"Software\\Classes\\" + scheme;

    if (RegCreateKeyExW(HKEY_CURRENT_USER, baseKeyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, NULL) != ERROR_SUCCESS)
    {
        return;
    }

    std::wstring description = L"URL:" + scheme + L" Protocol";
    RegSetValueExW(hKey, NULL, 0, REG_SZ, (BYTE*)description.c_str(), static_cast<DWORD>((description.size() + 1) * sizeof(wchar_t)));
    RegSetValueExW(hKey, L"URL Protocol", 0, REG_SZ, (BYTE*)L"", sizeof(wchar_t));

    HKEY hSubKey;
    std::wstring commandKeyPath = baseKeyPath + L"\\shell\\open\\command";
    if (RegCreateKeyExW(HKEY_CURRENT_USER, commandKeyPath.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hSubKey, NULL) != ERROR_SUCCESS)
    {
        RegCloseKey(hKey);
        return;
    }

    std::wstring command = L"\"" + exePath + L"\" \"%1\"";
    RegSetValueExW(hSubKey, NULL, 0, REG_SZ, (BYTE*)command.c_str(), static_cast<DWORD>((command.size() + 1) * sizeof(wchar_t)));

    RegCloseKey(hSubKey);
    RegCloseKey(hKey);

    return;
}

///
int main(int argc, char* argv[])
{
    bool startClosed = false;
    bool restartInstance = false;
    bool ignoreMutex = false;
    bool hasResetCamera = false;
    
    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "-startupLaunch") == 0)
        {
            startClosed = true;
        }

        if (std::strcmp(argv[i], "-elevated") == 0)
        {
            ignoreMutex = true;
        }

        if (std::strcmp(argv[i], "-restartCamera") == 0)
        {
            if (is_elevated)
            {
                hid_handler.startup_camera_restart.store(true);
            }

            hasResetCamera = true;
            ignoreMutex = true;
        }
    }

    // Register the eyetracking URI
    wchar_t cwd[MAX_PATH];
    GetModuleFileNameW(NULL, cwd, MAX_PATH);
    PathRemoveFileSpecW(cwd);

    std::wstring et_path = std::wstring(cwd) + L"\\eyetracking\\ETClient\\BeyondET.exe";
    RegisterUri(L"bset", et_path);

    // Get current executable path
    char exePath[MAX_PATH];
    DWORD len = GetModuleFileNameA(NULL, exePath, MAX_PATH);

    if (len != 0)
    {
#ifndef _DEBUG
        // Force working directory to the executable
        char exeDir[MAX_PATH];
        strcpy_s(exeDir, exePath);

        if (PathRemoveFileSpecA(exeDir))
        {
            SetCurrentDirectoryA(exeDir);
        }
#endif
    }

    // If program is already running, don't allow a second instance.
    LPCWSTR named_mutex = L"com_bigscreen_hmd_utility";
    HANDLE program_mutex = CreateMutex(NULL, TRUE, named_mutex);
    if (!ignoreMutex && ERROR_ALREADY_EXISTS == GetLastError())
    {
        // Restore window minimized / background
        if (!ui.bring_foreground(true))
        {
            ui.create_window(startClosed);
            MessageBox(glfwGetWin32Window(ui.window), LWSTR(StringID::ERROR_UTILITY_ALREADY_RUNNING).c_str(), LWSTR(StringID::POPUP_TITLE_ERROR).c_str(), MB_OK);
        }
        
        return 0;
    }
    // Create window
    else
    {
        ui.create_window(startClosed);
    }

    if (hasResetCamera && !is_elevated)
    {
        MessageBox(glfwGetWin32Window(ui.window), LWSTR(StringID::ERROR_ADMIN_PRIVILEGES_REQUIRED).c_str(), LWSTR(StringID::POPUP_TITLE_ERROR).c_str(), MB_OK);
    }

    if (!IsDebuggerPresent())
    {
        const wchar_t* progDataPath = _wgetenv(L"PROGRAMDATA");

        if (progDataPath)
        {
            std::wstring linkPath = std::wstring(progDataPath) + L"\\Microsoft\\Windows\\Start Menu\\Programs\\StartUp\\BeyondHID.lnk";

            DWORD dwAttrib = GetFileAttributes(linkPath.c_str());
            if (dwAttrib != INVALID_FILE_ATTRIBUTES && !(dwAttrib & FILE_ATTRIBUTE_DIRECTORY))
            {
                // If we're already in the startup folder, no need for elevated priviledges

            }
            else
            {
                // Restart with elevated priviledges.
                if (!is_elevated())
                {
                    ShellExecute(NULL, L"runas", L"BeyondHID.exe", L"-elevated", NULL, SW_SHOWNORMAL);
                    return 0;
                }

                // Save a shortcut into the startup folder so it runs on boot.
                CoInitialize(nullptr);
                char exe_path[MAX_PATH];
                GetModuleFileNameA(GetModuleHandle(NULL), exe_path, MAX_PATH);
                WCHAR exe_wstr[MAX_PATH];
                MultiByteToWideChar(CP_ACP, 0, exe_path, -1, exe_wstr, MAX_PATH);
                PathRemoveFileSpecA(exe_path);
                WCHAR dir_wstr[MAX_PATH];
                MultiByteToWideChar(CP_ACP, 0, exe_path, -1, dir_wstr, MAX_PATH);

                HRESULT hr = create_link(dir_wstr,
                    exe_wstr,
                    linkPath.c_str(), L"Bigscreen");
            }
        }
    }

    // Setup states.
#ifndef _DEBUG
    #ifdef APP_VERSION
        ui.app_version = APP_VERSION;
    #endif
#endif

    ui.alive.store(true);
    hid_handler.ready.store(false);

    // HID thread.
    hid_handler.run();

    // Process watch
    process_watch.run();

    // Wait for hid handler to be setup.
    while (!hid_handler.ready.load())
    {
        Sleep(10); // Don't clog the core.
    }

    // UI (main) thread.
    ui.run(); // Blocks until execution stops.

    // Join and cleanup threads.
    hid_handler.shutdown();
    ui.shutdown();
    utils.shutdown();
    process_watch.shutdown();

    CoUninitialize();

    ReleaseMutex(program_mutex);
    CloseHandle(program_mutex);
    return 0;
}