/* capture.c — coherent pull of a sub-rectangle of the CURRENT frame.
 *
 * A persistent capture window over a fixed rect of the captured surface. open()
 * allocates a stable buffer once (its address is fixed across syncs); sync()
 * re-fills it all-or-nothing from the current latest frame under one seqlock
 * window — the exact mirror of vgpup_sample_frame's discipline, but for a
 * sub-rectangle, and WITHOUT frame_id dedup (a repeat sync of the same frame
 * still yields the coherent rect — see vgpu_perception.h).
 *
 * Every read goes through gva_read under c->proc_cr3 into the stable buffer (or
 * a per-window scratch for the RGB24 conversion); no gva_ptr is ever held across
 * the seqlock window. NOT zero-copy by design (the rect crosses pages and must
 * survive the seq re-check). All GVA arithmetic is done in uint64_t.
 *
 * The window caches proc_cr3 + ring_gva (copied from the region handle at open);
 * the producer block lives at region_gva == ring_gva - VGPU_RING_OFFSET, where
 * latest / seq[] / desc[] are read (the same fields sample.c reads via region_gva).
 */

#include "perception-internal.h"

#include <stdlib.h>

/* Round `v` up to the next multiple of `a` (a is a power of two). */
static inline size_t vgpup_align_up(size_t v, size_t a)
{
    return (v + (a - 1u)) & ~(a - 1u);
}

vgpup_capture* vgpup_capture_open(vgpup_region* r, vgpup_rect rect,
                                  int layout, int format)
{
    vgpup_capture* c;
    uint32_t out_bpp, row_bytes;

    /* validate args and enums; fit against the ABI ceilings (NOT the current
     * desc — open may precede the first frame; per-frame fit is a sync check).
     * x+w / y+h computed in uint64_t to avoid overflow. */
    if (!r) { return NULL; }
    if (rect.w == 0u || rect.h == 0u) { return NULL; }
    if (layout != VGPUP_LAYOUT_FLAT && layout != VGPUP_LAYOUT_ROWS) { return NULL; }
    if (format != VGPUP_PX_BGRA && format != VGPUP_PX_RGB24) { return NULL; }
    if ((uint64_t)rect.x + rect.w > VGPU_MAX_WIDTH)  { return NULL; }
    if ((uint64_t)rect.y + rect.h > VGPU_MAX_HEIGHT) { return NULL; }

    out_bpp   = (format == VGPUP_PX_RGB24) ? 3u : 4u;
    row_bytes = rect.w * out_bpp;

    c = calloc(1, sizeof *c);
    if (!c) { return NULL; }

    c->proc_cr3  = r->proc_cr3;
    c->ring_gva  = r->ring_gva;
    c->x         = rect.x;
    c->y         = rect.y;
    c->w         = rect.w;
    c->h         = rect.h;
    c->out_bpp   = out_bpp;
    c->row_bytes = row_bytes;
    c->layout    = layout;
    c->format    = format;

    /* RGB24 stages each source BGRA row in a per-window scratch (the source row
     * is too large for the stack: up to VGPU_MAX_WIDTH*4). BGRA reads straight
     * into the destination row, no scratch needed. */
    if (format == VGPUP_PX_RGB24) {
        c->scratch = malloc((size_t)VGPU_MAX_WIDTH * 4u);
        if (!c->scratch) { vgpup_capture_close(c); return NULL; }
    }

    if (layout == VGPUP_LAYOUT_FLAT) {
        /* one contiguous block, rows back-to-back (sub-stride == row_bytes) */
        c->buf_bytes = (size_t)row_bytes * rect.h;
        c->buf = calloc(1, c->buf_bytes);
        if (!c->buf) { vgpup_capture_close(c); return NULL; }
    } else {
        /* ROWS: ONE backing block sliced into h cache-line-aligned sub-blocks +
         * an array of h pointers. One allocation (less fragmentation, addresses
         * stable since the backing is never reallocated); the padding between
         * row_bytes..row_stride is never filled and is excluded from info->bytes. */
        size_t   row_stride = vgpup_align_up(row_bytes, VGPUP_ROW_ALIGN);
        uint32_t i;
        c->buf_bytes = row_stride * rect.h;
        c->buf = calloc(1, c->buf_bytes);
        if (!c->buf) { vgpup_capture_close(c); return NULL; }
        c->rows = calloc(rect.h, sizeof *c->rows);
        if (!c->rows) { vgpup_capture_close(c); return NULL; }
        for (i = 0; i < rect.h; ++i) {
            c->rows[i] = c->buf + (size_t)i * row_stride;
        }
    }

    return c;
}

/* Convert w source BGRA pixels in `src` (B,G,R,A) to RGB24 (R,G,B) in `dst`. */
static inline void vgpup_bgra_to_rgb24(const uint8_t* src, uint8_t* dst, uint32_t w)
{
    uint32_t i;
    for (i = 0; i < w; ++i) {
        dst[3u * i + 0u] = src[4u * i + 2u];   /* R */
        dst[3u * i + 1u] = src[4u * i + 1u];   /* G */
        dst[3u * i + 2u] = src[4u * i + 0u];   /* B */
    }
}

