#include <string>
#include <thread>
#include <atomic>
#include <cassert>
#include <iostream>
#include <mutex>
#include <windows.h>
#include <io.h>
#include <winternl.h>
#define isatty                       _isatty
#define fileno                       _fileno
#define NtCurrentProcess()           ((HANDLE)(LONG_PTR) - 1)
#define RESIZE_CONHOST_SIGNAL_BUFFER 8
#ifndef PSEUDOCONSOLE_INHERIT_CURSOR
#    define PSEUDOCONSOLE_INHERIT_CURSOR 1
#endif
#define BUFF_SIZE 0x200

typedef struct _HPCON_INTERNAL
{
    HANDLE hWritePipe;
    HANDLE hConDrvReference;
    HANDLE hConHostProcess;
} HPCON_INTERNAL;

inline NTSTATUS WINAPI CreateHandle(PHANDLE FileHandle, PWSTR Buffer,
                                    ACCESS_MASK DesiredAccess,
                                    HANDLE      RootDirectory,
                                    BOOLEAN     IsInheritable,
                                    ULONG       OpenOptions)
{
    IO_STATUS_BLOCK IoStatusBlock;
    ULONG           Attributes = OBJ_CASE_INSENSITIVE;
    if (IsInheritable)
        Attributes |= OBJ_INHERIT;

    UNICODE_STRING ObjectName;
    memset(&ObjectName, 0, sizeof ObjectName);
    RtlInitUnicodeString(&ObjectName, Buffer);

    OBJECT_ATTRIBUTES ObjectAttributes;
    memset(&ObjectAttributes, 0, sizeof ObjectAttributes);
    ObjectAttributes.RootDirectory = RootDirectory;
    ObjectAttributes.Length = sizeof ObjectAttributes;
    ObjectAttributes.Attributes = Attributes;
    ObjectAttributes.ObjectName = &ObjectName;

    return NtOpenFile(
        FileHandle, DesiredAccess, &ObjectAttributes, &IoStatusBlock,
        FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
        OpenOptions);
}

class Terminal
{
  public:
    HPCON consoleHandle;

