// Kernel pipe object lifetime, shared backing, readiness, and wait integration.

#include <ntlite/atomic.h>
#include <stdbool.h>
#include <string.h>
#include <intrin.h>
#include <kernel/io_pipe.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_process.h>
#include <kernel/mm_pool.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>
#include <ntlite/syscall.h>

#define IO_PIPE_STATE_TAG MM_POOL_TAG('p', 'i', 'p', 's')
#define IO_PIPE_FRAME_TAG MM_POOL_TAG('p', 'i', 'p', 'f')

typedef enum IO_PIPE_DIRECTION {
    IoPipeReadEnd = 1,
    IoPipeWriteEnd = 2,
} IO_PIPE_DIRECTION;

typedef struct IO_PIPE_STATE {
    volatile int32_t RefCount;
    SPINLOCK Lock;
    EVENT* ReadEvent;
    EVENT* WriteEvent;
    PHYSICAL_ADDRESS* Frames;
    size_t PageCount;
    MM_KERNEL_VIEW* KernelView;
    SYS_PIPE_SHARED_STATE* Shared;
} IO_PIPE_STATE;

struct IO_PIPE_ENDPOINT {
    OBJECT Header;
    IO_PIPE_STATE* State;
    IO_PIPE_DIRECTION Direction;
};

// Returns the current buffered-byte count when the shared counters are sane.
static bool KiPipeUsedBytes(const SYS_PIPE_SHARED_STATE* shared, uint64_t* usedOut) {
    if (!shared || !usedOut || shared->Magic != SYS_PIPE_SHARED_MAGIC ||
        shared->Version != SYS_PIPE_SHARED_VERSION ||
        shared->Capacity != SYS_PIPE_BUFFER_CAPACITY ||
        shared->HeaderBytes != SYS_PIPE_SHARED_HEADER_BYTES)
        return false;

    uint64_t readPosition = shared->ReadPosition;
    uint64_t writePosition = shared->WritePosition;
    uint64_t used = writePosition - readPosition;
    if (used > SYS_PIPE_BUFFER_CAPACITY)
        return false;
    *usedOut = used;
    return true;
}

// Refreshes level-triggered endpoint events while the state lock is held.
static void KiPipeRefreshEventsLocked(IO_PIPE_STATE* state) {
    if (!state || !state->Shared)
        return;

    SYS_PIPE_SHARED_STATE* shared = state->Shared;
    uint64_t used = 0;
    if (!KiPipeUsedBytes(shared, &used)) {
        KeSignalEvent(state->ReadEvent);
        KeSignalEvent(state->WriteEvent);
        return;
    }

    if (used != 0 || shared->WriterOpen == 0)
        KeSignalEvent(state->ReadEvent);
    else
        KeClearEvent(state->ReadEvent);

    uint64_t freeBytes = SYS_PIPE_BUFFER_CAPACITY - used;
    if (shared->ReaderOpen == 0 || freeBytes >= SYS_PIPE_ATOMIC_WRITE_MAX)
        KeSignalEvent(state->WriteEvent);
    else
        KeClearEvent(state->WriteEvent);
}

// Releases shared pipe storage after both endpoint objects are gone.
static void KiPipeStateRelease(IO_PIPE_STATE* state) {
    if (!state || NtlAtomicDecrement32(&state->RefCount) != 0)
        return;

    if (state->KernelView) {
        MmDestroyKernelView(state->KernelView);
        state->KernelView = (MM_KERNEL_VIEW*)0;
        state->Shared = (SYS_PIPE_SHARED_STATE*)0;
    }
    if (state->ReadEvent)
        ObDereferenceObject(&state->ReadEvent->Header);
    if (state->WriteEvent)
        ObDereferenceObject(&state->WriteEvent->Header);
    if (state->Frames) {
        for (size_t i = 0; i < state->PageCount; ++i)
            if (state->Frames[i])
                (void)MmPfnRelease(state->Frames[i]);
        MmFreePool(state->Frames);
    }
    MmFreePool(state);
}

// Destroys one endpoint and publishes final-reader/final-writer state.
static void KiPipeEndpointDestroy(OBJECT* object) {
    if (!object)
        return;

    IO_PIPE_ENDPOINT* endpoint = CONTAINING_RECORD(object, IO_PIPE_ENDPOINT, Header);
    IO_PIPE_STATE* state = endpoint->State;
    endpoint->State = (IO_PIPE_STATE*)0;
    if (!state)
        return;

    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    if (state->Shared) {
        if (endpoint->Direction == IoPipeReadEnd)
            state->Shared->ReaderOpen = 0;
        else
            state->Shared->WriterOpen = 0;
        _ReadWriteBarrier();
        KiPipeRefreshEventsLocked(state);
    }
    KeReleaseSpinlock(&state->Lock, flags);
    KiPipeStateRelease(state);
}