int vgpup_capture_sync(vgpup_capture* c, vmie_mem* m, vgpup_capture_info* info)
{
    unsigned attempt;
    uint64_t region_gva;

    if (!c || !m || !info) { return -1; }

    /* the producer block sits one ring-offset below the cached ring base */
    region_gva = c->ring_gva - VGPU_RING_OFFSET;

    for (attempt = 0; attempt < VGPUP_SEQLOCK_RETRIES; ++attempt) {
        uint32_t    latest = 0, seq_before = 0, seq_after = 0;
        vgpu_desc_t d;
        uint64_t    seq_gva, desc_gva, slot_base, src_bytes;
        uint32_t    row;

        /* latest (acquire-equivalent: its own read) */
        if (vgpup_read_field(m, c->proc_cr3, region_gva,
                             offsetof(vgpu_producer_t, latest),
                             &latest, sizeof latest) < 0) {
            return -1;
        }
        if (latest == VGPU_LATEST_NONE || latest >= VGPU_SLOT_COUNT) {
            return 0;   /* no frame yet */
        }

        seq_gva  = region_gva + offsetof(vgpu_producer_t, seq)
                   + (uint64_t)latest * sizeof(uint32_t);
        desc_gva = region_gva + offsetof(vgpu_producer_t, desc)
                   + (uint64_t)latest * sizeof(vgpu_desc_t);

        if (vgpup_read_seq(m, c->proc_cr3, seq_gva, &seq_before) < 0) { return -1; }
        if (vgpup_seq_is_writing(seq_before)) { continue; }   /* writer in slot */

        if (gva_read(m, (uintptr_t)c->proc_cr3, (uintptr_t)desc_gva,
                     &d, sizeof d) < 0) {
            return -1;
        }

        /* descriptor sanity within the read window (tight BGRA, bounded dims) */
        if (d.format != VGPU_FMT_BGRA8888 || d.stride != d.width * 4u ||
            d.width == 0u || d.width > VGPU_MAX_WIDTH ||
            d.height == 0u || d.height > VGPU_MAX_HEIGHT) {
            continue;   /* likely a torn read; retry */
        }

        /* NO frame_id dedup (intentional): a repeat sync of the same frame must
         * still return the coherent rect (see vgpu_perception.h). */

        /* fit against THIS frame: rect must lie within the current dimensions.
         * A frame that became smaller is a strict lossy skip (no partial read). */
        if ((uint64_t)c->x + c->w > d.width || (uint64_t)c->y + c->h > d.height) {
            return 0;
        }

        /* fill the whole rect, row by row, from the latest slot under proc_cr3 */
        slot_base = c->ring_gva + (uint64_t)latest * VGPU_SLOT_STRIDE;
        src_bytes = (uint64_t)c->w * 4u;   /* source is BGRA, 4 bpp */
        for (row = 0; row < c->h; ++row) {
            uint64_t src_gva = slot_base
                             + (uint64_t)(c->y + row) * d.stride
                             + (uint64_t)c->x * 4u;
            uint8_t* dst = (c->layout == VGPUP_LAYOUT_ROWS)
                         ? (uint8_t*)c->rows[row]
                         : c->buf + (size_t)row * c->row_bytes;

            if (c->format == VGPUP_PX_RGB24) {
                /* stage the source BGRA row, then convert into the dst row */
                if (gva_read(m, (uintptr_t)c->proc_cr3, (uintptr_t)src_gva,
                             c->scratch, src_bytes) < 0) {
                    return -1;
                }
                vgpup_bgra_to_rgb24(c->scratch, dst, c->w);
            } else {
                /* BGRA verbatim: read straight into the dst row */
                if (gva_read(m, (uintptr_t)c->proc_cr3, (uintptr_t)src_gva,
                             dst, src_bytes) < 0) {
                    return -1;
                }
            }
        }

        /* re-check the slot seq ONCE after the whole rect is filled: unchanged
         * and still even → the rect is one coherent snapshot (all-or-nothing). */
        if (vgpup_read_seq(m, c->proc_cr3, seq_gva, &seq_after) < 0) { return -1; }
        if (seq_after != seq_before || vgpup_seq_is_writing(seq_after)) {
            continue;   /* the slot was rewritten under us — re-fill the rect */
        }

        info->desc.width        = d.width;
        info->desc.height       = d.height;
        info->desc.stride       = d.stride;
        info->desc.format       = d.format;
        info->desc.frame_id     = d.frame_id;
        info->desc.timestamp_ns = d.timestamp_ns;
        info->frame_id          = d.frame_id;
        info->bytes             = (size_t)c->row_bytes * c->h;   /* payload only */
        return 1;
    }
    return 0;   /* writer kept the slot busy past the retry limit — lossy skip */
}

uint32_t vgpup_capture_out_bpp(const vgpup_capture* c)
{
    return c ? c->out_bpp : 0u;
}

uint32_t vgpup_capture_row_bytes(const vgpup_capture* c)
{
    return c ? c->row_bytes : 0u;
}

uint32_t vgpup_capture_row_count(const vgpup_capture* c)
{
    return c ? c->h : 0u;
}

const uint8_t* vgpup_capture_mem(const vgpup_capture* c)
{
    if (!c || c->layout != VGPUP_LAYOUT_FLAT) { return NULL; }
    return c->buf;
}

size_t vgpup_capture_bytes(const vgpup_capture* c)
{
    if (!c || c->layout != VGPUP_LAYOUT_FLAT) { return 0u; }
    return (size_t)c->row_bytes * c->h;
}

const uint8_t* const* vgpup_capture_rows(const vgpup_capture* c)
{
    if (!c || c->layout != VGPUP_LAYOUT_ROWS) { return NULL; }
    return c->rows;
}

void vgpup_capture_close(vgpup_capture* c)
{
    if (!c) { return; }
    free(c->rows);
    free(c->scratch);
    free(c->buf);
    free(c);
}