// gdi32.c
#include <stdint.h>
#include <stddef.h>
#include <stdlib.h>
#include <ntlite/syscall.h>

#if defined(__SSE2__) || defined(_M_X64) || defined(_M_IX86)
#include <emmintrin.h>
#elif defined(__ARM_NEON) || defined(__ARM_NEON__) || defined(_M_ARM64)
#include <arm_neon.h>
#endif

// ---------------------------------------------------------------------------
// Win32 Base Types
// ---------------------------------------------------------------------------
typedef void* HDC;
typedef unsigned long  DWORD;
typedef unsigned short WORD;
typedef unsigned int   UINT;
typedef long           LONG;
typedef int            BOOL;

#define TRUE  1
#define FALSE 0

// StretchDIBits Constants
#define BI_RGB          0L
#define DIB_RGB_COLORS  0
#define SRCCOPY         (DWORD)0x00CC0020

// ---------------------------------------------------------------------------
// GDI Structs
// ---------------------------------------------------------------------------
typedef struct tagBITMAPINFOHEADER {
    DWORD biSize;
    LONG  biWidth;
    LONG  biHeight;
    WORD  biPlanes;
    WORD  biBitCount;
    DWORD biCompression;
    DWORD biSizeImage;
    LONG  biXPelsPerMeter;
    LONG  biYPelsPerMeter;
    DWORD biClrUsed;
    DWORD biClrImportant;
} BITMAPINFOHEADER;

typedef struct tagRGBQUAD {
    unsigned char rgbBlue;
    unsigned char rgbGreen;
    unsigned char rgbRed;
    unsigned char rgbReserved;
} RGBQUAD;

typedef struct tagBITMAPINFO {
    BITMAPINFOHEADER bmiHeader;
    RGBQUAD          bmiColors[1];
} BITMAPINFO;


// ---------------------------------------------------------------------------
// GDI Exports
// ---------------------------------------------------------------------------

// Static buffers to persist between function calls
static uint32_t* destBuffer = NULL;
static int* srcX_map = NULL;
static int currentDestCapacity = 0;
static int currentXMapCapacity = 0;

__declspec(dllexport) int StretchDIBits(
    HDC hdc,
    int xDest, int yDest, int DestWidth, int DestHeight,
    int xSrc, int ySrc, int SrcWidth, int SrcHeight,
    const void *lpBits, const BITMAPINFO *lpbmi,
    UINT iUsage, DWORD rop) 
{
    (void)hdc; (void)lpbmi; (void)iUsage; (void)rop;

    // Safety check to prevent division by zero
    if (DestWidth <= 0 || DestHeight <= 0 || SrcWidth <= 0 || SrcHeight <= 0) {
        return 0; 
    }

    // Handle trivial case (1:1 copy) to bypass allocation and scaling overhead
    if (DestWidth == SrcWidth && DestHeight == SrcHeight && xSrc == 0 && ySrc == 0) {
        SysBlitScreen(xDest, yDest, DestWidth, DestHeight, (const uint32_t*)lpBits);
        return DestHeight;
    }

    const uint32_t* srcPixels = (const uint32_t*)lpBits;

    // Only reallocate if the new size is LARGER than what we already have
    int requiredPixels = DestWidth * DestHeight;
    if (requiredPixels > currentDestCapacity) {
        free(destBuffer);
        destBuffer = (uint32_t*)malloc(requiredPixels * sizeof(uint32_t));
        currentDestCapacity = requiredPixels;
    }

    if (DestWidth > currentXMapCapacity) {
        free(srcX_map);
        srcX_map = (int*)malloc(DestWidth * sizeof(int));
        currentXMapCapacity = DestWidth;
    }

    for (int dx = 0; dx < DestWidth; ++dx) {
        // Nearest neighbor coordinate mapping
        srcX_map[dx] = xSrc + (dx * SrcWidth) / DestWidth;
        
        // Clamp to safe bounds
        if (srcX_map[dx] < 0) srcX_map[dx] = 0;
        if (srcX_map[dx] >= SrcWidth) srcX_map[dx] = SrcWidth - 1;
    }

    // Process the image row by row
    for (int dy = 0; dy < DestHeight; ++dy) {
        // Map destination Y to source Y
        int sy = ySrc + (dy * SrcHeight) / DestHeight;
        if (sy < 0) sy = 0;
        if (sy >= SrcHeight) sy = SrcHeight - 1;

        // Assumes stride == SrcWidth (tightly packed 32bpp)
        const uint32_t* srcRow = srcPixels + (sy * SrcWidth); 
        uint32_t* dstRow = destBuffer + (dy * DestWidth);

        int dx = 0;

#if defined(__SSE2__) || defined(_M_X64)
        // Process 4 pixels (16 bytes) per loop iteration using SSE2
        for (; dx <= DestWidth - 4; dx += 4) {
            // Gather 4 source pixels using our precomputed X map.
            // Note: _mm_set_epi32 takes arguments in reverse order (3, 2, 1, 0)
            __m128i pixels = _mm_set_epi32(
                srcRow[srcX_map[dx + 3]],
                srcRow[srcX_map[dx + 2]],
                srcRow[srcX_map[dx + 1]],
                srcRow[srcX_map[dx + 0]]
            );
            // Issue a single 128-bit unaligned store
            _mm_storeu_si128((__m128i*)&dstRow[dx], pixels);
        }
#elif defined(__ARM_NEON) || defined(__ARM_NEON__) || defined(_M_ARM64)
        // Process 4 pixels per loop iteration using ARM NEON
        for (; dx <= DestWidth - 4; dx += 4) {
            // Initialize a vector and populate lanes (avoids stack-memory hop)
            uint32x4_t pixels = vdupq_n_u32(0);
            pixels = vsetq_lane_u32(srcRow[srcX_map[dx + 0]], pixels, 0);
            pixels = vsetq_lane_u32(srcRow[srcX_map[dx + 1]], pixels, 1);
            pixels = vsetq_lane_u32(srcRow[srcX_map[dx + 2]], pixels, 2);
            pixels = vsetq_lane_u32(srcRow[srcX_map[dx + 3]], pixels, 3);
            
            // Issue a single 128-bit store
            vst1q_u32(&dstRow[dx], pixels);
        }
#endif

        // Epilogue: Handle remaining pixels if DestWidth isn't a multiple of 4
        for (; dx < DestWidth; ++dx) {
            dstRow[dx] = srcRow[srcX_map[dx]];
        }
    }

    // Direct mapping to kernel screen blit using the stretched buffer
    SysBlitScreen(xDest, yDest, DestWidth, DestHeight, destBuffer);

    return DestHeight; 
}