#include "update.h"
#include "json.hpp"

namespace ks
{
namespace
{
struct Handle
{
    HANDLE value{};
    ~Handle() { if (value) CloseHandle(value); }
};
std::wstring quoteArgument(const std::wstring &value)
{
    // Windows argv escaping, including a path ending with a backslash.
    std::wstring out = L"\"";
    size_t slashes = 0;
    for (auto ch : value)
    {
        if (ch == L'\\') { ++slashes; continue; }
        out.append(ch == L'"' ? slashes * 2 + 1 : slashes, L'\\');
        slashes = 0;
        out += ch;
    }
    out.append(slashes * 2, L'\\');
    return out + L'"';
}
}
UpdateResult runUpdate(bool install, const std::wstring &expectedVersion, bool showProgress)
{
    UpdateResult result;
    GUID guid{};
    if (FAILED(CoCreateGuid(&guid))) throw std::runtime_error("Cannot create update identifier");
    wchar_t id[40]{};
    StringFromGUID2(guid, id, 40);
    auto work = std::filesystem::temp_directory_path() / (L"KSwitcher-update-" + std::wstring(id));
    std::filesystem::create_directory(work);
    bool detached = false;
    Handle child;
    try
    {
        HRSRC resource = FindResourceW(nullptr, MAKEINTRESOURCEW(201), RT_RCDATA);
        HGLOBAL loaded = resource ? LoadResource(nullptr, resource) : nullptr;
        auto bytes = loaded ? static_cast<const unsigned char *>(LockResource(loaded)) : nullptr;
        if (!bytes) throw std::runtime_error("The updater resource is missing");
        atomicWrite(work / L"update.ps1", Bytes(bytes, bytes + SizeofResource(nullptr, resource)));
        std::wstring readyName = L"Local\\KSwitcher.Update.Ready." + std::wstring(id);
        Handle ready{CreateEventW(nullptr, TRUE, FALSE, readyName.c_str())};
        if (!ready.value) throw std::runtime_error("Cannot create updater handshake");
        wchar_t system[MAX_PATH]{};
        if (!GetSystemDirectoryW(system, MAX_PATH)) throw std::runtime_error("Cannot locate Windows PowerShell");
        auto powershell = std::filesystem::path(system) / L"WindowsPowerShell/v1.0/powershell.exe";
        std::wstring command = quoteArgument(powershell.wstring()) + L" -NoLogo -NoProfile -NonInteractive -ExecutionPolicy Bypass -File " +
            quoteArgument((work / L"update.ps1").wstring()) + L" -Mode " + (install ? L"Install" : L"Check") +
            L" -CurrentVersion " + quoteArgument(AppVersion) + L" -WorkDirectory " + quoteArgument(work.wstring()) +
            L" -InstallDirectory " + quoteArgument(std::filesystem::path(executable()).parent_path().wstring()) +
            L" -ParentProcessId " + std::to_wstring(GetCurrentProcessId()) + L" -ReadyEvent " + quoteArgument(readyName);
        if (install) command += L" -ExpectedVersion " + quoteArgument(expectedVersion);
        if (showProgress) command += L" -ShowProgress";
        STARTUPINFOW startup{sizeof(startup)};
        PROCESS_INFORMATION process{};
        if (!CreateProcessW(powershell.c_str(), command.data(), nullptr, nullptr, FALSE, CREATE_NO_WINDOW,
                            nullptr, work.c_str(), &startup, &process))
            throw std::runtime_error("Cannot start Windows PowerShell updater");
        child.value = process.hProcess;
        Handle thread{process.hThread};
        HANDLE waits[]{ready.value, child.value};
        DWORD waited = WaitForMultipleObjects(2, waits, FALSE, 240000);
        if (waited != WAIT_OBJECT_0 && waited != WAIT_OBJECT_0 + 1)
        {
            // The helper cannot replace files while this process is alive.
            TerminateProcess(child.value, 1);
            WaitForSingleObject(child.value, 5000);
            throw std::runtime_error("Update preparation timed out; the installed application was not changed");
        }
        auto data = readFile(work / L"result.json", 64 * 1024);
        auto json = nlohmann::json::parse(data.begin(), data.end());
        result.state = wide(json.at("state").get<std::string>());
        result.version = wide(json.value("version", std::string{}));
        result.error = wide(json.value("error", std::string{}));
        detached = install && waited == WAIT_OBJECT_0 && result.state == L"prepared";
        if (result.state == L"prepared" && !detached)
            throw std::runtime_error("Update helper exited before installation");
        if (!detached) std::filesystem::remove_all(work);
        return result;
    }
    catch (...)
    {
        if (!detached && child.value && WaitForSingleObject(child.value, 0) == WAIT_TIMEOUT)
        {
            TerminateProcess(child.value, 1);
            WaitForSingleObject(child.value, 5000);
        }
        if (!detached) { std::error_code ignored; std::filesystem::remove_all(work, ignored); }
        throw;
    }
}
bool confirmUpdatedStartup(const std::wstring &eventName)
{
    if (!eventName.starts_with(L"Local\\KSwitcher.Update.Started.")) return false;
    Handle event{OpenEventW(EVENT_MODIFY_STATE, FALSE, eventName.c_str())};
    return event.value && SetEvent(event.value);
}
}