/*
 * Wrapper application to set octave env variables and then run octave
 *
 * Copyright (C) 2020-2021 John Donoghue <john.donoghue@ieee.org>
 *
 * This program is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation; either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program; if not, see <http: *www.gnu.org/licenses/>.
*/

#include <windows.h>
#include <shlwapi.h>
#include <strsafe.h>

static int ParentDir (wchar_t *dir)
{
  int len = lstrlenW (dir);
  while (len > 0 && dir[len] != L'\\')
    len --;
  dir[len] = 0;
  return len;
}

static wchar_t * FilePart (wchar_t *dir)
{
  int len = lstrlenW (dir);
  while (len > 0 && dir[len-1] != L'\\')
    len --;
  return &dir[len];
}

static size_t get_num_physical_cores (void)
{
  DWORD length;
  char *lpi;
  BOOL res;
  size_t num_physical_cores;
  size_t offset;

  length = 0;
  GetLogicalProcessorInformationEx (RelationProcessorCore, NULL, &length);
  if (GetLastError () != ERROR_INSUFFICIENT_BUFFER)
    return 0;

  lpi = malloc (length);
  res = GetLogicalProcessorInformationEx
          (RelationProcessorCore,
           (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) lpi,
           &length);
  if (! res)
    {
      free (lpi);
      return 0;
    }

  num_physical_cores = 0;
  offset = 0;
  do
    {
      const PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX cur_lpi =
        (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) (lpi + offset);
      offset += cur_lpi->Size;
      num_physical_cores++;
    }
  while (offset < length);

  free (lpi);

  return num_physical_cores;
}

// Set up environment and launch octave.exe with appropriate
// arguments.  NOTE: There are corresponding VBS and BAT script
// versions of this program so any changes here may need to be made in
// those scripts as well.

