// Kernel process model. This file owns address-space metadata, handle tables, and process lifetime transitions.

#include <string.h>
#include <intrin.h>
#include <kernel/ke_process.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_event.h>
#include <kernel/ke_userenv.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>
#include <kernel/ke_cpu.h>
#include <kernel/hal/hal_apic.h>
#include <ntlite/syscall.h>

// Global process list.
static LIST_ENTRY gProcessList;
static SPINLOCK   gProcessListLock;
static uint32_t   gNextProcessId = 0;
static PROCESS*   gSystemProcess = (PROCESS*)0;

// Processes init.
static void KiProcessInit(void) {
    static int inited = 0;
    if (!inited) {
        LIST_INIT(&gProcessList);
        gProcessListLock = (SPINLOCK)SPINLOCK_INIT;
        inited = 1;
    }
}

// Returns process List.
LIST_ENTRY* KeGetProcessList(void) {
    KiProcessInit();
    return &gProcessList;
}

// Returns current Process.
PROCESS* KeGetCurrentProcess(void) {
    THREAD* currentThread = KeGetCurrentThread();
    if (currentThread) {
        return currentThread->Owner;
    }
    return NULL;
}

// Returns system Process.
PROCESS* KeGetSystemProcess(void) {
    return gSystemProcess;
}

// Creates system Process.
KSTATUS KeCreateSystemProcess(PROCESS** out) {
    if (!out) return STATUS_INVALID_PARAMETER;
    if (gSystemProcess) {
        *out = gSystemProcess;
        return STATUS_SUCCESS;
    }

    PROCESS* proc = (PROCESS*)0;
    KSTATUS s = KeCreateProcess("system", &proc);
    if (!KSUCCESS(s)) {
        *out = (PROCESS*)0;
        return s;
    }

    gSystemProcess = proc;
    *out = proc;
    return STATUS_SUCCESS;
}

// Implements the count Process Handles Locked operation.
static uint32_t KiCountProcessHandlesLocked(PROCESS* process) {
    uint32_t count = 0;
    if (!process) return 0;

    for (uint32_t i = 1; i < HANDLE_TABLE_SIZE; i++) {
        if (process->Handles[i].Object) count++;
    }
    return count;
}

// Copies truncated String.
static void KiCopyTruncatedString(char* dst, size_t dstSize, const char* src, size_t srcLen) {
    if (!dst || dstSize == 0) return;
    size_t n = 0;
    if (src) {
        while (n + 1 < dstSize && n < srcLen && src[n]) {
            dst[n] = src[n];
            n++;
        }
    }
    dst[n] = '\0';
}

