#include "stdafx.h"
#include "StartupManager.h"

#include <winrt/Windows.Foundation.h>
#include <winrt/Windows.ApplicationModel.h>

#pragma comment(lib, "windowsapp.lib")

using namespace winrt;
using namespace winrt::Windows::ApplicationModel;

namespace {
    const wchar_t* kTaskId = L"TaskbarCalculatorStartup";
    const wchar_t* kRunKey = L"Software\\Microsoft\\Windows\\CurrentVersion\\Run";
    const wchar_t* kRunVal = L"TaskbarCalculator";

    // --- HKCU\...\Run fallback (used when the app has no package identity) ---
    bool RunKeyExists() {
        HKEY h;
        if (RegOpenKeyExW(HKEY_CURRENT_USER, kRunKey, 0, KEY_QUERY_VALUE, &h) != ERROR_SUCCESS)
            return false;
        LONG r = RegQueryValueExW(h, kRunVal, nullptr, nullptr, nullptr, nullptr);
        RegCloseKey(h);
        return r == ERROR_SUCCESS;
    }

    bool RunKeySet(bool on) {
        HKEY h;
        if (RegCreateKeyExW(HKEY_CURRENT_USER, kRunKey, 0, nullptr, 0,
                            KEY_SET_VALUE, nullptr, &h, nullptr) != ERROR_SUCCESS)
            return false;
        LONG r;
        if (on) {
            wchar_t path[MAX_PATH];
            GetModuleFileNameW(nullptr, path, MAX_PATH);
            wchar_t quoted[MAX_PATH + 2];
            wsprintfW(quoted, L"\"%s\"", path);
            r = RegSetValueExW(h, kRunVal, 0, REG_SZ, (const BYTE*)quoted,
                               (DWORD)((lstrlenW(quoted) + 1) * sizeof(wchar_t)));
        } else {
            r = RegDeleteValueW(h, kRunVal);
            if (r == ERROR_FILE_NOT_FOUND) r = ERROR_SUCCESS;
        }
        RegCloseKey(h);
        return r == ERROR_SUCCESS;
    }
}

namespace Startup {

State Get() {
    try {
        auto task = StartupTask::GetAsync(kTaskId).get();
        switch (task.State()) {
            case StartupTaskState::Enabled:
            case StartupTaskState::EnabledByPolicy:
                return State::Enabled;
            default:
                return State::Disabled;
        }
    } catch (...) {
        // Unpackaged (classic install): use the HKCU\Run fallback.
        return RunKeyExists() ? State::Enabled : State::Disabled;
    }
}

bool Set(bool on) {
    try {
        auto task = StartupTask::GetAsync(kTaskId).get();
        if (on) {
            auto result = task.RequestEnableAsync().get();
            return result == StartupTaskState::Enabled
                || result == StartupTaskState::EnabledByPolicy;
        } else {
            task.Disable();
            return true;
        }
    } catch (...) {
        return RunKeySet(on);
    }
}

}