    HRESULT
    WINAPI
    InitializeConsole(COORD size, HANDLE hInput, HANDLE hOutput,
                      DWORD dwFlags)
    {
        NTSTATUS Status;
        BOOL     bRes;

        HANDLE InputHandle, OutputHandle;
        HANDLE hConServer, hConReference;
        HANDLE hProc = NtCurrentProcess();

        bRes = DuplicateHandle(hProc, hInput, hProc, &InputHandle, 0,
                               TRUE, DUPLICATE_SAME_ACCESS);
        assert(bRes != 0);

        bRes = DuplicateHandle(hProc, hOutput, hProc, &OutputHandle,
                               0, TRUE, DUPLICATE_SAME_ACCESS);
        assert(bRes != 0);

        Status = CreateHandle(&hConServer,
                              (LPWSTR)L"\\Device\\ConDrv\\Server",
                              GENERIC_ALL, NULL, TRUE, 0);

        std::cout << "Status: " << Status << '\n';

        assert(Status == 0);

        HANDLE              ReadPipeHandle, WritePipeHandle;
        SECURITY_ATTRIBUTES PipeAttributes = {};
        PipeAttributes.bInheritHandle = TRUE;
        PipeAttributes.lpSecurityDescriptor = NULL;
        PipeAttributes.nLength = sizeof PipeAttributes;

        bRes = CreatePipe(&ReadPipeHandle, &WritePipeHandle,
                          &PipeAttributes, 0);
        assert(bRes != 0);

        bRes = SetHandleInformation(
            ReadPipeHandle, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);
        assert(bRes != 0);

        PWSTR InheritCursor =
            (LPWSTR)L"--inheritcursor "; /* Requires one space */
        if (!(dwFlags & PSEUDOCONSOLE_INHERIT_CURSOR))
            InheritCursor = (LPWSTR)L"";

        wchar_t ConHostCommand[MAX_PATH];

        /* ConHost command format strings */
#define COMMAND_FORMAT                                          \
    L"\\\\?\\%s\\system32\\conhost.exe %s--width %hu --height " \
    L"%hu --signal 0x%x --server 0x%x"

        wchar_t winDir[MAX_PATH];
        GetWindowsDirectoryW(winDir, MAX_PATH);
        swprintf_s(ConHostCommand, MAX_PATH,
                   L"\"%s\\system32\\conhost.exe\" %s--width %hu "
                   L"--height %hu --signal 0x%x --server 0x%x",
                   winDir, InheritCursor, size.X, size.Y,
                   HandleToULong(ReadPipeHandle),
                   HandleToULong(hConServer));

        /* Initialize thread attribute list */
        HANDLE Values[4] = {hConServer, InputHandle, OutputHandle,
                            ReadPipeHandle};
        size_t AttrSize;
        LPPROC_THREAD_ATTRIBUTE_LIST AttrList = NULL;
        InitializeProcThreadAttributeList(NULL, 1, 0, &AttrSize);
        AttrList = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(
            GetProcessHeap(), 0, AttrSize);
        InitializeProcThreadAttributeList(AttrList, 1, 0, &AttrSize);
        bRes = UpdateProcThreadAttribute(
            AttrList, 0,
            PROC_THREAD_ATTRIBUTE_HANDLE_LIST, /* 0x20002u */
            Values, sizeof Values, NULL, NULL);
        assert(bRes != 0);

        /* Assign members of STARTUPINFOEXW */
        PROCESS_INFORMATION ProcInfo;
        memset(&ProcInfo, 0, sizeof ProcInfo);
        STARTUPINFOEXW SInfoEx;
        memset(&SInfoEx, 0, sizeof SInfoEx);
        SInfoEx.StartupInfo.cb = sizeof SInfoEx;
        SInfoEx.StartupInfo.dwFlags |= STARTF_USESTDHANDLES;
        SInfoEx.StartupInfo.hStdInput = InputHandle;
        SInfoEx.StartupInfo.hStdOutput = OutputHandle;
        SInfoEx.StartupInfo.hStdError = OutputHandle;
        SInfoEx.lpAttributeList = AttrList;

        bRes = CreateProcessW(
            NULL, ConHostCommand, NULL, NULL, TRUE,
            EXTENDED_STARTUPINFO_PRESENT, NULL, NULL,
            &SInfoEx.StartupInfo, &ProcInfo);
        assert(bRes != 0);

        Status = CreateHandle(
            &hConReference, (LPWSTR)L"\\Reference",
            GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE, hConServer,
            FALSE, FILE_SYNCHRONOUS_IO_NONALERT);
        assert(Status == 0);

        /* Return handles to caller */
        HPCON_INTERNAL* consoleHandleInt = (HPCON_INTERNAL*)HeapAlloc(
            GetProcessHeap(), 0, sizeof(*consoleHandleInt));
        consoleHandleInt->hWritePipe = WritePipeHandle;
        consoleHandleInt->hConDrvReference = hConReference;
        consoleHandleInt->hConHostProcess = ProcInfo.hProcess;

        consoleHandle = consoleHandleInt;

        /* Cleanup */
        HeapFree(GetProcessHeap(), 0, AttrList);
        CloseHandle(InputHandle);
        CloseHandle(OutputHandle);
        CloseHandle(hConServer);
        CloseHandle(ProcInfo.hThread);
        CloseHandle(ProcInfo.hProcess);

        return 0;
    }
};

HANDLE hPipeIn = NULL, hPipeOut;
HANDLE hPipePTYIn, hPipePTYOut;
Terminal terminal;