int wmain (int argc, wchar_t **argv)
{
  PROCESS_INFORMATION pi;

  // FIXME: There are currently have no checks to ensure that the
  // following fixed-size buffers are sufficiently large.  MAX_PATH is
  // 260 on Winndows by default.  But a user might have increased that
  // limit to 32767 characters by manually changing policy settings.
  // Maybe it would be better if we used C++ and std::string objects?
  // Note that the use of sizeof(path)-1 will fail if we switch to
  // using
  //
  //   wchar_t *path = (wchar_t *) malloc (...);
  //
  // instead of literal character arrays.

#define PATH_SZ 1024
  wchar_t msyspath[PATH_SZ];
  wchar_t rootpath[PATH_SZ];
  wchar_t path[PATH_SZ];
  wchar_t binpath[PATH_SZ];

  int gui_arg_found = 0;
  int no_gui_arg_found = 0;
  int no_gui_libs = 0;
  int i;

  /* get path of us which we will assume is the root of the install */

  DWORD nSize = GetModuleFileNameW (NULL, path, PATH_SZ-1);
  if (nSize)
    {
      while (nSize > 0 && path[nSize] != L'\\')
        nSize --;
      path[nSize] = L'\0';
    }

#ifdef NO_SHORT_PATH_NAMES
  StringCchCopyW (rootpath, PATH_SZ, path);
#else
  /* transform to short paths to work around issues with spaces in
     paths */
  /* FIXME: This won't help on systems with de-activated short paths */
  nSize = GetShortPathNameW (path, rootpath, PATH_SZ-1);
  if (nSize == 0)
    StringCchCopyW (rootpath, PATH_SZ, path);
#endif

  SetEnvironmentVariableW (L"MXE_ROOT", rootpath);

  /* get msys path and set env */

  StringCchCopyW (msyspath, PATH_SZ, rootpath);

  StringCchCopyW (path, PATH_SZ, rootpath);
  StringCchCatW (path, PATH_SZ, L"\\mingw64\\bin\\octave.exe");
  if (PathFileExistsW (path))
    {
      SetEnvironmentVariableW (L"MSYSTEM", L"MINGW64");
      StringCchCatW (msyspath, PATH_SZ, L"\\usr");
      SetEnvironmentVariableW (L"MSYSPATH", msyspath);

      StringCchCatW (msyspath, PATH_SZ, L"\\bin");

      StringCchCopyW (binpath, PATH_SZ, rootpath);
      StringCchCatW (binpath, PATH_SZ, L"\\mingw64\\bin");
    }
  else
    {
      StringCchCopyW (path, PATH_SZ, rootpath);
      StringCchCatW (path, PATH_SZ, L"\\mingw32\\bin\\octave.exe");
      if (PathFileExistsW (path))
        {
          SetEnvironmentVariableW (L"MSYSTEM", L"MINGW32");
          StringCchCatW (msyspath, PATH_SZ, L"\\usr");
          SetEnvironmentVariableW (L"MSYSPATH", msyspath);

          StringCchCatW (msyspath, PATH_SZ, L"\\bin");

          StringCchCopyW (binpath, PATH_SZ, rootpath);
          StringCchCatW (binpath, PATH_SZ, L"\\mingw32\\bin");
        }
      else
        {
          SetEnvironmentVariableW (L"MSYSTEM", L"MSYS");
          SetEnvironmentVariableW (L"MSYSPATH", msyspath);

          StringCchCatW (msyspath, PATH_SZ, L"\\bin");

          StringCchCopyW (binpath, PATH_SZ, rootpath);
          StringCchCatW (binpath, PATH_SZ, L"\\bin");
        }
    }

  /* binpath is to the octave bin dir so get the parent */
  StringCchCopyW (path, PATH_SZ, binpath);
  ParentDir (path);

#ifdef SHOW_PATHS
  MessageBoxW (NULL, rootpath, L"octroot", MB_OK);
  MessageBoxW (NULL, path, L"octhome", MB_OK);
  MessageBoxW (NULL, binpath, L"octbin", MB_OK);
  MessageBoxW (NULL, msyspath, L"msysbin", MB_OK);
#endif

  /* qt paths
   * either %OCT_HOME%\qt5\plugins
   * or %OCT_HOME\plugins
   */
  nSize = lstrlenW (path);
  StringCchCatW (path, PATH_SZ, L"\\qt5\\plugins");
  SetEnvironmentVariableW (L"QT_PLUGIN_PATH", path);

  path[nSize] = L'\0';

  /* exe paths */
  {
#define PATHBUF_SZ 8096
    wchar_t pathbuffer[PATHBUF_SZ];
    wchar_t newpathbuffer[PATHBUF_SZ];

    nSize = GetEnvironmentVariableW (L"PATH", pathbuffer, PATHBUF_SZ-1);
    pathbuffer[nSize] = '\0';

    /* set our paths first */
    StringCchCopyW (newpathbuffer, PATHBUF_SZ, binpath);
    StringCchCatW (newpathbuffer, PATHBUF_SZ, L";");
    StringCchCatW (newpathbuffer, PATHBUF_SZ, msyspath);
    StringCchCatW (newpathbuffer, PATHBUF_SZ, L";");
    StringCchCatW (newpathbuffer, PATHBUF_SZ, pathbuffer);

    SetEnvironmentVariableW (L"PATH", newpathbuffer);
  }

  /* other env */
  SetEnvironmentVariableW (L"TERM", L"cygwin");
  SetEnvironmentVariableW (L"GNUTERM", L"wxt");
  SetEnvironmentVariableW (L"GS", L"gs.exe");

  /* home set */
  nSize = GetEnvironmentVariableW (L"HOME", path, PATH_SZ-1);
  if (nSize == 0 || path[0] == 0)
    {
      wchar_t newhome[PATH_SZ];

      nSize = GetEnvironmentVariableW (L"USERPROFILE", path, PATH_SZ-1);
      if (nSize == 0 || path[0] == 0)
        {
          /* build dir from drive and path */
          wchar_t tmpbuff[PATH_SZ];
          nSize = GetEnvironmentVariableW (L"HOMEDRIVE", tmpbuff, PATH_SZ-1);
          if (nSize)
            StringCchCopyW (path, PATH_SZ, tmpbuff);
          else
            path[0] = '\0';

          nSize = GetEnvironmentVariableW (L"HOMEPATH", tmpbuff, PATH_SZ-1);
          if (nSize)
            StringCchCopyW (path, PATH_SZ, tmpbuff);
        }

#ifdef NO_SHORT_PATH_NAMES
      StringCchCopyW (newhome, PATH_SZ, path);
#else
      /* transform to short paths to work around issues with spaces in
         paths */
      /* FIXME: This won't help on systems with de-activated short
         paths */
      nSize = GetShortPathNameW (path, newhome, PATH_SZ-1);

      if (nSize == 0)
        StringCchCopyW (newhome, PATH_SZ, path);
#endif

      SetEnvironmentVariableW (L"HOME", newhome);
    }

  /* set number of OpenBLAS threads */
  nSize = GetEnvironmentVariableW (L"OPENBLAS_NUM_THREADS", NULL, 0);
  if (nSize == 0)
    {
      /* Only set if it wasn't already set in the environment */
      size_t num_threads;
      num_threads = get_num_physical_cores ();

      // Setting OPENBLAS_NUM_THREADS to something higher than what NUM_THREADS
      // was set to when OpenBLAS was configured, can lead to errors.
      // FIXME: Can/should we get that number from the library?
#define MAX_NUM_THREADS 24
      if (num_threads > MAX_NUM_THREADS)
        num_threads = MAX_NUM_THREADS;

      if (num_threads > 0)
        {
#define THREADS_SZ 64
          wchar_t buffer[THREADS_SZ];
          StringCchPrintfW (buffer, THREADS_SZ, L"%zu", num_threads);
          SetEnvironmentVariableW (L"OPENBLAS_NUM_THREADS", buffer);
        }
    }

  /* check for gui mode */
  for (i = 1; i < argc; i++)
    {
      if (StrCmpW (argv[i], L"--no-gui-libs") == 0)
        {
          if (gui_arg_found)
            {
              /* Inconsistent options.  We can't start the GUI without gui libs.
                 How should we fail?   For now, --no-gui-libs will
                 override the other options.  */
            }

          no_gui_libs = 1;
        }
      else if (StrCmpW (argv[i], L"--gui") == 0
               || StrCmpW (argv[i], L"--force-gui") == 0)
        {
          if (no_gui_libs)
            {
              /* Inconsistent options.  We can't start the GUI without gui libs.
                 How should we fail?   For now, --no-gui-libs will
                 override the other options.  */
            }

          gui_arg_found = 1;
        }
      else if (StrCmpW (argv[i], L"--no-gui") == 0)
        no_gui_arg_found = 1;

      /* NOTE: specifying both --no-gui and --gui is also an
         inconsistent set of options but we leave it to octave.exe to
         detect that.  */
    }

#ifdef FIRST_TIME
  {
    /* change directory to USERPROFILE before starting Octave */
    wchar_t tmpbuff[PATH_SZ];
    nSize = GetEnvironmentVariableW (L"USERPROFILE", tmpbuff, PATH_SZ-1);
    if (nSize)
      StringCchCopyW (path, PATH_SZ, tmpbuff);

    SetCurrentDirectoryW (path);
  }
#endif

  /* set up process args and start it */
  {
    STARTUPINFO si;
#define ARGBUF_SZ 4096
    wchar_t argbuffer[ARGBUF_SZ];

    ZeroMemory (&si, sizeof (si));
    si.cb = sizeof (si);
    ZeroMemory (&pi, sizeof (pi));

    StringCchCopyW (path, PATH_SZ, binpath);

    StringCchCopyW (argbuffer, ARGBUF_SZ, L"octave.exe ");
    StringCchCatW (path, PATH_SZ, L"\\octave.exe");

    /* If parent process has a console, attach to it.
       Let the function fail silently, when parent has no console
       (e.g., when program has been started from link in start menu).
       No console will be shown in this case. */
    AttachConsole (ATTACH_PARENT_PROCESS);

    if (! (no_gui_libs || no_gui_arg_found))
      {
        /* Unless --no-gui or --no-gui-libs is specified, we will use a GUI window.  */
        si.dwFlags = STARTF_USESHOWWINDOW;

        /* If none of the options --no-gui, --gui, or --force-gui
           were specified, then we'll add --gui to start the gui as
           most Windows users would expect.  */

        if (! gui_arg_found)
          StringCchCatW (argbuffer, ARGBUF_SZ, L"--gui ");
      }

    /* quote and append each arg */
    for (i = 1; i < argc; i++)
      {
        StringCchCatW (argbuffer, ARGBUF_SZ, L"\"");
        StringCchCatW (argbuffer, ARGBUF_SZ, argv[i]);
        StringCchCatW (argbuffer, ARGBUF_SZ, L"\" ");
      }

    /* Start the child process */
    int status = CreateProcessW (path,      // Module name
                                 argbuffer, // Command line
                                 NULL,      // Process handle not inheritable
                                 NULL,      // Thread handle not inheritable
                                 FALSE,     // Set handle inheritance to FALSE
                                 0,         // No creation flags
                                 NULL,      // Use parent's environment block
                                 NULL,      // Use parent's starting directory
                                 &si,       // Pointer to STARTUPINFO
                                 &pi);      // Pointer to PROCESS_INFORMATION

    if (! status)
      return 1;
  }

  /* Wait until child process exits. */
  WaitForSingleObject (pi.hProcess, INFINITE);

  /* Close process and thread handles */
  CloseHandle (pi.hProcess);
  CloseHandle (pi.hThread);

  return 0;
}