// Kernel surface storage. This file owns shared backing frames, mappings, commit generations, and privileged presentation.

#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <kernel/ke_surface.h>
#include <kernel/ke_process.h>
#include <kernel/ke_mutex.h>
#include <kernel/ke_event.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_display.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_pool.h>
#include <ntlite/syscall.h>

typedef struct KE_SURFACE {
    OBJECT Header;
    uint32_t Width;
    uint32_t Height;
    uint32_t PitchBytes;
    uint32_t Format;
    uint32_t Flags;
    size_t PayloadSize;
    size_t PageCount;
    bool StatsRegistered;
    bool LastDamageValid;
    uint16_t Reserved;
    uint32_t Reserved2;
    uint64_t MapCount;
    uint64_t ContentGeneration;
    uint64_t CommitCount;
    uint64_t CommitRectCount;
    uint64_t CommitClippedRectCount;
    uint64_t CommitEmptyClipCount;
    DISPLAY_RECT LastDamageBounds;
    PHYSICAL_ADDRESS* Frames;
    MM_KERNEL_VIEW* KernelView;
    void* KernelBase;
    EVENT* CommitEvent;
    MUTEX* Lock;
} KE_SURFACE;

static SPINLOCK gSurfaceStatsLock = SPINLOCK_INIT;
static SYS_SURFACE_STATS gSurfaceStats;

// Multiplies two sizes while rejecting overflow.
static bool KiSurfaceMulSize(size_t a, size_t b, size_t* out) {
    if (!out || (a != 0 && b > SIZE_MAX / a)) return false;
    *out = a * b;
    return true;
}

// Aligns one size upward while rejecting wraparound.
static size_t KiSurfaceAlignUp(size_t value, size_t alignment) {
    if (!alignment || (alignment & (alignment - 1u))) return 0;
    size_t mask = alignment - 1u;
    if (value > SIZE_MAX - mask) return 0;
    return (value + mask) & ~mask;
}

// Converts a generic object reference to a surface only when its type is correct.
static KE_SURFACE* KiSurfaceFromObject(OBJECT* object) {
    if (!object || object->Type != ObjectTypeSurface) return (KE_SURFACE*)0;
    return CONTAINING_RECORD(object, KE_SURFACE, Header);
}

// Registers a surface in the global telemetry snapshot.
static void KiSurfaceStatsRegister(KE_SURFACE* surface) {
    if (!surface || surface->StatsRegistered) return;
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    gSurfaceStats.CreatedSurfaces++;
    gSurfaceStats.LiveSurfaces++;
    gSurfaceStats.LivePages += surface->PageCount;
    gSurfaceStats.LiveBytes += surface->PayloadSize;
    surface->StatsRegistered = true;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}

// Removes a surface from the global telemetry snapshot.
static void KiSurfaceStatsUnregister(KE_SURFACE* surface) {
    if (!surface || !surface->StatsRegistered) return;
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    gSurfaceStats.DestroyedSurfaces++;
    if (gSurfaceStats.LiveSurfaces) gSurfaceStats.LiveSurfaces--;
    gSurfaceStats.LivePages = gSurfaceStats.LivePages >= surface->PageCount
        ? gSurfaceStats.LivePages - surface->PageCount : 0;
    gSurfaceStats.LiveBytes = gSurfaceStats.LiveBytes >= surface->PayloadSize
        ? gSurfaceStats.LiveBytes - surface->PayloadSize : 0;
    surface->StatsRegistered = false;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}

// Records one successful surface mapping operation.
static void KiSurfaceStatsRecordMap(KE_SURFACE* surface) {
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    surface->MapCount++;
    gSurfaceStats.MapCalls++;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}

// Records one commit after rectangle validation and clipping.
static void KiSurfaceStatsRecordCommit(uint64_t inputRects, uint64_t clippedRects, uint64_t emptyClips) {
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    gSurfaceStats.CommitCalls++;
    gSurfaceStats.CommitRectInputs += inputRects;
    gSurfaceStats.CommitClippedRects += clippedRects;
    gSurfaceStats.CommitEmptyClips += emptyClips;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}

