/* Copyright (C) 2014 InfiniDB, Inc.

   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; version 2 of
   the License.

   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, write to the Free Software
   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
   MA 02110-1301, USA. */

#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#include <conio.h>
#include <malloc.h>
#include <unistd.h>
#include <io.h>
#include <direct.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <psapi.h>

#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <sstream>
//#define NDEBUG
#include <cassert>
#include <unistd.h>
using namespace std;

#include <boost/algorithm/string/predicate.hpp>
namespace ba = boost::algorithm;

#include "idbregistry.h"
#include "syslog.h"

namespace
{
bool vFlg;
bool useSCMAPI;

SERVICE_STATUS_HANDLE svcStatusHandle;
SERVICE_STATUS svcStatus;
HANDLE svcStopEvent;

struct ProcInfo
{
  ProcInfo()
  {
  }
  ProcInfo(const string& name) : pName(name)
  {
  }
  string pName;
  string pCmdLine;
};

typedef vector<ProcInfo> ProcInfoVec;

const int interProcessSleepTime = 2;
const size_t cmdLineLen = 128;

string installDir;

bool shuttingDown = false;

int runIt(const string& pName)
{
  int rc = 0;
  char* cmdLine = (char*)alloca(cmdLineLen);
  strncpy_s(cmdLine, cmdLineLen, pName.c_str(), pName.size());
  PROCESS_INFORMATION pInfo;
  ZeroMemory(&pInfo, sizeof(pInfo));
  STARTUPINFO sInfo;
  ZeroMemory(&sInfo, sizeof(sInfo));

  try
  {
    if (CreateProcess(0, cmdLine, 0, 0, false, 0, 0, 0, &sInfo, &pInfo) == 0)
      return -1;
  }
  catch (exception& e)
  {
    cout << e.what() << endl;
  }

  if (WaitForSingleObject(pInfo.hProcess, INFINITE) != WAIT_OBJECT_0)
  {
    rc = -1;
    goto out;
  }

  DWORD exitCode;

  if (GetExitCodeProcess(pInfo.hProcess, &exitCode) == 0)
  {
    rc = -1;
    goto out;
  }

  if (exitCode != 0)
    rc = -1;

out:
  CloseHandle(pInfo.hProcess);
  return rc;
}

int loadBRM()
{
  // if no save file just return
  string saveFile = installDir + "\\dbrm\\BRM_saves_current";

  if (GetFileAttributes(saveFile.c_str()) == INVALID_FILE_ATTRIBUTES)
    return 0;

  // read contents of save file
  ifstream ifs(saveFile.c_str());

  if (!ifs)
    return -1;

  string saveFilePfx;
  getline(ifs, saveFilePfx);

  if (saveFilePfx.empty())
    return 0;

  // run load_brm and wait for it to finish
  string brmCmd = "load_brm \"" + saveFilePfx + "\"";

  if (runIt(brmCmd))
    return -1;

  return 0;
}

void ReportSvcStatus(DWORD currentState, DWORD win32ExitCode, DWORD waitHint)
{
  static DWORD checkPoint = 1;

  svcStatus.dwCurrentState = currentState;
  svcStatus.dwWin32ExitCode = win32ExitCode;
  svcStatus.dwWaitHint = waitHint;

  if (currentState == SERVICE_START_PENDING)
    svcStatus.dwControlsAccepted = 0;
  else
    svcStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;

  if (currentState == SERVICE_RUNNING || currentState == SERVICE_STOPPED)
    svcStatus.dwCheckPoint = 0;
  else
    svcStatus.dwCheckPoint = checkPoint++;

  SetServiceStatus(svcStatusHandle, &svcStatus);
}

extern "C" DWORD WINAPI procRunner(LPVOID);

int startUp()
{
  int rc;

  syslog(LOG_INFO, "System is starting");

  if (runIt("clearShm"))
    return -1;

  rc = loadBRM();

  if (rc)
    return rc;

  ProcInfoVec procInfo;
  procInfo.push_back(ProcInfo("workernode DBRM_Worker1 fg"));
  procInfo.push_back(ProcInfo("controllernode fg"));
  procInfo.push_back(ProcInfo("DecomSvr"));
  procInfo.push_back(ProcInfo("PrimProc"));
  procInfo.push_back(ProcInfo("WriteEngineServer"));
  procInfo.push_back(ProcInfo("ExeMgr"));
  procInfo.push_back(ProcInfo("DDLProc"));
  procInfo.push_back(ProcInfo("DMLProc"));
  string mysqldCmd = "mysqld --defaults-file=" + installDir + "\\my.ini";
  procInfo.push_back(ProcInfo(mysqldCmd));

  const ProcInfoVec::size_type numProcs = procInfo.size();

  for (unsigned pidx = 0; pidx < numProcs; pidx++)
  {
    cout << "Starting " << procInfo[pidx].pName << "...";

    HANDLE thd;
    DWORD tid;
    thd = CreateThread(0, 0, procRunner, &procInfo[pidx], 0, &tid);

    if (thd == NULL)
    {
      LPTSTR lpBuffer;
      LPTSTR* lppBuffer = &lpBuffer;
      DWORD fmRes = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0,
                                  GetLastError(), 0, (LPSTR)lppBuffer, 0, 0);

      cerr << endl;
      ostringstream ostr;
      ostr << "Failed to start process runner thread for " << procInfo[pidx].pName << ": ";

      if (fmRes > 0)
        ostr << lpBuffer;
      else
        ostr << "Unknown error";

      cerr << ostr.str() << endl;
      syslog(LOG_ERR, ostr.str().c_str());
      return -1;
    }

    Sleep(interProcessSleepTime * 1000);
    cout << endl;

    if (useSCMAPI)
      ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR,
                      interProcessSleepTime * (static_cast<DWORD>(numProcs) - pidx - 1) * 1000);
  }

  return 0;
}

