/* 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;
    }
}

}

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);
}
