import signal
from msvcrt import get_osfhandle
from ctypes.wintypes import *
from ctypes import *
import os
import psutil
import struct


if sizeof(c_ulong) == sizeof(c_void_p):
    ULONG_PTR = c_ulong
elif sizeof(c_ulonglong) == sizeof(c_void_p):
    ULONG_PTR = c_ulonglong
DWORD_PTR = ULONG_PTR
SIZE_T = ULONG_PTR
PVOID = c_void_p

class PROC_THREAD_ATTRIBUTE_LIST(Structure):
    pass

class SECURITY_ATTRIBUTES(Structure):
    pass

class STARTUPINFOW(Structure):
    _fields_ = [
        ('cb',               DWORD),
        ('lpReserved',       LPWSTR),
        ('lpDesktop',        LPWSTR),
        ('lpTitle',          LPWSTR),
        ('dwX',              DWORD),
        ('dwY',              DWORD),
        ('dwXSize',          DWORD),
        ('dwYSize',          DWORD),
        ('dwXCountChars',    DWORD),
        ('dwYCountChars',    DWORD),
        ('dwFillAttributes', DWORD),
        ('dwFlags',          DWORD),
        ('wShowWindow',      WORD),
        ('cbReserved2',      WORD),
        ('lpReserved2',      LPBYTE),
        ('hStdInput',        HANDLE),
        ('hStdOutput',       HANDLE),
        ('hStdError',        HANDLE),
    ]

class STARTUPINFOEXW(Structure):
    _fields_ = [
        ('StartupInfo',     STARTUPINFOW),
        ('lpAttributeList', POINTER(PROC_THREAD_ATTRIBUTE_LIST)),
    ]

class PROCESS_INFORMATION(Structure):
    _fields_ = [
        ('hProcess',    HANDLE),
        ('hThread',     HANDLE),
        ('dwProcessId', DWORD),
        ('dwThreadId',  DWORD),
    ]


class UnbufferedProcess:
    def __init__(self, path):
        assert os.path.exists(path.split()[0]), f"{path.split()[0]} is not a valid exe"
        self.path = path
        self.exe_name = os.path.basename(path.split()[0])
        self.process_info = None
        self.stdin = None
        self.stdout = None

    def run(self):
        kernel32 = WinDLL('kernel32', use_last_error=True)
        CreateProcessW = kernel32.CreateProcessW
        CreateProcessW.argtypes = [LPCWSTR, LPWSTR,
                                   POINTER(SECURITY_ATTRIBUTES),
                                   POINTER(SECURITY_ATTRIBUTES), BOOL, DWORD,
                                   LPVOID, LPCWSTR, POINTER(STARTUPINFOW),
                                   POINTER(PROCESS_INFORMATION)]
        CreateProcessW.restype = BOOL
        self.CloseHandle = kernel32.CloseHandle
        self.CloseHandle.argtypes = [HANDLE]
        self.CloseHandle.restype = BOOL
        InitializeProcThreadAttributeList = kernel32.InitializeProcThreadAttributeList
        InitializeProcThreadAttributeList.argtypes = [
            POINTER(PROC_THREAD_ATTRIBUTE_LIST), DWORD, DWORD, POINTER(SIZE_T)]
        InitializeProcThreadAttributeList.restype = BOOL
        UpdateProcThreadAttribute = kernel32.UpdateProcThreadAttribute
        UpdateProcThreadAttribute.argtypes = [
            POINTER(PROC_THREAD_ATTRIBUTE_LIST), DWORD, DWORD_PTR, PVOID,
            SIZE_T, PVOID, POINTER(SIZE_T)]
        UpdateProcThreadAttribute.restype = BOOL

        self.stdin_r, self.stdin = os.pipe()
        self.stdout, self.stdout_w = os.pipe()

        os.set_inheritable(self.stdin_r, True)
        os.set_inheritable(self.stdout_w, True)

        stdin_handle = get_osfhandle(self.stdin_r)
        stdout_handle = get_osfhandle(self.stdout_w)

        PROC_THREAD_ATTRIBUTE_HANDLE_LIST = 0x20002
        size = SIZE_T()
        InitializeProcThreadAttributeList(None, 1, 0, byref(size))
        attr_list = cast(create_string_buffer(size.value),
                         POINTER(PROC_THREAD_ATTRIBUTE_LIST))
        handle_list = (HANDLE * 2)(stdin_handle, stdout_handle)
        UpdateProcThreadAttribute(attr_list, 0,
                                  PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
                                  handle_list, sizeof(handle_list), None, None)

        FOPEN = 0x01
        FDEV = 0x40
        buf = struct.pack('I', 2)  # cfi_len, number of handles
        buf += 2 * struct.pack('B', FOPEN | FDEV)  # file info
        buf += struct.pack('PP', stdin_handle, stdout_handle)  # os handle

        startup_info_ex = STARTUPINFOEXW()
        startup_info_ex.lpAttributeList = attr_list
        startup_info = startup_info_ex.StartupInfo
        startup_info.cb = sizeof(startup_info_ex)
        startup_info.cbReserved2 = len(buf)
        startup_info.lpReserved2 = cast(buf, LPBYTE)

        EXTENDED_STARTUPINFO_PRESENT = 0x00080000
        self.process_info = PROCESS_INFORMATION()
        CreateProcessW(None, self.path, None, None, True,
                       EXTENDED_STARTUPINFO_PRESENT, None, None,
                       byref(startup_info), byref(self.process_info))
        #print(f"Opened Process: {self.exe_name} with PID {self.process_info.dwProcessId}")

        os.close(self.stdin_r)
        os.close(self.stdout_w)

    def stop(self):
        #print(f"Killing Process: {self.exe_name} with PID {self.process_info.dwProcessId}")
        os.kill(self.process_info.dwProcessId, signal.SIGTERM)
        self.CloseHandle(self.process_info.hThread)
        self.CloseHandle(self.process_info.hProcess)
        self.process_info = None

        # for p in psutil.process_iter():
        #     if p.name() == self.exe_name:
        #         p.kill()
        #         p.wait(timeout=10)