DWORD WINAPI procRunner(LPVOID parms)
{
  ProcInfo* pip = reinterpret_cast<ProcInfo*>(parms);
  ProcInfo pi = *pip;
  char* cmdLine = (char*)alloca(cmdLineLen);
  strncpy_s(cmdLine, cmdLineLen, pi.pName.c_str(), pi.pName.size());
  BOOL cpRc;
  PROCESS_INFORMATION pInfo;
  STARTUPINFO sInfo;

  for (;;)
  {
    ZeroMemory(&sInfo, sizeof(sInfo));
    ZeroMemory(&pInfo, sizeof(pInfo));
    cpRc = CreateProcess(0, cmdLine, 0, 0, false, 0, 0, 0, &sInfo, &pInfo);

    if (cpRc == 0)
    {
      LPTSTR lpBuffer;
      LPTSTR* lppBuffer = &lpBuffer;
      DWORD fmRes = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0,
                                  GetLastError(), 0, (LPSTR)lppBuffer, 0, 0);
      ostringstream ostr;
      cerr << endl;
      ostr << "Failed to start process " << pi.pName << ": ";

      if (fmRes > 0)
        ostr << lpBuffer;
      else
        ostr << "Unknown error";

      cerr << ostr.str() << endl;
      syslog(LOG_ERR, ostr.str().c_str());
      return -1;
    }
    else
    {
      ostringstream ostr;
      ostr << pi.pName << " started ";
      cerr << ostr.str() << endl;
      syslog(LOG_INFO, ostr.str().c_str());
    }

    WaitForSingleObject(pInfo.hProcess, INFINITE);
    ostringstream ostr;
    ostr << pi.pName;

    if (shuttingDown)
    {
      ostr << " shut down";
      syslog(LOG_INFO, ostr.str().c_str());
    }
    else
    {
      ostr << " has died unexpectedly";
      syslog(LOG_ERR, ostr.str().c_str());
    }

    cerr << ostr.str() << endl;
    CloseHandle(pInfo.hProcess);

    if (shuttingDown)
      return 0;

    Sleep(10 * 1000);
  }

  return -1;
}

int killProcByPid(DWORD pid)
{
  int rc = -1;
  HANDLE hProc;
  hProc = OpenProcess(PROCESS_TERMINATE, false, pid);

  if (hProc != NULL)
  {
    if (TerminateProcess(hProc, 0) != 0)
      rc = 0;

    CloseHandle(hProc);
  }

  return rc;
}

DWORD maxPids = 64;
DWORD* pids = 0;

int killProcByName(const string& pname)
{
  if (!pids)
    pids = (DWORD*)malloc(maxPids * sizeof(DWORD));

  DWORD needed = 0;

  if (EnumProcesses(pids, maxPids * sizeof(DWORD), &needed) == 0)
    return -1;

  while (needed == maxPids * sizeof(DWORD))
  {
    maxPids *= 2;
    pids = (DWORD*)realloc(pids, maxPids * sizeof(DWORD));

    if (EnumProcesses(pids, maxPids * sizeof(DWORD), &needed) == 0)
      return -1;
  }

  int rc = -1;
  DWORD numPids = needed / sizeof(DWORD);
  DWORD i;

  for (i = 0; i < numPids; i++)
  {
    bool found = false;

    if (pids[i] != 0)
    {
      TCHAR szProcessName[MAX_PATH] = TEXT("<unknown>");

      // Get a handle to the process.
      HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pids[i]);

      // Get the process name.
      if (NULL != hProcess)
      {
        HMODULE hMod;
        DWORD cbNeeded;

        if (EnumProcessModules(hProcess, &hMod, sizeof(hMod), &cbNeeded))
        {
          GetModuleBaseName(hProcess, hMod, szProcessName, sizeof(szProcessName) / sizeof(TCHAR));
        }
      }

      if (pname == szProcessName)
        found = true;

      CloseHandle(hProcess);
    }

    if (found)
    {
      rc = killProcByPid(pids[i]);
      break;
    }
  }

  return rc;
}