// Copies process Info.
size_t KeCopyProcessInfo(struct SYS_PROCESS_INFO* infoOut, size_t maxCount) {
    KiProcessInit();

    PROCESS* current = KeGetCurrentProcess();
    PROCESS* system = KeGetSystemProcess();

    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 (infoOut && count < maxCount) {
            SYS_PROCESS_INFO* info = &infoOut[count];
            memset(info, 0, sizeof(*info));
            info->Size = sizeof(*info);
            info->ProcessId = p->ProcessId;
            info->ImageBase = p->ImageBase;
            info->ImageSize = p->ImageSize;
            info->PebBase = p->UserPebBase;
            info->ModuleCount = p->ModuleCount;

            uintptr_t handleFlags = KeAcquireSpinlock(&p->HandleLock);
            uint32_t threadCount = p->ThreadCount;
            uint32_t handleCount = KiCountProcessHandlesLocked(p);
            int exitCode = p->ExitCode;
            long terminating = p->Terminating;
            long exited = p->Exited;
            KeReleaseSpinlock(&p->HandleLock, handleFlags);

            info->ThreadCount = threadCount;
            info->HandleCount = handleCount;
            info->ExitCode = exitCode;

            if (exited) {
                info->State = SYS_PROCESS_STATE_EXITED;
                info->Flags |= SYS_PROCESS_FLAG_EXITED;
            } else if (terminating) {
                info->State = SYS_PROCESS_STATE_TERMINATING;
                info->Flags |= SYS_PROCESS_FLAG_TERMINATING;
            } else {
                info->State = SYS_PROCESS_STATE_RUNNING;
            }

            if (p == system) info->Flags |= SYS_PROCESS_FLAG_SYSTEM;
            if (p == current) info->Flags |= SYS_PROCESS_FLAG_CURRENT;

            MmQueryProcessUserMemory(p, &info->UserRegionCount,
                                     &info->UserReservedBytes,
                                     &info->UserCommittedBytes);

            KiCopyTruncatedString(info->Name, sizeof(info->Name), p->Name, sizeof(p->Name));

            uintptr_t cwdFlags = KeAcquireSpinlock(&p->CwdLock);
            info->CurrentDrive = p->CurrentDrive;
            KiCopyTruncatedString(info->CurrentDirectory, sizeof(info->CurrentDirectory),
                                  p->CurrentDirectory, p->CurrentDirectoryLength);
            KeReleaseSpinlock(&p->CwdLock, cwdFlags);
        }
        count++;
    }
    KeReleaseSpinlock(&gProcessListLock, flags);
    return count;
}

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

    if (proc == gSystemProcess)
        gSystemProcess = (PROCESS*)0;

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

    // Release every user VA region owned by this process. This includes the
    // main image, DLLs, TEBs/stacks that survived thread teardown, process
    // environment strings, and leaked user allocations.
    MmDestroyAllUserRegions(proc);
    proc->ImageBase = 0;
    proc->ImageSize = 0;
    proc->ModuleCount = 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;

    // 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) {
    KiProcessInit();

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

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

    proc->Header.Destroy = KiProcessDestroy;

    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->Terminating = 0;
    proc->Exited      = 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->CurrentDirectoryLength = 0;
    proc->CurrentDrive = 0;
    proc->CwdLock = (SPINLOCK)SPINLOCK_INIT;
    LIST_INIT(&proc->UserRegions);
    proc->UserRegionLock = (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 KiCopyEnvironmentStringView(PROCESS* proc,
                                           SYS_STRING_VIEW src,
                                           uintptr_t* baseOut,
                                           size_t* sizeOut) {
    if (!baseOut || !sizeOut) return STATUS_INVALID_PARAMETER;
    *baseOut = 0;
    *sizeOut = 0;

    if (!SysSvIsValid(src)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(src, '\0')) return STATUS_INVALID_PARAMETER;
    if (src.Length == (size_t)-1) return STATUS_BUFFER_TOO_SMALL;

    size_t len = src.Length + 1;
    char* dst = (char*)MmAllocateUserMemoryForProcess(proc, len);
    if (!dst) return STATUS_NO_MEMORY;

    if (src.Length) memcpy(dst, src.Data, src.Length);
    dst[src.Length] = '\0';

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

static KSTATUS KiCopyEnvironmentStringViewFixed(PROCESS* proc,
                                                SYS_STRING_VIEW src,
                                                size_t capacity,
                                                uintptr_t* baseOut,
                                                size_t* sizeOut) {
    if (!baseOut || !sizeOut || capacity == 0) return STATUS_INVALID_PARAMETER;
    *baseOut = 0;
    *sizeOut = 0;

    if (!SysSvIsValid(src)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(src, '\0')) return STATUS_INVALID_PARAMETER;
    if (src.Length + 1 > capacity) return STATUS_BUFFER_TOO_SMALL;

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

    if (src.Length) memcpy(dst, src.Data, src.Length);
    dst[src.Length] = '\0';

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

// Releases partial Process Environment.
static void KiFreePartialProcessEnvironment(PROCESS* proc) {
    if (!proc) return;

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

static KSTATUS KiFormatNativeCwdV(const PROCESS* proc,
                                  SYS_STRING_VIEW currentDirectory,
                                  char* out,
                                  size_t outSize,
                                  size_t* lengthOut) {
    if (lengthOut) *lengthOut = 0;
    if (!out || outSize == 0) return STATUS_INVALID_PARAMETER;
    out[0] = '\0';

    if (!SysSvIsValid(currentDirectory)) return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(currentDirectory, '\0')) return STATUS_INVALID_PARAMETER;

    uint8_t drive = proc ? proc->CurrentDrive : 0;
    if (drive > 7) drive = 0;

    if (currentDirectory.Length + 4 > outSize) return STATUS_BUFFER_TOO_SMALL;

    out[0] = (char)('0' + drive);
    out[1] = ':';
    out[2] = '/';
    if (currentDirectory.Length)
        memcpy(out + 3, currentDirectory.Data, currentDirectory.Length);
    out[3 + currentDirectory.Length] = '\0';

    if (lengthOut) *lengthOut = 3 + currentDirectory.Length;
    return STATUS_SUCCESS;
}

KSTATUS KeSetupUserProcessEnvironment(PROCESS* proc,
                                       SYS_STRING_VIEW imagePath,
                                       SYS_STRING_VIEW commandLine,
                                       SYS_STRING_VIEW currentDirectory) {
    if (!proc || imagePath.Length == 0) return STATUS_INVALID_PARAMETER;
    if (!SysSvIsValid(imagePath) || !SysSvIsValid(commandLine) || !SysSvIsValid(currentDirectory))
        return STATUS_INVALID_PARAMETER;
    if (SysSvContainsByte(imagePath, '\0') ||
        SysSvContainsByte(commandLine, '\0') ||
        SysSvContainsByte(currentDirectory, '\0'))
        return STATUS_INVALID_PARAMETER;

    SYS_STRING_VIEW finalCommandLine = commandLine.Length ? commandLine : imagePath;
    SYS_STRING_VIEW finalCwdRel = currentDirectory;
    if (finalCwdRel.Length == 0 && proc->CurrentDirectoryLength)
        finalCwdRel = SysSvFromParts(proc->CurrentDirectory, proc->CurrentDirectoryLength);

    char finalCwd[SYS_CWD_MAX];
    size_t finalCwdLength = 0;
    KSTATUS s = KiFormatNativeCwdV(proc, finalCwdRel, finalCwd, sizeof(finalCwd), &finalCwdLength);
    if (!KSUCCESS(s)) return s;

    KiFreePartialProcessEnvironment(proc);

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

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

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

    SYS_PEB* peb = (SYS_PEB*)MmAllocateUserMemoryForProcess(proc, 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 = SysSvFromParts((const char*)proc->UserImagePathBase, imagePath.Length);
    peb->CommandLine = SysSvFromParts((const char*)proc->UserCmdlineBase, finalCommandLine.Length);
    peb->CurrentDirectory = SysSvFromParts((const char*)proc->UserCwdBase, finalCwdLength);

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

fail:
    KiFreePartialProcessEnvironment(proc);
    return s;
}

// Process exit / termination.
static void KiCompleteProcessExitIfLastThread(PROCESS* process) {
    if (!process) return;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    uint32_t threadCount = process->ThreadCount;
    KeReleaseSpinlock(&process->HandleLock, flags);

    if (threadCount == 0 && _InterlockedExchange(&process->Exited, 1) == 0) {
        if (process->ExitEvent) KeSignalEvent(process->ExitEvent);
    }
}

// Implements the exit Process operation.
void KeExitProcess(PROCESS* process, int exitCode) {
    if (!process) return;
    process->ExitCode = exitCode;
    _InterlockedExchange(&process->Terminating, 1);

    // The calling thread exits immediately after this in SysExitProcess. Do
    // not drop the process object here: each thread owns one process reference,
    // and KiThreadDestroy releases that reference after unlinking the thread.
    KiCompleteProcessExitIfLastThread(process);
}

// Implements the terminate Process operation.
KSTATUS KeTerminateProcess(PROCESS* process, int exitCode) {
    if (!process) return STATUS_INVALID_PARAMETER;

    process->ExitCode = exitCode;
    _InterlockedExchange(&process->Terminating, 1);

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    LIST_ENTRY* head = &process->ThreadList;
    THREAD* self = KeGetCurrentThread();
    KCPU* selfCpu = KeGetCurrentCpu();
    uint32_t selfIndex = selfCpu ? selfCpu->CpuIndex : KE_INVALID_CPU;

    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        THREAD* t = CONTAINING_RECORD(node, THREAD, ProcessLink);
        if (t == self)
            continue;

        (void)KeAlertThread(t);

        if (t->RunningCpu != KE_INVALID_CPU && t->RunningCpu != selfIndex && t->RunningCpu < gCpuCount) {
            KCPU* cpu = &gCpus[t->RunningCpu];
            if (cpu->Present && cpu->Started && cpu->Online) {
                _InterlockedExchange(&cpu->ReschedulePending, 1);
                HalApicSendRescheduleIpi(cpu->LapicId);
            }
        }
    }
    KeReleaseSpinlock(&process->HandleLock, flags);

    KiCompleteProcessExitIfLastThread(process);
    return STATUS_SUCCESS;
}

// Checks whether process Terminating.
int KeIsProcessTerminating(PROCESS* process) {
    return process && process->Terminating != 0;
}

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

// Handle table.
KSTATUS KeAllocateHandleWithRights(PROCESS* process, OBJECT* object, uint32_t rights, uint32_t* handleOut) {
    if (!process || !object || !handleOut || rights == 0) return STATUS_INVALID_PARAMETER;

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

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

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

// Creates handle.
KSTATUS KeAllocateHandle(PROCESS* process, OBJECT* object, uint32_t* handleOut) {
    return KeAllocateHandleWithRights(process, object, HANDLE_RIGHT_ALL, handleOut);
}

// Releases handle.
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].Object;
    if (obj) {
        process->Handles[handle].Object = (OBJECT*)0;
        process->Handles[handle].Rights = 0;
        process->Handles[handle].Flags = 0;
        KeReleaseSpinlock(&process->HandleLock, flags);
        ObDereferenceObject(obj);
        return;
    }
    KeReleaseSpinlock(&process->HandleLock, flags);
}

// Returns handle With Rights.
OBJECT* KeLookupHandleWithRights(PROCESS* process, uint32_t handle, uint32_t requiredRights) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE)
        return (OBJECT*)0;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    HANDLE_ENTRY* entry = &process->Handles[handle];
    OBJECT* obj = entry->Object;
    if (obj && requiredRights && ((entry->Rights & requiredRights) != requiredRights))
        obj = (OBJECT*)0;
    if (obj) ObReferenceObject(obj);
    KeReleaseSpinlock(&process->HandleLock, flags);
    return obj;
}

// Returns handle.
OBJECT* KeLookupHandle(PROCESS* process, uint32_t handle) {
    return KeLookupHandleWithRights(process, handle, 0);
}

// Returns handle Rights.
uint32_t KeGetHandleRights(PROCESS* process, uint32_t handle) {
    if (!process || handle == 0 || handle >= HANDLE_TABLE_SIZE) return 0;

    uintptr_t flags = KeAcquireSpinlock(&process->HandleLock);
    uint32_t rights = process->Handles[handle].Object ? process->Handles[handle].Rights : 0;
    KeReleaseSpinlock(&process->HandleLock, flags);
    return rights;
}

// Returns process By Id.
PROCESS* KeGetProcessById(uint32_t pid) {
    KiProcessInit();
    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;
}