#include <stdio.h>
#include <assert.h>
#include <windows.h>

static int update_start(const char *path)
{
    HANDLE hProcess = NULL;
    STARTUPINFO startupInfo = {0};
    PROCESS_INFORMATION processInfo;

    char module[256];
    char cmdline[8192];
    startupInfo.cb = sizeof(startupInfo);
    hProcess = OpenProcess(SYNCHRONIZE, TRUE, GetCurrentProcessId());
    size_t count = GetModuleFileName(NULL, module, sizeof(module));
    assert(count < sizeof(module));
    snprintf(cmdline, sizeof(cmdline),
            "update.exe %p -wait %s -delete %s -copy %s -final",
            hProcess, module, module, module);
    if (CreateProcess(path, cmdline, NULL, NULL, TRUE, NORMAL_PRIORITY_CLASS,
                NULL, NULL, &startupInfo, &processInfo) != TRUE) {
        printf("Create Process Fail: %d!\n", GetLastError());
        return -1;
    }
    CloseHandle(processInfo.hThread);
    CloseHandle(processInfo.hProcess);
    CloseHandle(hProcess);
    return 0;
}

static int update_wait(const char *handle)
{
    HANDLE hProcess = NULL;
    int count = sscanf(handle, "%p", &hProcess);
    assert(count == 1);
    WaitForSingleObject(hProcess, INFINITE);
    CloseHandle(hProcess);
    return 0;
}

static int update_final(const char *path)
{
    HANDLE hProcess = NULL;
    STARTUPINFO startupInfo = {0};
    PROCESS_INFORMATION processInfo;

    char module[256];
    char cmdline[8192];
    startupInfo.cb = sizeof(startupInfo);
    hProcess = OpenProcess(SYNCHRONIZE, TRUE, GetCurrentProcessId());
    size_t count = GetModuleFileName(NULL, module, sizeof(module));
    assert(count < sizeof(module));
    snprintf(cmdline, sizeof(cmdline), "update.exe %p -wait %s -delete -1 -clear",
            hProcess, module, module);
    if (CreateProcess(path, cmdline, NULL, NULL, TRUE, NORMAL_PRIORITY_CLASS,
                NULL, NULL, &startupInfo, &processInfo) != TRUE) {
        CloseHandle(hProcess);
        return -1;
    }
    CloseHandle(processInfo.hThread);
    CloseHandle(processInfo.hProcess);
    CloseHandle(hProcess);
    return 0;
}

static int update_interpret(int argc, char *argv[])
{
    int i;
    char module[8192];
    size_t count = GetModuleFileName(NULL, module, sizeof(module));
    assert (count < sizeof(module));

    for (i=1; i<argc; i++) {
        if (!strcmp(argv[i], "-wait")) {
            if (update_wait(argv[i-1]) != 0)
                printf("update wait fail: %d\n", GetLastError());
        }else if (!strcmp(argv[i], "-copy")){
            if (CopyFile(module, argv[i-1], TRUE) != TRUE)
                printf("copy %s to %s fail\n", module, argv[i-1]);
        }else if (!strcmp(argv[i], "-delete")){
            if (DeleteFile(argv[i-1]) != TRUE)
                printf("delete %s fail\n", argv[i-1]);
        }else if (!strcmp(argv[i], "-final")){
            if (update_final(argv[i-1]) != 0)
                printf("update final fail: %d\n", GetLastError());
        }else if (!strcmp(argv[i], "-clear")){
            return atoi(argv[i-1]);
        }
    }
    return 0;
}

static int autorun_set(const char *regpath)
{
    HKEY hAutoRunKey = NULL;
    if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, regpath, 0,
               KEY_ALL_ACCESS, &hAutoRunKey) != ERROR_SUCCESS){
    printf("RegOpenKeyEx Fail\n");
    return -1;
    }
    char module[8192];
    size_t count = GetModuleFileName(NULL, module, sizeof(module));
    assert(count < sizeof(module));
    RegSetValueEx(hAutoRunKey, "rpcserver", 0, REG_SZ,
        (LPBYTE)module, strlen(module));
    RegCloseKey(hAutoRunKey);
    return 0;
}

int main(int argc, char *argv[])
{
    //FreeConsole();

    if (argc > 1 ) {
        if (!update_interpret(argc, argv))
            return 0;
    }else{
        if (!update_start("update.exe"))
            return 0;
    }

    autorun_set("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run");

    HANDLE hMutex = CreateMutex(NULL, FALSE, "rg4.net/rpcserver");
    if (GetLastError() == ERROR_ALREADY_EXISTS){
    CloseHandle(hMutex);
    return -1;
    }

    for (;;) {
        Sleep(1000);
        printf("Hello World!\n");
    }

    CloseHandle(hMutex);
    return 0;
}