package com.fatico.winthing.systems.system;

import com.fatico.winthing.windows.SystemException;
import com.fatico.winthing.windows.jna.Advapi32;
import com.fatico.winthing.windows.jna.Kernel32;
import com.google.common.collect.ImmutableList;
import com.google.inject.Inject;
import com.sun.jna.platform.win32.Kernel32Util;
import com.sun.jna.platform.win32.Shell32;
import com.sun.jna.platform.win32.Tlhelp32;
import com.sun.jna.platform.win32.WinDef;
import com.sun.jna.platform.win32.WinNT;
import com.sun.jna.platform.win32.WinUser;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

public class SystemService {

    private static final List<String> REQUIRED_PRIVILEGES = ImmutableList.of(
        WinNT.SE_SHUTDOWN_NAME
    );

    private final Kernel32 kernel32;
    private final Advapi32 advapi32;
    private final Shell32 shell32;

    @Inject
    public SystemService(final Kernel32 kernel32, final Advapi32 advapi32,
            final Shell32 shell32) throws SystemException {
        this.kernel32 = Objects.requireNonNull(kernel32);
        this.advapi32 = Objects.requireNonNull(advapi32);
        this.shell32 = Objects.requireNonNull(shell32);
        escalatePrivileges(REQUIRED_PRIVILEGES);
    }

    public void shutdown() throws SystemException {
        final boolean success = advapi32.InitiateSystemShutdown(
            null,
            null,
            new WinDef.DWORD(0),
            true,
            false
        );
        if (!success) {
            throw new SystemException(Kernel32Util.formatMessage(kernel32.GetLastError()));
        }
    }

    public void reboot() throws SystemException {
        final boolean success = advapi32.InitiateSystemShutdown(
            null,
            null,
            new WinDef.DWORD(0),
            true,
            true
        );
        if (!success) {
            throw new SystemException(Kernel32Util.formatMessage(kernel32.GetLastError()));
        }
    }
    
    public void suspend() throws SystemException {
        final boolean success = kernel32.SetSystemPowerState(
            true,
            false
        );
        if (!success) {
            throw new SystemException(Kernel32Util.formatMessage(kernel32.GetLastError()));
        }
    }

    public void hibernate() throws SystemException {
        final boolean success = kernel32.SetSystemPowerState(
                false,
                false
            );
        if (!success) {
            throw new SystemException(Kernel32Util.formatMessage(kernel32.GetLastError()));
        }
    }

    public void run(final String command, final String parameters, final String workingDirectory)
            throws SystemException {
        final WinDef.INT_PTR result = shell32.ShellExecute(
            null,
            "open",
            Objects.requireNonNull(command),
            Objects.requireNonNull(parameters),
            workingDirectory,
            WinUser.SW_SHOWNORMAL
        );
        if (result.intValue() <= 32) {
            throw new SystemException("Could not run command: " + command + " " + parameters);
        }
    }

    public void open(final String uri) throws SystemException {
        final WinDef.INT_PTR result = shell32.ShellExecute(
            null,
            "open",
            Objects.requireNonNull(uri),
            null,
            null,
            WinUser.SW_SHOWNORMAL
        );
        if (result.intValue() <= 32) {
            throw new SystemException("Cannot open URI: " + uri);
        }
    }

    @SuppressFBWarnings("DM_CONVERT_CASE")
    public Map<Integer, String> findProcesses(final String nameFragment) {
        Objects.requireNonNull(nameFragment);

        final String lowercaseNameFragment = nameFragment.toLowerCase();
        final Map<Integer, String> processIds = new HashMap<>();

        final WinNT.HANDLE snapshot = kernel32.CreateToolhelp32Snapshot(
            Tlhelp32.TH32CS_SNAPPROCESS,
            null
        );
        try {
            final Tlhelp32.PROCESSENTRY32.ByReference entryReference =
                new Tlhelp32.PROCESSENTRY32.ByReference();
            if (kernel32.Process32First(snapshot, entryReference)) {
                while (kernel32.Process32Next(snapshot, entryReference)) {
                    final String processName = new String(entryReference.szExeFile).trim();
                    if (processName.toLowerCase().contains(lowercaseNameFragment)) {
                        processIds.put(entryReference.th32ProcessID.intValue(), processName);
                    }
                }
            }
        } finally {
            kernel32.CloseHandle(snapshot);
        }

        return processIds;
    }

    private void escalatePrivileges(final List<String> requiredPrivilegeNames)
            throws SystemException {
        final WinNT.HANDLE accessToken;
        {
            final WinNT.HANDLEByReference tokenReference = new WinNT.HANDLEByReference();
            final boolean success = advapi32.OpenProcessToken(
                kernel32.GetCurrentProcess(),
                WinNT.TOKEN_ADJUST_PRIVILEGES | WinNT.TOKEN_QUERY,
                tokenReference
            );
            if (!success) {
                throw new SystemException("Cannot open access token");
            }
            accessToken = tokenReference.getValue();
        }

        final WinNT.TOKEN_PRIVILEGES privileges = new WinNT.TOKEN_PRIVILEGES(
            requiredPrivilegeNames.size()
        );
        {
            privileges.PrivilegeCount.setValue(requiredPrivilegeNames.size());
            int index = 0;
            for (final String privilegeName : requiredPrivilegeNames) {
                final WinNT.LUID luid = new WinNT.LUID();
                {
                    final boolean success = advapi32.LookupPrivilegeValue(
                        null,
                        privilegeName,
                        luid
                    );
                    if (!success) {
                        throw new SystemException("Cannot find privilege " + privilegeName);
                    }
                }
                privileges.Privileges[index] = new WinNT.LUID_AND_ATTRIBUTES();
                privileges.Privileges[index].Luid = luid;
                privileges.Privileges[index].Attributes.setValue(WinNT.SE_PRIVILEGE_ENABLED);
                index++;
            }
        }

        {
            final boolean success = advapi32.AdjustTokenPrivileges(
                accessToken,
                false,
                privileges,
                privileges.size(),
                null,
                null
            );
            if (!success) {
                throw new SystemException(
                    "Cannot obtain required privileges: "
                        + Kernel32Util.formatMessage(kernel32.GetLastError())
                );
            }
        }
    }

}