// Records one presentation performed through DisplayManager authority.
static void KiSurfaceStatsRecordPresent(uint64_t bytesCopied) {
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    gSurfaceStats.PhysicalPresentCalls++;
    gSurfaceStats.PhysicalPresentBytesCopied += bytesCopied;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}

// Releases surface storage. User VADs retain their own backing-frame references until unmapped.
static void KiSurfaceDestroy(OBJECT* object) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    if (!surface) return;

    KiSurfaceStatsUnregister(surface);
    if (surface->KernelView) {
        MmDestroyKernelView(surface->KernelView);
        surface->KernelView = (MM_KERNEL_VIEW*)0;
        surface->KernelBase = (void*)0;
    }
    if (surface->CommitEvent) {
        ObDereferenceObject(&surface->CommitEvent->Header);
        surface->CommitEvent = (EVENT*)0;
    }
    if (surface->Lock) {
        ObDereferenceObject(&surface->Lock->Header);
        surface->Lock = (MUTEX*)0;
    }
    if (surface->Frames) {
        for (size_t i = 0; i < surface->PageCount; ++i)
            if (surface->Frames[i]) (void)MmPfnRelease(surface->Frames[i]);
        MmFreePool(surface->Frames);
        surface->Frames = (PHYSICAL_ADDRESS*)0;
    }
}

// Clips one rectangle to the surface extent.
static bool KiSurfaceClipRect(const KE_SURFACE* surface, const DISPLAY_RECT* input, DISPLAY_RECT* output) {
    if (!surface || !output) return false;
    DISPLAY_RECT r;
    if (input) {
        r = *input;
    } else {
        r.X = 0;
        r.Y = 0;
        r.W = (int32_t)surface->Width;
        r.H = (int32_t)surface->Height;
    }
    if (r.W <= 0 || r.H <= 0) return false;

    int64_t x0 = r.X;
    int64_t y0 = r.Y;
    int64_t x1 = x0 + r.W;
    int64_t y1 = y0 + r.H;
    if (x0 < 0) x0 = 0;
    if (y0 < 0) y0 = 0;
    if (x1 > surface->Width) x1 = surface->Width;
    if (y1 > surface->Height) y1 = surface->Height;
    if (x0 >= x1 || y0 >= y1) return false;

    output->X = (int32_t)x0;
    output->Y = (int32_t)y0;
    output->W = (int32_t)(x1 - x0);
    output->H = (int32_t)(y1 - y0);
    return true;
}

// Expands one damage bound to include another clipped rectangle.
static void KiSurfaceUnionRect(DISPLAY_RECT* bounds, const DISPLAY_RECT* rect) {
    int32_t x0 = bounds->X < rect->X ? bounds->X : rect->X;
    int32_t y0 = bounds->Y < rect->Y ? bounds->Y : rect->Y;
    int32_t bx1 = bounds->X + bounds->W;
    int32_t by1 = bounds->Y + bounds->H;
    int32_t rx1 = rect->X + rect->W;
    int32_t ry1 = rect->Y + rect->H;
    int32_t x1 = bx1 > rx1 ? bx1 : rx1;
    int32_t y1 = by1 > ry1 ? by1 : ry1;
    bounds->X = x0;
    bounds->Y = y0;
    bounds->W = x1 - x0;
    bounds->H = y1 - y0;
}

