#include "Common.h"
#include "ServiceManager.h"
#include "Logger.h"

bool EnsureWin32Service(
    const wchar_t* serviceName,
    const wchar_t* binaryPath,
    bool startNow)
{
    if (!serviceName || !binaryPath)
        return false;

    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        Log(L"[-] OpenSCManager failed: %lu", GetLastError());
        return false;
    }

    SC_HANDLE svc = CreateServiceW(
        scm,
        serviceName,
        serviceName,
        SERVICE_CHANGE_CONFIG | SERVICE_START | SERVICE_QUERY_STATUS,
        SERVICE_WIN32_OWN_PROCESS,
        SERVICE_AUTO_START,
        SERVICE_ERROR_NORMAL,
        binaryPath,
        nullptr,
        nullptr,
        nullptr,
        nullptr,
        nullptr);

    DWORD lastErr = GetLastError();

    if (!svc && lastErr == ERROR_SERVICE_EXISTS)
    {
        svc = OpenServiceW(
            scm,
            serviceName,
            SERVICE_CHANGE_CONFIG | SERVICE_START | SERVICE_QUERY_STATUS);
    }

    if (!svc)
    {
        Log(L"[-] Open/Create service failed: %lu", lastErr);
        CloseServiceHandle(scm);
        return false;
    }

    if (!ChangeServiceConfigW(
        svc,
        SERVICE_NO_CHANGE,
        SERVICE_AUTO_START,
        SERVICE_ERROR_NORMAL,
        binaryPath,
        nullptr,
        nullptr,
        nullptr,
        nullptr,
        nullptr,
        nullptr))
    {
        Log(L"[!] ChangeServiceConfigW failed: %lu", GetLastError());
    }

    if (startNow)
    {
        SERVICE_STATUS st{};
        bool running = QueryServiceStatus(svc, &st) && st.dwCurrentState == SERVICE_RUNNING;

        if (!running)
        {
            if (!StartServiceW(svc, 0, nullptr))
            {
                DWORD e = GetLastError();

                if (e != ERROR_SERVICE_ALREADY_RUNNING)
                    Log(L"[-] StartService failed: %lu", e);
            }
            else
            {
                Log(L"[+] Service started");
            }
        }
        else
        {
            Log(L"[+] Service already running");
        }
    }

    CloseServiceHandle(svc);
    CloseServiceHandle(scm);

    return true;
}

bool IsServiceRunning(const wchar_t* serviceName)
{
    if (!serviceName)
        return false;

    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_CONNECT);
    if (!scm)
        return false;

    SC_HANDLE svc = OpenServiceW(scm, serviceName, SERVICE_QUERY_STATUS);
    if (!svc)
    {
        CloseServiceHandle(scm);
        return false;
    }

    SERVICE_STATUS st{};
    bool running = QueryServiceStatus(svc, &st) && st.dwCurrentState == SERVICE_RUNNING;

    CloseServiceHandle(svc);
    CloseServiceHandle(scm);

    return running;
}

bool EnsureLocalSystemService(
    const wchar_t* serviceName,
    const wchar_t* binaryPath,
    bool startNow,
    bool restartOnFailure,
    const wchar_t* displayName)
{
    if (!serviceName || !binaryPath)
        return false;

    const wchar_t* display = (displayName && *displayName) ? displayName : serviceName;

    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_ALL_ACCESS);
    if (!scm)
    {
        Log(L"[-] EnsureLocalSystemService: OpenSCManager failed: %lu", GetLastError());
        return false;
    }
    SC_HANDLE svc = CreateServiceW(
        scm,
        serviceName,
        display,
        SERVICE_CHANGE_CONFIG | SERVICE_START | SERVICE_QUERY_STATUS,
        SERVICE_WIN32_OWN_PROCESS,
        SERVICE_AUTO_START,
        SERVICE_ERROR_NORMAL,
        binaryPath,
        nullptr,
        nullptr,
        nullptr,
        L"LocalSystem",
        nullptr);

    DWORD createError = GetLastError();

    if (!svc && createError == ERROR_SERVICE_EXISTS)
    {
        svc = OpenServiceW(
            scm,
            serviceName,
            SERVICE_CHANGE_CONFIG | SERVICE_START | SERVICE_QUERY_STATUS);
    }

    if (!svc)
    {
        Log(L"[-] EnsureLocalSystemService: create/open failed: %lu", createError);
        CloseServiceHandle(scm);
        return false;
    }

    if (!ChangeServiceConfigW(
        svc,
        SERVICE_NO_CHANGE,
        SERVICE_AUTO_START,
        SERVICE_ERROR_NORMAL,
        binaryPath,
        nullptr,
        nullptr,
        nullptr,
        L"LocalSystem",
        nullptr,
        nullptr))
    {
        Log(L"[!] EnsureLocalSystemService: ChangeServiceConfigW failed: %lu", GetLastError());
    }

    if (restartOnFailure)
    {
        SC_ACTION actions[3]{};
        actions[0].Type = SC_ACTION_RESTART;
        actions[0].Delay = 5000;
        actions[1].Type = SC_ACTION_RESTART;
        actions[1].Delay = 10000;
        actions[2].Type = SC_ACTION_RESTART;
        actions[2].Delay = 30000;

        SERVICE_FAILURE_ACTIONSW failure{};
        failure.dwResetPeriod = 86400;
        failure.cActions = 3;
        failure.lpsaActions = actions;

        ChangeServiceConfig2W(svc, SERVICE_CONFIG_FAILURE_ACTIONS, &failure);

        SERVICE_FAILURE_ACTIONS_FLAG flag{};
        flag.fFailureActionsOnNonCrashFailures = TRUE;
        ChangeServiceConfig2W(svc, SERVICE_CONFIG_FAILURE_ACTIONS_FLAG, &flag);
    }

    if (startNow)
    {
        SERVICE_STATUS st{};
        bool running = QueryServiceStatus(svc, &st) && st.dwCurrentState == SERVICE_RUNNING;

        if (!running)
        {
            if (!StartServiceW(svc, 0, nullptr))
            {
                DWORD e = GetLastError();

                if (e != ERROR_SERVICE_ALREADY_RUNNING)
                {
                    Log(L"[-] EnsureLocalSystemService: StartService failed: %lu", e);
                    CloseServiceHandle(svc);
                    CloseServiceHandle(scm);
                    return false;
                }
            }
        }
    }

    CloseServiceHandle(svc);
    CloseServiceHandle(scm);
    return true;
}

bool SetServiceDescription(
    const wchar_t* serviceName,
    const wchar_t* description)
{
    if (!serviceName || !description)
        return false;

    SC_HANDLE scm = OpenSCManagerW(nullptr, nullptr, SC_MANAGER_CONNECT);

    if (!scm)
        return false;

    SC_HANDLE svc = OpenServiceW(scm, serviceName, SERVICE_CHANGE_CONFIG);

    bool ok = false;

    if (svc)
    {
        SERVICE_DESCRIPTIONW desc{};
        desc.lpDescription = const_cast<LPWSTR>(description);

        ok = ChangeServiceConfig2W(svc, SERVICE_CONFIG_DESCRIPTION, &desc) != FALSE;
        CloseServiceHandle(svc);
    }

    CloseServiceHandle(scm);
    return ok;
}