#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <wtsapi32.h>

// PROC_THREAD_ATTRIBUTE_JOB_LIST is currently missing from MinGW headers
#ifndef PROC_THREAD_ATTRIBUTE_JOB_LIST
#define PROC_THREAD_ATTRIBUTE_JOB_LIST ProcThreadAttributeValue(13, FALSE, TRUE, FALSE)
#endif

SERVICE_STATUS_HANDLE service_status_handle;
SERVICE_STATUS service_status;
HANDLE stop_event;

#define SERVICE_NAME "SunshineSvc"

DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpContext) {
  switch(dwControl) {
  case SERVICE_CONTROL_INTERROGATE:
    return NO_ERROR;

  case SERVICE_CONTROL_STOP:
    // Let SCM know we're stopping in up to 30 seconds
    service_status.dwCurrentState     = SERVICE_STOP_PENDING;
    service_status.dwControlsAccepted = 0;
    service_status.dwWaitHint         = 30 * 1000;
    SetServiceStatus(service_status_handle, &service_status);

    // Trigger ServiceMain() to start cleanup
    SetEvent(stop_event);
    return NO_ERROR;

  default:
    return NO_ERROR;
  }
}

HANDLE CreateJobObjectForChildProcess() {
  HANDLE job_handle = CreateJobObjectW(NULL, NULL);
  if(!job_handle) {
    return NULL;
  }

  JOBOBJECT_EXTENDED_LIMIT_INFORMATION job_limit_info = {};

  // Kill Sunshine.exe when the final job object handle is closed (which will happen if we terminate unexpectedly).
  // This ensures we don't leave an orphaned Sunshine.exe running with an inherited handle to our log file.
  job_limit_info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;

  // Allow Sunshine.exe to use CREATE_BREAKAWAY_FROM_JOB when spawning processes to ensure they can to live beyond
  // the lifetime of SunshineSvc.exe. This avoids unexpected user data loss if we crash or are killed.
  job_limit_info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_BREAKAWAY_OK;

  if(!SetInformationJobObject(job_handle, JobObjectExtendedLimitInformation, &job_limit_info, sizeof(job_limit_info))) {
    CloseHandle(job_handle);
    return NULL;
  }

  return job_handle;
}

LPPROC_THREAD_ATTRIBUTE_LIST AllocateProcThreadAttributeList(DWORD attribute_count) {
  SIZE_T size;
  InitializeProcThreadAttributeList(NULL, attribute_count, 0, &size);

  auto list = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(GetProcessHeap(), 0, size);
  if(list == NULL) {
    return NULL;
  }

  if(!InitializeProcThreadAttributeList(list, attribute_count, 0, &size)) {
    HeapFree(GetProcessHeap(), 0, list);
    return NULL;
  }

  return list;
}

HANDLE DuplicateTokenForConsoleSession() {
  auto console_session_id = WTSGetActiveConsoleSessionId();
  if(console_session_id == 0xFFFFFFFF) {
    // No console session yet
    return NULL;
  }

  HANDLE current_token;
  if(!OpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE, &current_token)) {
    return NULL;
  }

  // Duplicate our own LocalSystem token
  HANDLE new_token;
  if(!DuplicateTokenEx(current_token, TOKEN_ALL_ACCESS, NULL, SecurityImpersonation, TokenPrimary, &new_token)) {
    CloseHandle(current_token);
    return NULL;
  }

  CloseHandle(current_token);

  // Change the duplicated token to the console session ID
  if(!SetTokenInformation(new_token, TokenSessionId, &console_session_id, sizeof(console_session_id))) {
    CloseHandle(new_token);
    return NULL;
  }

  return new_token;
}

HANDLE OpenLogFileHandle() {
  WCHAR log_file_name[MAX_PATH];

  // Create sunshine.log in the Temp folder (usually %SYSTEMROOT%\Temp)
  GetTempPathW(_countof(log_file_name), log_file_name);
  wcscat_s(log_file_name, L"sunshine.log");

  // The file handle must be inheritable for our child process to use it
  SECURITY_ATTRIBUTES security_attributes = { sizeof(security_attributes), NULL, TRUE };

  // Overwrite the old sunshine.log
  return CreateFileW(log_file_name,
    GENERIC_WRITE,
    FILE_SHARE_READ,
    &security_attributes,
    CREATE_ALWAYS,
    0,
    NULL);
}

VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
  service_status_handle = RegisterServiceCtrlHandlerEx(SERVICE_NAME, HandlerEx, NULL);
  if(service_status_handle == NULL) {
    // Nothing we can really do here but terminate ourselves
    ExitProcess(GetLastError());
    return;
  }

  // Tell SCM we're starting
  service_status.dwServiceType             = SERVICE_WIN32_OWN_PROCESS;
  service_status.dwServiceSpecificExitCode = 0;
  service_status.dwWin32ExitCode           = NO_ERROR;
  service_status.dwWaitHint                = 0;
  service_status.dwControlsAccepted        = 0;
  service_status.dwCheckPoint              = 0;
  service_status.dwCurrentState            = SERVICE_START_PENDING;
  SetServiceStatus(service_status_handle, &service_status);

  stop_event = CreateEventA(NULL, TRUE, FALSE, NULL);
  if(stop_event == NULL) {
    // Tell SCM we failed to start
    service_status.dwWin32ExitCode = GetLastError();
    service_status.dwCurrentState  = SERVICE_STOPPED;
    SetServiceStatus(service_status_handle, &service_status);
    return;
  }

  auto log_file_handle = OpenLogFileHandle();
  if(log_file_handle == INVALID_HANDLE_VALUE) {
    // Tell SCM we failed to start
    service_status.dwWin32ExitCode = GetLastError();
    service_status.dwCurrentState  = SERVICE_STOPPED;
    SetServiceStatus(service_status_handle, &service_status);
    return;
  }

  auto job_handle = CreateJobObjectForChildProcess();
  if(job_handle == NULL) {
    // Tell SCM we failed to start
    service_status.dwWin32ExitCode = GetLastError();
    service_status.dwCurrentState  = SERVICE_STOPPED;
    SetServiceStatus(service_status_handle, &service_status);
    return;
  }

  // We can use a single STARTUPINFOEXW for all the processes that we launch
  STARTUPINFOEXW startup_info         = {};
  startup_info.StartupInfo.cb         = sizeof(startup_info);
  startup_info.StartupInfo.lpDesktop  = (LPWSTR)L"winsta0\\default";
  startup_info.StartupInfo.dwFlags    = STARTF_USESTDHANDLES;
  startup_info.StartupInfo.hStdInput  = NULL;
  startup_info.StartupInfo.hStdOutput = log_file_handle;
  startup_info.StartupInfo.hStdError  = log_file_handle;

  // Allocate an attribute list with space for 2 entries
  startup_info.lpAttributeList = AllocateProcThreadAttributeList(2);
  if(startup_info.lpAttributeList == NULL) {
    // Tell SCM we failed to start
    service_status.dwWin32ExitCode = GetLastError();
    service_status.dwCurrentState  = SERVICE_STOPPED;
    SetServiceStatus(service_status_handle, &service_status);
    return;
  }

  // Only allow Sunshine.exe to inherit the log file handle, not all inheritable handles
  UpdateProcThreadAttribute(startup_info.lpAttributeList,
    0,
    PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
    &log_file_handle,
    sizeof(log_file_handle),
    NULL,
    NULL);

  // Start Sunshine.exe inside our job object
  UpdateProcThreadAttribute(startup_info.lpAttributeList,
    0,
    PROC_THREAD_ATTRIBUTE_JOB_LIST,
    &job_handle,
    sizeof(job_handle),
    NULL,
    NULL);

  // Tell SCM we're running (and stoppable now)
  service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP;
  service_status.dwCurrentState     = SERVICE_RUNNING;
  SetServiceStatus(service_status_handle, &service_status);

  // Loop every 3 seconds until the stop event is set or Sunshine.exe is running
  while(WaitForSingleObject(stop_event, 3000) != WAIT_OBJECT_0) {
    auto console_token = DuplicateTokenForConsoleSession();
    if(console_token == NULL) {
      continue;
    }

    PROCESS_INFORMATION process_info;
    if(!CreateProcessAsUserW(console_token,
         L"Sunshine.exe",
         NULL,
         NULL,
         NULL,
         TRUE,
         ABOVE_NORMAL_PRIORITY_CLASS | CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW | EXTENDED_STARTUPINFO_PRESENT,
         NULL,
         NULL,
         (LPSTARTUPINFOW)&startup_info,
         &process_info)) {
      CloseHandle(console_token);
      continue;
    }

    // Close handles that are no longer needed
    CloseHandle(console_token);
    CloseHandle(process_info.hThread);

    // Wait for either the stop event to be set or Sunshine.exe to terminate
    const HANDLE wait_objects[] = { stop_event, process_info.hProcess };
    switch(WaitForMultipleObjects(_countof(wait_objects), wait_objects, FALSE, INFINITE)) {
    case WAIT_OBJECT_0:
      // The service is shutting down, so terminate Sunshine.exe.
      // TODO: Send a graceful exit request and only terminate forcefully as a last resort.
      TerminateProcess(process_info.hProcess, ERROR_PROCESS_ABORTED);
      break;

    case WAIT_OBJECT_0 + 1:
      // Sunshine terminated itself.
      break;
    }

    CloseHandle(process_info.hProcess);
  }

  // Let SCM know we've stopped
  service_status.dwCurrentState = SERVICE_STOPPED;
  SetServiceStatus(service_status_handle, &service_status);
}

int main(int argc, char *argv[]) {
  static const SERVICE_TABLE_ENTRY service_table[] = {
    { (LPSTR)SERVICE_NAME, ServiceMain },
    { NULL, NULL }
  };

  // By default, services have their current directory set to %SYSTEMROOT%\System32.
  // We want to use the directory where Sunshine.exe is located instead of system32.
  // This requires stripping off 2 path components: the file name and the last folder
  WCHAR module_path[MAX_PATH];
  GetModuleFileNameW(NULL, module_path, _countof(module_path));
  for(auto i = 0; i < 2; i++) {
    auto last_sep = wcsrchr(module_path, '\\');
    if(last_sep) {
      *last_sep = 0;
    }
  }
  SetCurrentDirectoryW(module_path);

  // Trigger our ServiceMain()
  return StartServiceCtrlDispatcher(service_table);
}