// Clips one kernel-captured rectangle list. A zero count means the whole surface.
static KSTATUS KiSurfaceCaptureRects(KE_SURFACE* surface, const DISPLAY_RECT* rects, uint32_t rectCount,
                                     DISPLAY_RECT* clippedOut, uint32_t* clippedCountOut,
                                     uint64_t* emptyOut) {
    if (!surface || !clippedOut || !clippedCountOut || !emptyOut)
        return NTLSTATUS_INVALID_PARAMETER;
    if (rectCount > SYS_SURFACE_MAX_RECTS || (rectCount && !rects))
        return NTLSTATUS_INVALID_PARAMETER;

    DISPLAY_RECT whole = {0, 0, (int32_t)surface->Width, (int32_t)surface->Height};
    if (!rectCount) {
        rects = &whole;
        rectCount = 1;
    }

    uint32_t clippedCount = 0;
    uint64_t empty = 0;
    for (uint32_t i = 0; i < rectCount; ++i) {
        DISPLAY_RECT clipped;
        if (!KiSurfaceClipRect(surface, &rects[i], &clipped)) {
            empty++;
            continue;
        }
        clippedOut[clippedCount++] = clipped;
    }
    *clippedCountOut = clippedCount;
    *emptyOut = empty;
    return NTLSTATUS_SUCCESS;
}

// Creates one surface storage object and returns its initial object reference.
KSTATUS KeSurfaceCreate(uint32_t width, uint32_t height, uint32_t format, uint32_t flags,
                        OBJECT** surfaceOut) {
    if (!surfaceOut) return NTLSTATUS_INVALID_PARAMETER;
    *surfaceOut = (OBJECT*)0;
    if (!width || !height || width > SYS_SURFACE_MAX_DIMENSION || height > SYS_SURFACE_MAX_DIMENSION)
        return NTLSTATUS_INVALID_PARAMETER;
    if (format != SYS_SURFACE_FORMAT_XRGB8888 || flags != SYS_SURFACE_FLAG_NONE)
        return NTLSTATUS_NOT_SUPPORTED;

    size_t rowBytes = 0;
    if (!KiSurfaceMulSize((size_t)width, sizeof(uint32_t), &rowBytes))
        return NTLSTATUS_INVALID_PARAMETER;
    size_t pitchBytes = KiSurfaceAlignUp(rowBytes, 64u);
    if (!pitchBytes || pitchBytes < rowBytes) return NTLSTATUS_INVALID_PARAMETER;

    size_t payloadSize = 0;
    if (!KiSurfaceMulSize(pitchBytes, (size_t)height, &payloadSize) || payloadSize > SYS_SURFACE_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;
    size_t pageCount = (payloadSize + PAGE_SIZE - 1u) / PAGE_SIZE;
    if (!pageCount || pageCount > SIZE_MAX / sizeof(PHYSICAL_ADDRESS))
        return NTLSTATUS_INVALID_PARAMETER;

    KE_SURFACE* surface = (KE_SURFACE*)ObAllocateObject(ObjectTypeSurface, sizeof(KE_SURFACE));
    if (!surface) return NTLSTATUS_NO_MEMORY;
    surface->Header.Destroy = KiSurfaceDestroy;
    surface->Width = width;
    surface->Height = height;
    surface->PitchBytes = (uint32_t)pitchBytes;
    surface->Format = format;
    surface->Flags = flags;
    surface->PayloadSize = payloadSize;
    surface->PageCount = pageCount;

    surface->Frames = (PHYSICAL_ADDRESS*)MmAllocatePoolTag(pageCount * sizeof(PHYSICAL_ADDRESS),
                                                            MM_POOL_TAG('s', 'f', 'r', 'm'));
    if (surface->Frames)
        memset(surface->Frames, 0, pageCount * sizeof(PHYSICAL_ADDRESS));

    KSTATUS status = KeMutexCreate(&surface->Lock);
    if (KSUCCESS(status)) status = KeEventCreate(EventAutoReset, 0, &surface->CommitEvent);
    if (!surface->Frames || !KSUCCESS(status)) {
        ObDereferenceObject(&surface->Header);
        return !surface->Frames ? NTLSTATUS_NO_MEMORY : status;
    }
    for (size_t i = 0; i < pageCount; ++i) {
        PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseSurface);
        if (!frame) {
            ObDereferenceObject(&surface->Header);
            return NTLSTATUS_NO_MEMORY;
        }
        surface->Frames[i] = frame;
        void* kva = MmPfnToVirtual(frame);
        if (kva) memset(kva, 0, PAGE_SIZE);
    }

    status = MmCreateKernelViewFromFrames(surface->Frames, surface->PageCount, &surface->KernelView);
    if (!KSUCCESS(status)) {
        ObDereferenceObject(&surface->Header);
        return status;
    }
    surface->KernelBase = MmKernelViewBase(surface->KernelView);
    KiSurfaceStatsRegister(surface);
    *surfaceOut = &surface->Header;
    return NTLSTATUS_SUCCESS;
}