void PipeListener()
{
    HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
    char   szBuffer[BUFF_SIZE];
    DWORD  dwBytesWritten, dwBytesRead, dwAvailable;

    int loopCount = 0;
    while (true)
    {
        loopCount++;

        // Check if data is available to read
        if (PeekNamedPipe(hPipeIn, NULL, 0, NULL, &dwAvailable, NULL))
        {
            if (loopCount % 100 == 0)
            {
                printf("Loop %d: PeekNamedPipe succeeded, "
                       "dwAvailable: %lu\n",
                       loopCount, dwAvailable);
            }

            if (dwAvailable > 0)
            {
                if (ReadFile(hPipeIn, szBuffer, BUFF_SIZE - 1,
                             &dwBytesRead, NULL))
                {
                    if (dwBytesRead == 0)
                        break;

                    szBuffer[dwBytesRead] =
                        '\0';

                    WriteFile(hConsole, szBuffer, dwBytesRead,
                              &dwBytesWritten, NULL);
                }
                else
                {
                    // ReadFile failed
                    DWORD error = GetLastError();
                    if (error == ERROR_BROKEN_PIPE ||
                        error == ERROR_PIPE_NOT_CONNECTED)
                        break;
                }
            }
            else
            {
                // No data available, sleep briefly to avoid busy
                // waiting
                Sleep(10);
            }
        }
        else
        {
            printf("ERROR: DATA IS UNAVAILABLE TO READ FROM PIPE.");


            DWORD error = GetLastError();
            if (error == ERROR_BROKEN_PIPE ||
                error == ERROR_PIPE_NOT_CONNECTED)
            {
                printf("FATAL ERROR: BROKEN PIPE OR PIPE NOT "
                       "CONNECTED. SHUTTING DOWN LISTENERS. PREPARE "
                       "FOR MELTDOWN.\n");

                break;
            }

            Sleep(10);
        }
    }
}

int main(int argc, char** argv)
{
    BOOL bRes;

    std::thread listenerThread(PipeListener);

    SECURITY_ATTRIBUTES attributes = {};
    attributes.nLength = sizeof(SECURITY_ATTRIBUTES);
    attributes.lpSecurityDescriptor = NULL;
    attributes.bInheritHandle = TRUE;

    /* Create the pipes to which the ConPTY will connect */
    if (CreatePipe(&hPipePTYIn, &hPipeOut, &attributes, 0) &&
        CreatePipe(&hPipeIn, &hPipePTYOut, &attributes, 0))
    {
        /* Create the Pseudo Console attached to the PTY-end of the
         * pipes */
        COORD                      consoleSize = {0};
        CONSOLE_SCREEN_BUFFER_INFO csbi = {0};
        if (GetConsoleScreenBufferInfo(
                GetStdHandle(STD_OUTPUT_HANDLE), &csbi))
        {
            consoleSize.X =
                csbi.srWindow.Right - csbi.srWindow.Left + 1;
            consoleSize.Y =
                csbi.srWindow.Bottom - csbi.srWindow.Top + 1;
        }

        terminal.InitializeConsole(consoleSize, hPipePTYIn,
                                   hPipePTYOut, 0);

    }

    /* Initialize thread attribute */
    size_t                       AttrSize;
    LPPROC_THREAD_ATTRIBUTE_LIST AttrList = NULL;
    InitializeProcThreadAttributeList(NULL, 1, 0, &AttrSize);
    AttrList = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(
        GetProcessHeap(), HEAP_ZERO_MEMORY, AttrSize);

    InitializeProcThreadAttributeList(AttrList, 1, 0, &AttrSize);

    bRes = UpdateProcThreadAttribute(
        AttrList, 0,
        PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE, /* 0x20016u */
        terminal.consoleHandle, sizeof terminal.consoleHandle, NULL,
        NULL);
    assert(bRes != 0);

    /* Initialize startup info struct */
    PROCESS_INFORMATION ProcInfo;
    memset(&ProcInfo, 0, sizeof ProcInfo);
    STARTUPINFOEXW SInfoEx;
    memset(&SInfoEx, 0, sizeof SInfoEx);
    SInfoEx.StartupInfo.cb = sizeof SInfoEx;
    SInfoEx.lpAttributeList = AttrList;

    wchar_t Command[] = L"cmd /c echo Hello World";

    bRes = CreateProcessW(NULL, Command, NULL, NULL, TRUE,
                          EXTENDED_STARTUPINFO_PRESENT, NULL, NULL,
                          &SInfoEx.StartupInfo, &ProcInfo);
    assert(bRes != 0);

    while (1) {}

    /* Cleanup */
    CloseHandle(ProcInfo.hThread);
    CloseHandle(ProcInfo.hProcess);
    HeapFree(GetProcessHeap(), 0, AttrList);
    CloseHandle(hPipeOut);
    CloseHandle(hPipeIn);
    CloseHandle(hPipePTYOut);
    CloseHandle(hPipePTYIn);
    listenerThread.join();

    return 0;
}