// Allocates one endpoint object for existing shared state.
static IO_PIPE_ENDPOINT* KiPipeCreateEndpoint(IO_PIPE_STATE* state, IO_PIPE_DIRECTION direction) {
    IO_PIPE_ENDPOINT* endpoint =
        (IO_PIPE_ENDPOINT*)ObAllocateObject(ObjectTypePipeEndpoint, sizeof(IO_PIPE_ENDPOINT));
    if (!endpoint)
        return (IO_PIPE_ENDPOINT*)0;
    endpoint->Header.Destroy = KiPipeEndpointDestroy;
    endpoint->State = state;
    endpoint->Direction = direction;
    return endpoint;
}

// Allocates the physical pages and stable kernel view for one pipe data plane.
static KSTATUS KiPipeCreateBacking(IO_PIPE_STATE* state) {
    size_t payloadSize = sizeof(SYS_PIPE_SHARED_STATE);
    size_t pageCount = (payloadSize + PAGE_SIZE - 1u) / PAGE_SIZE;
    state->Frames = (PHYSICAL_ADDRESS*)MmAllocatePoolTag(pageCount * sizeof(PHYSICAL_ADDRESS),
                                                         IO_PIPE_FRAME_TAG);
    if (!state->Frames)
        return NTLSTATUS_NO_MEMORY;
    memset(state->Frames, 0, pageCount * sizeof(PHYSICAL_ADDRESS));
    state->PageCount = pageCount;

    for (size_t i = 0; i < pageCount; ++i) {
        PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseAnonymous);
        if (!frame)
            return NTLSTATUS_NO_MEMORY;
        state->Frames[i] = frame;
        void* page = MmPfnToVirtual(frame);
        if (page)
            memset(page, 0, PAGE_SIZE);
    }

    KSTATUS status = MmCreateKernelViewFromFrames(state->Frames, state->PageCount,
                                                   &state->KernelView);
    if (!KSUCCESS(status))
        return status;
    state->Shared = (SYS_PIPE_SHARED_STATE*)MmKernelViewBase(state->KernelView);
    if (!state->Shared)
        return NTLSTATUS_FAILURE;

    state->Shared->Magic = SYS_PIPE_SHARED_MAGIC;
    state->Shared->Version = SYS_PIPE_SHARED_VERSION;
    state->Shared->Capacity = SYS_PIPE_BUFFER_CAPACITY;
    state->Shared->HeaderBytes = SYS_PIPE_SHARED_HEADER_BYTES;
    state->Shared->ReaderOpen = 1;
    state->Shared->WriterOpen = 1;
    return NTLSTATUS_SUCCESS;
}

