// ke_process.c - process model
#include <string.h>
#include <intrin.h>
#include <ke_process.h>
#include <ke_thread.h>
#include <ke_sched.h>
#include <ke_event.h>
#include <ke_userenv.h>
#include <mm_virt.h>
#include <ob_object.h>

// ---------------------------------------------------------------------------
// Global process list
// ---------------------------------------------------------------------------
static LIST_ENTRY gProcessList;
static SPINLOCK   gProcessListLock;
static uint32_t   gNextProcessId = 1;

static void ProcessInit(void) {
    static int inited = 0;
    if (!inited) {
        LIST_INIT(&gProcessList);
        gProcessListLock = (SPINLOCK)SPINLOCK_INIT;
        inited = 1;
    }
}

LIST_ENTRY* KeGetProcessList(void) {
    ProcessInit();
    return &gProcessList;
}

PROCESS* KeGetCurrentProcess(void) {
    THREAD* currentThread = KeGetCurrentThread();
    if (currentThread) {
        return currentThread->Owner;
    }
    return NULL;
}

size_t KeCopyProcessIds(uint32_t* pidsOut, size_t maxCount) {
    ProcessInit();

    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    size_t count = 0;
    LIST_ENTRY* head = &gProcessList;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        PROCESS* p = CONTAINING_RECORD(node, PROCESS, GlobalLink);
        if (pidsOut && count < maxCount) pidsOut[count] = p->ProcessId;
        count++;
    }
    KeReleaseSpinlock(&gProcessListLock, flags);
    return count;
}

// ---------------------------------------------------------------------------
// Destructor (called by ObDereferenceObject when RefCount hits zero)
// ---------------------------------------------------------------------------
static void ProcessDestroy(OBJECT* obj) {
    PROCESS* proc = CONTAINING_RECORD(obj, PROCESS, Header);

    // Release all remaining handles.
    for (int i = 1; i < HANDLE_TABLE_SIZE; i++) {
        if (proc->Handles[i]) {
            ObDereferenceObject(proc->Handles[i]);
            proc->Handles[i] = (OBJECT *)0;
        }
    }

    // Release the main executable image.
    if (proc->ImageBase && proc->ImageSize) {
        MmFreeUserMemory((void *)proc->ImageBase, proc->ImageSize);
        proc->ImageBase = 0;
        proc->ImageSize = 0;
    }

    // Release per-process modules.
    for (uint32_t i = 0; i < proc->ModuleCount; i++) {
        if (proc->Modules[i].Base && proc->Modules[i].Size) {
            MmFreeUserMemory((void *)proc->Modules[i].Base, proc->Modules[i].Size);
        }
        proc->Modules[i].Base = 0;
        proc->Modules[i].Size = 0;
        proc->Modules[i].Name[0] = '\0';
    }
    proc->ModuleCount = 0;

    // Release the user-mode process environment.
    if (proc->UserPebBase && proc->UserPebSize) {
        MmFreeUserMemory((void *)proc->UserPebBase, proc->UserPebSize);
        proc->UserPebBase = 0;
        proc->UserPebSize = 0;
    }
    if (proc->UserImagePathBase && proc->UserImagePathSize) {
        MmFreeUserMemory((void *)proc->UserImagePathBase, proc->UserImagePathSize);
        proc->UserImagePathBase = 0;
        proc->UserImagePathSize = 0;
    }
    if (proc->UserCmdlineBase && proc->UserCmdlineSize) {
        MmFreeUserMemory((void *)proc->UserCmdlineBase, proc->UserCmdlineSize);
        proc->UserCmdlineBase = 0;
        proc->UserCmdlineSize = 0;
    }
    if (proc->UserCwdBase && proc->UserCwdSize) {
        MmFreeUserMemory((void *)proc->UserCwdBase, proc->UserCwdSize);
        proc->UserCwdBase = 0;
        proc->UserCwdSize = 0;
    }

    // Release the exit event.
    if (proc->ExitEvent) {
        ObDereferenceObject(&proc->ExitEvent->Header);
        proc->ExitEvent = (EVENT *)0;
    }

    // Remove from global list.
    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    ListRemove(&proc->GlobalLink);
    KeReleaseSpinlock(&gProcessListLock, flags);

    // Memory for the struct itself is freed by ObDereferenceObject after
    // Destroy returns.
}