// Maps one surface into the calling process's VAD namespace with attenuated page permissions.
KSTATUS KeSurfaceMap(PROCESS* process, OBJECT* object, uint32_t flags,
                     void** baseOut, uint32_t* pitchBytesOut) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    if (!process || !surface || !baseOut || !pitchBytesOut || !(flags & SYS_SURFACE_MAP_READ) ||
        (flags & ~(SYS_SURFACE_MAP_READ | SYS_SURFACE_MAP_WRITE)))
        return NTLSTATUS_INVALID_PARAMETER;

    uint64_t pageFlags = PAGE_PRESENT | PAGE_USER | PAGE_NX;
    if (flags & SYS_SURFACE_MAP_WRITE) pageFlags |= PAGE_WRITE;
    KSTATUS status = MmVadCreateSharedViewFromFrames(process, surface->Frames, surface->PageCount,
                                                      surface->PayloadSize, pageFlags,
                                                      MM_VAD_SURFACE, MmVadShared, baseOut);
    if (KSUCCESS(status)) {
        *pitchBytesOut = surface->PitchBytes;
        KiSurfaceStatsRecordMap(surface);
    }
    return status;
}

// Removes one mapping only when the supplied VAD is a surface shared view owned by the process.
KSTATUS KeSurfaceUnmap(PROCESS* process, void* base) {
    if (!process || !base) return NTLSTATUS_INVALID_PARAMETER;
    void* vadBase = (void*)0;
    size_t vadSize = 0;
    uint64_t pageFlags = 0;
    MM_VAD_TYPE type = MmVadGeneric;
    uint32_t vadFlags = 0;
    KSTATUS status = MmVadQuery(process, base, &vadBase, &vadSize, &pageFlags, &type, &vadFlags);
    (void)vadSize;
    (void)pageFlags;
    if (!KSUCCESS(status)) return status;
    if (vadBase != base || type != MmVadShared || !(vadFlags & MM_VAD_SURFACE))
        return NTLSTATUS_INVALID_PARAMETER;
    return MmVadDestroy(process, base);
}

// Snapshots one surface's storage and commit metadata.
KSTATUS KeSurfaceQuery(OBJECT* object, SYS_SURFACE_DESC* descOut) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    if (!surface || !descOut) return NTLSTATUS_INVALID_PARAMETER;
    memset(descOut, 0, sizeof(*descOut));

    KSTATUS status = KeAcquireMutex(surface->Lock, KE_WAIT_INFINITE);
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(status)) return status;
    descOut->Width = surface->Width;
    descOut->Height = surface->Height;
    descOut->PitchBytes = surface->PitchBytes;
    descOut->Format = surface->Format;
    descOut->PayloadBytes = surface->PayloadSize;
    descOut->PageCount = surface->PageCount;
    descOut->ContentGeneration = surface->ContentGeneration;
    descOut->CommitCount = surface->CommitCount;
    descOut->CommitRectCount = surface->CommitRectCount;
    descOut->CommitClippedRectCount = surface->CommitClippedRectCount;
    descOut->CommitEmptyClipCount = surface->CommitEmptyClipCount;
    descOut->LastDamageBounds = surface->LastDamageBounds;
    descOut->LastDamageValid = surface->LastDamageValid;
    KeReleaseMutex(surface->Lock);

    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    descOut->MapCount = surface->MapCount;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
    return NTLSTATUS_SUCCESS;
}