int shutDown()
{
  shuttingDown = true;
  syslog(LOG_INFO, "System is shutting down");

  vector<string> pList;

  pList.push_back("mysqld.exe");
  pList.push_back("DMLProc.exe");
  pList.push_back("DDLProc.exe");
  pList.push_back("ExeMgr.exe");
  pList.push_back("WriteEngineServer.exe");
  pList.push_back("PrimProc.exe");
  pList.push_back("DecomSvr.exe");
  pList.push_back("controllernode.exe");
  pList.push_back("workernode.exe");

  vector<string>::iterator iter = pList.begin();
  vector<string>::iterator end = pList.end();
  vector<string>::size_type i = pList.size();

  while (iter != end)
  {
    if (killProcByName(*iter) == 0)
      Sleep(interProcessSleepTime * 1000);

    if (useSCMAPI)
      ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, static_cast<DWORD>(--i) * interProcessSleepTime * 1000);

    ++iter;
  }

  if (runIt("save_brm") == 0)
    runIt("clearShm");

  if (useSCMAPI)
    ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 0);

  return 0;
}

void svcStop()
{
  useSCMAPI = true;
  ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 20 * 1000);
  SetEvent(svcStopEvent);
}

extern "C" VOID WINAPI svcMain(DWORD, LPSTR*);

// The SCM has asked us to start as a service
int svcStart()
{
  SERVICE_TABLE_ENTRY st[] = {
      {"InfiniDB", svcMain},
      {0, 0},
  };

  // This call returns when the service has stopped
  StartServiceCtrlDispatcher(st);

  return 0;
}

extern "C" VOID WINAPI svcCtrlHandler(DWORD);

VOID WINAPI svcMain(DWORD dwArgc, LPSTR* lpszArgv)
{
  useSCMAPI = true;

  svcStatusHandle = RegisterServiceCtrlHandler("InfiniDB", svcCtrlHandler);

  svcStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
  svcStatus.dwServiceSpecificExitCode = 0;

  ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, interProcessSleepTime * 7 * 1000);

  svcStopEvent = CreateEvent(0, 1, 0, 0);

  if (startUp() != 0)
  {
    ReportSvcStatus(SERVICE_STOPPED, 1, 0);
    return;
  }

  ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);

  WaitForSingleObject(svcStopEvent, INFINITE);

  shutDown();

  CloseHandle(svcStopEvent);

  ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
}

VOID WINAPI svcCtrlHandler(DWORD dwCtrl)
{
  // Handle the requested control code.

  switch (dwCtrl)
  {
    case SERVICE_CONTROL_STOP:      // Notifies a service that it should stop
    case SERVICE_CONTROL_SHUTDOWN:  // Notifies a service that the system is shutting down so the service can
                                    // perform cleanup tasks
      svcStop();
      break;

    case SERVICE_CONTROL_INTERROGATE:  // Notifies a service that it should report its current status
                                       // information to the service control manager
      break;

    case SERVICE_CONTROL_CONTINUE:  // Notifies a paused service that it should resume
      break;

    case SERVICE_CONTROL_PARAMCHANGE:  // Notifies a service that its startup parameters have changed. The
                                       // service should reread its startup parameters
      break;

    case SERVICE_CONTROL_PAUSE:  // Notifies a service that it should pause
      break;

    default: break;
  }
}

}  // namespace

int main(int argc, char** argv)
{
  opterr = 0;
  vFlg = false;
  int c;
  int rc;
  useSCMAPI = false;

  while ((c = getopt(argc, argv, "v")) != EOF)
    switch (c)
    {
      case 'v': vFlg = true; break;

      case '?':
      default: break;
    }

  if (!vFlg)
  {
    _close(0);
    _close(1);
    _close(2);
    int fd = -1;
    _sopen_s(&fd, "NUL", _O_TEXT | _O_RDONLY, _SH_DENYNO, _S_IREAD);
    assert(fd == 0);
    fd = -1;
    _sopen_s(&fd, "NUL", _O_TEXT | _O_APPEND | _O_WRONLY, _SH_DENYNO, _S_IWRITE);
    assert(fd == 1);
    fd = -1;
    _sopen_s(&fd, "NUL", _O_TEXT | _O_APPEND | _O_WRONLY, _SH_DENYNO, _S_IWRITE);
    assert(fd == 2);
  }

  int (*cmdFp)();

  string command;

  if (argc - optind < 1)
    command = "svcstart";
  else
    command = argv[optind];

  if (ba::istarts_with(command, "sta"))
  {
    cmdFp = startUp;
  }
  else if (ba::istarts_with(command, "sto") || ba::istarts_with(command, "sh"))
  {
    cmdFp = shutDown;
  }
  else if (ba::istarts_with(command, "sv"))
  {
    cmdFp = svcStart;
  }
  else
  {
    cerr << "Unknown command " << command << " (try 'start' or 'shutdown')" << endl;
    return 1;
  }

  installDir = IDBreadRegistry("", true);

  string newDir = installDir + "\\bin";
  rc = _chdir(newDir.c_str());

  rc = cmdFp();

  return (rc ? 1 : 0);
}