// ---------------------------------------------------------------------------
// KeCreateProcess
// ---------------------------------------------------------------------------
KSTATUS KeCreateProcess(const char* name, PROCESS **out) {
    ProcessInit();

    if (!name || !out) return STATUS_INVALID_PARAMETER;

    PROCESS* proc = (PROCESS *)ObAllocateObject(ObjectTypeProcess, sizeof(PROCESS));
    if (!proc) return STATUS_NO_MEMORY;

    proc->Header.Destroy = ProcessDestroy;

    proc->ProcessId  = _InterlockedExchangeAdd((volatile long *)&gNextProcessId, 1);
    strncpy(proc->Name, name, PROCESS_NAME_MAX - 1);
    proc->Name[PROCESS_NAME_MAX - 1] = '\0';

    LIST_INIT(&proc->ThreadList);
    proc->ThreadCount = 0;
    proc->HandleLock  = (SPINLOCK)SPINLOCK_INIT;
    proc->ExitCode    = 0;
    proc->ImageBase   = 0;
    proc->ImageSize   = 0;
    proc->UserPebBase = 0;
    proc->UserPebSize = 0;
    proc->UserImagePathBase = 0;
    proc->UserImagePathSize = 0;
    proc->UserCmdlineBase = 0;
    proc->UserCmdlineSize = 0;
    proc->UserCwdBase = 0;
    proc->UserCwdSize = 0;
    proc->ModuleCount = 0;
    proc->CurrentDirectory[0] = '\0';
    proc->CwdLock = (SPINLOCK)SPINLOCK_INIT;

    // Slot 0 is the invalid handle - leave it NULL.
    memset(proc->Handles, 0, sizeof(proc->Handles));

    LIST_INIT(&proc->GlobalLink);

    // Create the manual-reset exit event (initially unsignaled).
    if (!KSUCCESS(KeEventCreate(EventManualReset, 0, &proc->ExitEvent))) {
        ObDereferenceObject(&proc->Header);
        return STATUS_NO_MEMORY;
    }

    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    ListInsertTail(&gProcessList, &proc->GlobalLink);
    KeReleaseSpinlock(&gProcessListLock, flags);

    *out = proc;
    return STATUS_SUCCESS;
}


// ---------------------------------------------------------------------------
// KeSetupUserProcessEnvironment
// ---------------------------------------------------------------------------
static KSTATUS CopyEnvironmentString(const char* src, uintptr_t* baseOut, size_t* sizeOut) {
    if (!baseOut || !sizeOut) return STATUS_INVALID_PARAMETER;
    if (!src) src = "";

    size_t len = strlen(src) + 1;
    char* dst = (char*)MmAllocateUserMemory(len);
    if (!dst) return STATUS_NO_MEMORY;

    memcpy(dst, src, len);
    *baseOut = (uintptr_t)dst;
    *sizeOut = len;
    return STATUS_SUCCESS;
}

static KSTATUS CopyEnvironmentStringFixed(const char* src, size_t capacity,
                                          uintptr_t* baseOut, size_t* sizeOut) {
    if (!baseOut || !sizeOut || capacity == 0) return STATUS_INVALID_PARAMETER;
    if (!src) src = "";

    char* dst = (char*)MmAllocateUserMemory(capacity);
    if (!dst) return STATUS_NO_MEMORY;
    memset(dst, 0, capacity);

    size_t len = strlen(src);
    if (len >= capacity) len = capacity - 1;
    memcpy(dst, src, len);
    dst[len] = '\0';

    *baseOut = (uintptr_t)dst;
    *sizeOut = capacity;
    return STATUS_SUCCESS;
}

static void FreePartialProcessEnvironment(PROCESS* proc) {
    if (!proc) return;

    if (proc->UserPebBase && proc->UserPebSize) {
        MmFreeUserMemory((void*)proc->UserPebBase, proc->UserPebSize);
        proc->UserPebBase = 0;
        proc->UserPebSize = 0;
    }
    if (proc->UserImagePathBase && proc->UserImagePathSize) {
        MmFreeUserMemory((void*)proc->UserImagePathBase, proc->UserImagePathSize);
        proc->UserImagePathBase = 0;
        proc->UserImagePathSize = 0;
    }
    if (proc->UserCmdlineBase && proc->UserCmdlineSize) {
        MmFreeUserMemory((void*)proc->UserCmdlineBase, proc->UserCmdlineSize);
        proc->UserCmdlineBase = 0;
        proc->UserCmdlineSize = 0;
    }
    if (proc->UserCwdBase && proc->UserCwdSize) {
        MmFreeUserMemory((void*)proc->UserCwdBase, proc->UserCwdSize);
        proc->UserCwdBase = 0;
        proc->UserCwdSize = 0;
    }
}

