#include "platform.h"
#include <sddl.h>
#include <shellapi.h>
#include <shlobj.h>
#include <taskschd.h>
namespace ks
{
namespace
{
template <class T> struct Com
{
    T *p{};
    ~Com()
    {
        if (p)
            p->Release();
    }
    T **out()
    {
        return &p;
    }
    T *operator->()
    {
        return p;
    }
};
struct Str
{
    BSTR value;
    Str(const wchar_t *s) : value(SysAllocString(s))
    {
        if (!value)
            throw std::bad_alloc();
    }
    ~Str()
    {
        SysFreeString(value);
    }
    operator BSTR() const
    {
        return value;
    }
};
void check(HRESULT hr)
{
    if (FAILED(hr))
        throw std::runtime_error("Cannot configure Windows autostart (HRESULT " +
                                 std::to_string((unsigned long)hr) + ")");
}
bool missing(HRESULT hr)
{
    return hr == HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND) || hr == HRESULT_FROM_WIN32(ERROR_PATH_NOT_FOUND);
}
void runEntry(bool enabled)
{
    HKEY key{};
    if (RegCreateKeyExW(HKEY_CURRENT_USER, L"Software\\Microsoft\\Windows\\CurrentVersion\\Run", 0, nullptr,
                        0, KEY_SET_VALUE, nullptr, &key, nullptr) != ERROR_SUCCESS)
        throw std::runtime_error("Cannot access autostart settings");
    auto command = L"\"" + executable() + L"\"";
    auto status = enabled ? RegSetValueExW(key, L"KSwitcher", 0, REG_SZ, (const BYTE *)command.c_str(),
                                           DWORD((command.size() + 1) * sizeof(wchar_t)))
                          : RegDeleteValueW(key, L"KSwitcher");
    RegCloseKey(key);
    if (status != ERROR_SUCCESS && status != ERROR_FILE_NOT_FOUND)
        throw std::runtime_error("Cannot save autostart settings");
}
} // namespace
std::wstring currentUserSid()
{
    HANDLE token{};
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
        throw std::runtime_error("Cannot identify Windows account");
    DWORD size{};
    GetTokenInformation(token, TokenUser, nullptr, 0, &size);
    Bytes data(size);
    BOOL ok = GetTokenInformation(token, TokenUser, data.data(), size, &size);
    CloseHandle(token);
    if (!ok)
        throw std::runtime_error("Cannot identify Windows account");
    LPWSTR sid{};
    if (!ConvertSidToStringSidW(((TOKEN_USER *)data.data())->User.Sid, &sid))
        throw std::runtime_error("Cannot read account SID");
    std::wstring result = sid;
    LocalFree(sid);
    return result;
}
void configureAutostart(int mode)
{
    Com<ITaskService> service;
    check(CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService,
                           (void **)service.out()));
    VARIANT empty;
    VariantInit(&empty);
    check(service->Connect(empty, empty, empty, empty));
    Com<ITaskFolder> folder;
    check(service->GetFolder(Str(L"\\"), folder.out()));
    Com<IRegisteredTask> existing;
    HRESULT found = folder->GetTask(Str(L"KSwitcher"), existing.out());
    if (FAILED(found) && !missing(found))
        check(found);
    if (mode != 2)
    {
        if (SUCCEEDED(found))
            check(existing->put_Enabled(VARIANT_FALSE));
        runEntry(mode == 1);
        return;
    }
    auto sid = currentUserSid();
    Str user(sid.c_str());
    Com<ITaskDefinition> definition;
    check(service->NewTask(0, definition.out()));
    Com<IPrincipal> principal;
    check(definition->get_Principal(principal.out()));
    check(principal->put_UserId(user));
    check(principal->put_LogonType(TASK_LOGON_INTERACTIVE_TOKEN));
    check(principal->put_RunLevel(TASK_RUNLEVEL_HIGHEST));
    Com<ITaskSettings> settings;
    check(definition->get_Settings(settings.out()));
    check(settings->put_DisallowStartIfOnBatteries(VARIANT_FALSE));
    check(settings->put_StopIfGoingOnBatteries(VARIANT_FALSE));
    check(settings->put_ExecutionTimeLimit(Str(L"PT0S")));
    check(settings->put_MultipleInstances(TASK_INSTANCES_IGNORE_NEW));
    Com<ITriggerCollection> triggers;
    check(definition->get_Triggers(triggers.out()));
    Com<ITrigger> trigger;
    check(triggers->Create(TASK_TRIGGER_LOGON, trigger.out()));
    Com<ILogonTrigger> logon;
    check(trigger->QueryInterface(IID_ILogonTrigger, (void **)logon.out()));
    check(logon->put_UserId(user));
    Com<IActionCollection> actions;
    check(definition->get_Actions(actions.out()));
    Com<IAction> action;
    check(actions->Create(TASK_ACTION_EXEC, action.out()));
    Com<IExecAction> exec;
    check(action->QueryInterface(IID_IExecAction, (void **)exec.out()));
    auto path = executable();
    check(exec->put_Path(Str(path.c_str())));
    check(exec->put_WorkingDirectory(Str(std::filesystem::path(path).parent_path().c_str())));
    VARIANT account;
    VariantInit(&account);
    account.vt = VT_BSTR;
    account.bstrVal = user;
    Com<IRegisteredTask> installed;
    check(folder->RegisterTaskDefinition(Str(L"KSwitcher"), definition.p, TASK_CREATE_OR_UPDATE, account,
                                         empty, TASK_LOGON_INTERACTIVE_TOKEN, empty, installed.out()));
    runEntry(false);
}
void setAutostart(bool enabled, bool elevated)
{
    int mode = enabled ? (elevated ? 2 : 1) : 0;
    if (IsUserAnAdmin())
    {
        configureAutostart(mode);
        return;
    }
    // Ordinary startup needs no UAC unless an elevated task must be changed.
    if (mode != 2)
    {
        Com<ITaskService> service;
        HRESULT hr = CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService,
                                      (void **)service.out());
        VARIANT empty;
        VariantInit(&empty);
        Com<ITaskFolder> folder;
        Com<IRegisteredTask> task;
        if (SUCCEEDED(hr))
            hr = service->Connect(empty, empty, empty, empty);
        if (SUCCEEDED(hr))
            hr = service->GetFolder(Str(L"\\"), folder.out());
        if (SUCCEEDED(hr))
            hr = folder->GetTask(Str(L"KSwitcher"), task.out());
        if (missing(hr))
        {
            runEntry(mode == 1);
            return;
        }
        check(hr);
    }
    auto file = executable(),
         args = L"--configure-autostart=" + std::to_wstring(mode) + L" " + currentUserSid();
    SHELLEXECUTEINFOW request{sizeof(request)};
    request.fMask = SEE_MASK_NOCLOSEPROCESS;
    request.lpVerb = L"runas";
    request.lpFile = file.c_str();
    request.lpParameters = args.c_str();
    request.nShow = SW_HIDE;
    if (!ShellExecuteExW(&request) || !request.hProcess)
        throw std::runtime_error("Administrator approval was cancelled or unavailable");
    // Continue servicing sent hook messages while UAC/helper runs, without
    // dispatching another Settings Save or global command reentrantly.
    DWORD wait{};
    while ((wait = MsgWaitForMultipleObjects(1, &request.hProcess, FALSE, INFINITE, QS_SENDMESSAGE)) ==
           WAIT_OBJECT_0 + 1)
    {
        MSG message;
        PeekMessageW(&message, nullptr, WM_NULL, WM_NULL, PM_NOREMOVE);
    }
    DWORD code = 1;
    GetExitCodeProcess(request.hProcess, &code);
    CloseHandle(request.hProcess);
    if (wait != WAIT_OBJECT_0 || code != 0)
        throw std::runtime_error("Elevated autostart was not changed");
}
} // namespace ks