// Creates one anonymous pipe with kernel-owned endpoints and a user-mapped data plane.
KSTATUS IoCreatePipe(IO_PIPE_ENDPOINT** readEndOut, IO_PIPE_ENDPOINT** writeEndOut) {
    if (!readEndOut || !writeEndOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *readEndOut = (IO_PIPE_ENDPOINT*)0;
    *writeEndOut = (IO_PIPE_ENDPOINT*)0;

    IO_PIPE_STATE* state = (IO_PIPE_STATE*)MmAllocatePoolTag(sizeof(*state), IO_PIPE_STATE_TAG);
    if (!state)
        return NTLSTATUS_NO_MEMORY;
    memset(state, 0, sizeof(*state));
    state->RefCount = 2;
    state->Lock = (SPINLOCK)SPINLOCK_INIT;

    KSTATUS status = KeEventCreate(EventManualReset, 0, &state->ReadEvent);
    if (KSUCCESS(status))
        status = KeEventCreate(EventManualReset, 1, &state->WriteEvent);
    if (KSUCCESS(status))
        status = KiPipeCreateBacking(state);
    if (!KSUCCESS(status)) {
        state->RefCount = 1;
        KiPipeStateRelease(state);
        return status;
    }

    IO_PIPE_ENDPOINT* readEnd = KiPipeCreateEndpoint(state, IoPipeReadEnd);
    IO_PIPE_ENDPOINT* writeEnd = KiPipeCreateEndpoint(state, IoPipeWriteEnd);
    if (!readEnd || !writeEnd) {
        if (readEnd)
            ObDereferenceObject(&readEnd->Header);
        else
            KiPipeStateRelease(state);
        if (writeEnd)
            ObDereferenceObject(&writeEnd->Header);
        else
            KiPipeStateRelease(state);
        return NTLSTATUS_NO_MEMORY;
    }

    *readEndOut = readEnd;
    *writeEndOut = writeEnd;
    return NTLSTATUS_SUCCESS;
}

// Maps the shared pipe data plane into one process-owned VAD.
KSTATUS IoMapPipe(IO_PIPE_ENDPOINT* endpoint, PROCESS* process, void** sharedStateOut) {
    if (!endpoint || !process || !sharedStateOut || !endpoint->State)
        return NTLSTATUS_INVALID_PARAMETER;
    *sharedStateOut = (void*)0;

    IO_PIPE_STATE* state = endpoint->State;
    KSTATUS status = MmVadCreateSharedViewFromFrames(process, state->Frames, state->PageCount,
                                                      sizeof(SYS_PIPE_SHARED_STATE),
                                                      PAGE_PRESENT | PAGE_USER | PAGE_WRITE | PAGE_NX,
                                                      MM_VAD_PIPE, MmVadShared, sharedStateOut);
    if (!KSUCCESS(status))
        return status;

    status = MmVadProtect(process, *sharedStateOut, SYS_PIPE_SHARED_HEADER_BYTES,
                          PAGE_PRESENT | PAGE_USER | PAGE_NX);
    if (!KSUCCESS(status)) {
        (void)MmVadDestroy(process, *sharedStateOut);
        *sharedStateOut = (void*)0;
    }
    return status;
}

// Refreshes readiness after a user-mode ring transition.
KSTATUS IoNotifyPipe(IO_PIPE_ENDPOINT* endpoint) {
    if (!endpoint || !endpoint->State)
        return NTLSTATUS_INVALID_PARAMETER;
    IO_PIPE_STATE* state = endpoint->State;
    uintptr_t flags = KeAcquireSpinlock(&state->Lock);
    KiPipeRefreshEventsLocked(state);
    KeReleaseSpinlock(&state->Lock, flags);
    return NTLSTATUS_SUCCESS;
}

// Queries readiness for poll-style events.
KSTATUS IoPollPipe(IO_PIPE_ENDPOINT* endpoint, uint32_t events, uint32_t* returnedEventsOut) {
    if (!endpoint || !returnedEventsOut || !endpoint->State || !endpoint->State->Shared)
        return NTLSTATUS_INVALID_PARAMETER;

    SYS_PIPE_SHARED_STATE* shared = endpoint->State->Shared;
    uint64_t used = 0;
    if (!KiPipeUsedBytes(shared, &used))
        return NTLSTATUS_FAILURE;

    uint32_t returned = 0;
    if (endpoint->Direction == IoPipeReadEnd) {
        if ((events & SYS_POLL_READ) && (used != 0 || shared->WriterOpen == 0))
            returned |= SYS_POLL_READ;
        if (shared->WriterOpen == 0)
            returned |= SYS_POLL_HANGUP;
    } else {
        if (shared->ReaderOpen == 0) {
            returned |= SYS_POLL_ERROR | SYS_POLL_HANGUP;
        } else if ((events & SYS_POLL_WRITE) && used < SYS_PIPE_BUFFER_CAPACITY) {
            returned |= SYS_POLL_WRITE;
        }
    }
    *returnedEventsOut = returned;
    return NTLSTATUS_SUCCESS;
}

// Returns the endpoint event used by generic waits.
EVENT* IoPipeWaitEvent(IO_PIPE_ENDPOINT* endpoint) {
    if (!endpoint || !endpoint->State)
        return (EVENT*)0;
    return endpoint->Direction == IoPipeReadEnd
         ? endpoint->State->ReadEvent
         : endpoint->State->WriteEvent;
}

// Publishes one kernel waiter before entering the generic event wait engine.
void IoPipeBeginWait(IO_PIPE_ENDPOINT* endpoint) {
    if (!endpoint || !endpoint->State || !endpoint->State->Shared)
        return;
    SYS_PIPE_SHARED_STATE* shared = endpoint->State->Shared;
    if (endpoint->Direction == IoPipeReadEnd)
        NtlAtomicIncrement32(&shared->ReadWaiters);
    else
        NtlAtomicIncrement32(&shared->WriteWaiters);

    uintptr_t flags = KeAcquireSpinlock(&endpoint->State->Lock);
    KiPipeRefreshEventsLocked(endpoint->State);
    KeReleaseSpinlock(&endpoint->State->Lock, flags);
}

// Withdraws one kernel waiter after the generic event wait completes.
void IoPipeEndWait(IO_PIPE_ENDPOINT* endpoint) {
    if (!endpoint || !endpoint->State || !endpoint->State->Shared)
        return;
    SYS_PIPE_SHARED_STATE* shared = endpoint->State->Shared;
    if (endpoint->Direction == IoPipeReadEnd)
        NtlAtomicDecrement32(&shared->ReadWaiters);
    else
        NtlAtomicDecrement32(&shared->WriteWaiters);
}

// Returns whether this endpoint is the read side of the pipe.
bool IoPipeIsReadEnd(IO_PIPE_ENDPOINT* endpoint) {
    return endpoint && endpoint->Direction == IoPipeReadEnd;
}