#include <Windows.h>
#include <string>
#include <stdexcept>
#include <filesystem>
#include <Psapi.h>

void ShowError(DWORD);
void ShowLastError();

using std::filesystem::path;

int __stdcall wWinMain(
    HINSTANCE hInstance,
    HINSTANCE hPrevInstance,
    LPWSTR     lpCmdLine,
    int       nShowCmd
)
{
    int numArgs = 0;
    LPWSTR* args = CommandLineToArgvW(lpCmdLine, &numArgs);
    if (numArgs < 1 || std::wcslen(lpCmdLine) == 0)
    {
        MessageBoxA(NULL, "No PID", "Error", MB_OK);
        return -1;
    }

    int pid = -1;
    try
    {
        pid = std::stoi(args[0]);
    }
    catch (std::invalid_argument e)
    {
        char msg[256];
        sprintf_s(msg, "Invalid process id: %s", e.what());
        MessageBoxA(NULL, msg, "Error", MB_OK);
        return -1;
    }

    CreateProcess(L"%localappdata%/wsltty/bin/mintty.exe", L"", NULL, NULL, FALSE, 0, NULL, NULL,)

    HANDLE process = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION, FALSE, pid);
    if (process == NULL)
    {
        ShowLastError();
        return -1;
    }

    HMODULE kernel = GetModuleHandle(L"kernel32.dll");
    if (kernel == NULL)
    {
        ShowLastError();
        return -1;
    }
    
    path lib;
    try
    {
        lib = std::filesystem::canonical(path("InjectAppUserModelId_x64.dll"));
    }
    catch (std::filesystem::filesystem_error e)
    {
        ShowError(e.code().value());
        return -1;
    }
    const std::string s = lib.string();
    char* libCstr = new char[s.length() + 1];
    memcpy(libCstr, s.c_str(), s.length() + 1);

    FARPROC loadLibraryAddr = GetProcAddress(kernel, "LoadLibraryA");
    if (loadLibraryAddr == NULL)
    {
        ShowLastError();
        return -1;
    }
    
    HMODULE injected = LoadLibrary(lib.c_str());
    if (injected == NULL)
    {
        ShowLastError();
        return -1;
    }

    FARPROC targetAddr = GetProcAddress(injected, "DoWork");
    if (targetAddr == NULL)
    {
        ShowLastError();
        return -1;
    }

    INT_PTR offset = (INT_PTR)targetAddr - (INT_PTR)injected;

    void* remoteMem = VirtualAllocEx(process, NULL, strlen(libCstr) + 1, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    if (!WriteProcessMemory(process, remoteMem, libCstr, strlen(libCstr) + 1, NULL))
    {
        ShowLastError();
        return -1;
    }

    HANDLE thread = CreateRemoteThread(process, NULL, 0, (LPTHREAD_START_ROUTINE)loadLibraryAddr, remoteMem, 0, NULL);
    WaitForSingleObject(thread, INFINITE);

    DWORD ret = 0;
    GetExitCodeThread(thread, &ret);

    if (ret == NULL)
    {
        ShowError(ret);
        return -1;
    }

    HMODULE* remoteModuleList = new HMODULE[1000];
    DWORD needed = 0;
    EnumProcessModules(process, remoteModuleList, 1000, &needed);
    int numModules = needed / sizeof(HMODULE);

    bool found = false;

    wchar_t* name = new wchar_t[1024];
    for (int i = 0; i < numModules; i++)
    {
        GetModuleFileNameEx(process, remoteModuleList[i], name, 1023);
        std::wstring modName(name);

        if (modName.find(L"Inject") != std::wstring::npos)
        {
            const char* msg = "Message from here";
            void* remoteMem = VirtualAllocEx(process, NULL, strlen(msg) + 1, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
            if (!WriteProcessMemory(process, remoteMem, msg, strlen(msg) + 1, NULL))
            {
                ShowLastError();
                return -1;
            }

            HANDLE thread = CreateRemoteThread(process, NULL, 0, (LPTHREAD_START_ROUTINE)((INT_PTR)remoteModuleList[i] + offset), remoteMem, 0, NULL);
            WaitForSingleObject(thread, INFINITE);

            DWORD ret = 0;
            GetExitCodeThread(thread, &ret);

            if (ret != 0)
            {
                char err[128];
                sprintf_s(err, "Error from injected function: %d", ret);
                MessageBoxA(NULL, err, "Error", MB_OK);
                return -1;
            }

            found = true;
            break;
        }
    }

    if (!found)
    {
        MessageBox(NULL, L"Couldn't find injected dll", L"Error", MB_OK);
        return -1;
    }

    MessageBoxA(NULL, "Success!", "OK", MB_OK);

    return 0;
}

void ShowLastError()
{
    ShowError(GetLastError());
}

void ShowError(DWORD error)
{
    LPTSTR err;
    FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, NULL, error, 0, (LPTSTR)&err, 0, NULL);
    MessageBoxW(NULL, err, L"Error", MB_OK);
}