KSTATUS KeSetupUserProcessEnvironment(PROCESS* proc,
                                       const char* imagePath,
                                       const char* commandLine,
                                       const char* currentDirectory) {
    if (!proc || !imagePath) return STATUS_INVALID_PARAMETER;

    const char* finalCommandLine = (commandLine && commandLine[0]) ? commandLine : imagePath;
    const char* finalCwd = currentDirectory ? currentDirectory : proc->CurrentDirectory;

    FreePartialProcessEnvironment(proc);

    KSTATUS s = CopyEnvironmentString(imagePath, &proc->UserImagePathBase, &proc->UserImagePathSize);
    if (!KSUCCESS(s)) goto fail;

    s = CopyEnvironmentString(finalCommandLine, &proc->UserCmdlineBase, &proc->UserCmdlineSize);
    if (!KSUCCESS(s)) goto fail;

    s = CopyEnvironmentStringFixed(finalCwd, SYS_CWD_MAX, &proc->UserCwdBase, &proc->UserCwdSize);
    if (!KSUCCESS(s)) goto fail;

    SYS_PEB* peb = (SYS_PEB*)MmAllocateUserMemory(sizeof(SYS_PEB));
    if (!peb) { s = STATUS_NO_MEMORY; goto fail; }
    memset(peb, 0, sizeof(SYS_PEB));

    peb->Magic = SYS_PEB_MAGIC;
    peb->Version = SYS_ENV_VERSION;
    peb->ProcessId = proc->ProcessId;
    peb->ImagePath = (const char*)proc->UserImagePathBase;
    peb->ImagePathLength = proc->UserImagePathSize ? proc->UserImagePathSize - 1 : 0;
    peb->CommandLine = (const char*)proc->UserCmdlineBase;
    peb->CommandLineLength = proc->UserCmdlineSize ? proc->UserCmdlineSize - 1 : 0;
    peb->CurrentDirectory = (const char*)proc->UserCwdBase;
    peb->CurrentDirectoryLength = strlen(finalCwd);

    proc->UserPebBase = (uintptr_t)peb;
    proc->UserPebSize = sizeof(SYS_PEB);
    return STATUS_SUCCESS;

fail:
    FreePartialProcessEnvironment(proc);
    return s;
}

// ---------------------------------------------------------------------------
// KeExitProcess
// ---------------------------------------------------------------------------
void KeExitProcess(PROCESS* process, int exitCode) {
    if (!process) return;
    process->ExitCode = exitCode;
    // Wake any threads waiting for this process to finish.
    if (process->ExitEvent) KeSignalEvent(process->ExitEvent);
    // The last thread exiting via KeExitThread will drop the process refcount.
    // This function is a policy hook; actual cleanup is in ProcessDestroy.
    ObDereferenceObject(&process->Header);
}

// ---------------------------------------------------------------------------
// KeGetProcessExitCode
// ---------------------------------------------------------------------------
KSTATUS KeGetProcessExitCode(PROCESS* process, int* exitCodeOut) {
    if (!process || !exitCodeOut) return STATUS_INVALID_PARAMETER;
    *exitCodeOut = process->ExitCode;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// Handle table
// ---------------------------------------------------------------------------
KSTATUS KeAllocateHandle(PROCESS* process, OBJECT* object, uint32_t* handleOut) {
    if (!process || !object || !handleOut) return STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);

    for (int i = 1; i < HANDLE_TABLE_SIZE; i++) {
        if (!process->Handles[i]) {
            ObReferenceObject(object);
            process->Handles[i] = object;
            *handleOut = (uint32_t)i;
            KeReleaseSpinlock(&process->HandleLock, flags);
            return STATUS_SUCCESS;
        }
    }

    KeReleaseSpinlock(&process->HandleLock, flags);
    return STATUS_NO_MEMORY;
}

void KeCloseHandle(PROCESS* process, uint32_t handle) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE) return;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    OBJECT* obj = process->Handles[handle];
    if (obj) {
        process->Handles[handle] = (OBJECT *)0;
        KeReleaseSpinlock(&process->HandleLock, flags);
        ObDereferenceObject(obj);
        return;
    }
    KeReleaseSpinlock(&process->HandleLock, flags);
}

OBJECT* KeLookupHandle(PROCESS* process, uint32_t handle) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE)
        return (OBJECT *)0;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    OBJECT* obj = process->Handles[handle];
    if (obj) ObReferenceObject(obj);
    KeReleaseSpinlock(&process->HandleLock, flags);
    return obj;
}

PROCESS* KeGetProcessById(uint32_t pid) {
    ProcessInit();
    uintptr_t flags = KeAcquireSpinlock(&gProcessListLock);
    LIST_ENTRY* head = &gProcessList;
    LIST_ENTRY* node = head->Fwd;
    while (node != head) {
        PROCESS* p = CONTAINING_RECORD(node, PROCESS, GlobalLink);
        if (p->ProcessId == pid) {
            ObReferenceObject(&p->Header);
            KeReleaseSpinlock(&gProcessListLock, flags);
            return p;
        }
        node = node->Fwd;
    }
    KeReleaseSpinlock(&gProcessListLock, flags);
    return (PROCESS*)0;
}