// Announces changed pixels without touching physical scanout.
KSTATUS KeSurfaceCommit(OBJECT* object, const DISPLAY_RECT* rects, uint32_t rectCount,
                        uint64_t* generationOut) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    if (!surface) return NTLSTATUS_INVALID_PARAMETER;

    uint32_t requestedCount = rectCount ? rectCount : 1u;
    DISPLAY_RECT clipped[SYS_SURFACE_MAX_RECTS];
    uint32_t clippedCount = 0;
    uint64_t emptyCount = 0;
    KSTATUS status = KiSurfaceCaptureRects(surface, rects, rectCount, clipped, &clippedCount, &emptyCount);
    if (!KSUCCESS(status)) return status;

    status = KeAcquireMutex(surface->Lock, KE_WAIT_INFINITE);
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(status)) return status;
    surface->ContentGeneration++;
    surface->CommitCount++;
    surface->CommitRectCount += requestedCount;
    surface->CommitClippedRectCount += clippedCount;
    surface->CommitEmptyClipCount += emptyCount;
    surface->LastDamageValid = clippedCount ? 1u : 0u;
    if (clippedCount) {
        surface->LastDamageBounds = clipped[0];
        for (uint32_t i = 1; i < clippedCount; ++i)
            KiSurfaceUnionRect(&surface->LastDamageBounds, &clipped[i]);
    } else {
        memset(&surface->LastDamageBounds, 0, sizeof(surface->LastDamageBounds));
    }
    uint64_t generation = surface->ContentGeneration;
    KeReleaseMutex(surface->Lock);

    KiSurfaceStatsRecordCommit(requestedCount, clippedCount, emptyCount);
    KeSignalEvent(surface->CommitEvent);
    if (generationOut) *generationOut = generation;
    return NTLSTATUS_SUCCESS;
}

// Presents one surface after a caller has separately proved DisplayManager PRESENT authority.
KSTATUS KeSurfacePresent(OBJECT* object, const DISPLAY_RECT* rects, uint32_t rectCount,
                         uint64_t* bytesCopiedOut) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    if (!surface || !bytesCopiedOut) return NTLSTATUS_INVALID_PARAMETER;

    DISPLAY_RECT clipped[SYS_SURFACE_MAX_RECTS];
    uint32_t clippedCount = 0;
    uint64_t emptyCount = 0;
    KSTATUS status = KiSurfaceCaptureRects(surface, rects, rectCount, clipped, &clippedCount, &emptyCount);
    (void)emptyCount;
    if (!KSUCCESS(status)) return status;
    if (!surface->KernelBase) return NTLSTATUS_DEVICE_ERROR;

    const uint8_t* sourceBase = (const uint8_t*)surface->KernelBase;
    uint64_t bytesCopied = 0;
    for (uint32_t i = 0; i < clippedCount; ++i) {
        const DISPLAY_RECT* rect = &clipped[i];
        const uint32_t* source = (const uint32_t*)(sourceBase +
            (size_t)(uint32_t)rect->Y * surface->PitchBytes +
            (size_t)(uint32_t)rect->X * sizeof(uint32_t));
        status = KeDisplayPresentPitched(rect->X, rect->Y, rect->W, rect->H, source, surface->PitchBytes);
        if (!KSUCCESS(status)) return status;
        bytesCopied += (uint64_t)(uint32_t)rect->W * (uint64_t)(uint32_t)rect->H * sizeof(uint32_t);
    }

    *bytesCopiedOut = bytesCopied;
    KiSurfaceStatsRecordPresent(bytesCopied);
    return NTLSTATUS_SUCCESS;
}

// Returns the passive event used to wait for commit-generation changes.
EVENT* KeSurfaceWaitEvent(OBJECT* object) {
    KE_SURFACE* surface = KiSurfaceFromObject(object);
    return surface ? surface->CommitEvent : (EVENT*)0;
}

// Snapshots global surface telemetry.
void KeSurfaceQueryStats(SYS_SURFACE_STATS* statsOut) {
    if (!statsOut) return;
    uintptr_t irq = KeAcquireSpinlock(&gSurfaceStatsLock);
    *statsOut = gSurfaceStats;
    statsOut->Size = sizeof(*statsOut);
    statsOut->Version = SYS_STATS_VERSION;
    KeReleaseSpinlock(&gSurfaceStatsLock, irq);
}