/*
 * Unit test: DRM display subsystem
 * Tests connector detection, CRTC, plane enumeration, and framebuffer creation.
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <poll.h>
#include <linux/dma-heap.h>
#include <xf86drm.h>
#include <xf86drmMode.h>
#include <drm_fourcc.h>

static int dma_heap_alloc(const char *heap_name, size_t size, int *fd) {
    char path[64];
    snprintf(path, sizeof(path), "/dev/dma_heap/%s", heap_name);
    int heap_fd = open(path, O_RDWR | O_CLOEXEC);
    if (heap_fd < 0) return -1;
    struct dma_heap_allocation_data data = { .len = size, .fd_flags = O_RDWR | O_CLOEXEC };
    int ret = ioctl(heap_fd, DMA_HEAP_IOCTL_ALLOC, &data);
    close(heap_fd);
    if (ret < 0) return -1;
    *fd = data.fd;
    return 0;
}

static int create_fb(int drm_fd, int buf_fd, uint32_t w, uint32_t h,
                     uint32_t stride, uint32_t format, uint32_t *fb_id) {
    uint32_t handles[4] = {0}, pitches[4] = {0}, offsets[4] = {0};
    if (drmPrimeFDToHandle(drm_fd, buf_fd, &handles[0])) {
        fprintf(stderr, "  drmPrimeFDToHandle failed\n");
        return -1;
    }
    if (format == DRM_FORMAT_NV12) {
        handles[1] = handles[0];
        pitches[0] = stride;
        pitches[1] = stride;
        offsets[0] = 0;
        offsets[1] = stride * h;
    } else {
        pitches[0] = stride;
        offsets[0] = 0;
    }
    int ret = drmModeAddFB2(drm_fd, w, h, format, handles, pitches, offsets, fb_id, 0);
    if (ret)
        fprintf(stderr, "  drmModeAddFB2(%ux%u fmt=0x%x stride=%u) failed: %s\n",
                w, h, format, stride, strerror(errno));
    return ret;
}

static void page_flip_handler(int fd, unsigned int seq,
                              unsigned int tv_sec, unsigned int tv_usec,
                              void *user_data) {
    (void)fd; (void)seq; (void)tv_sec; (void)tv_usec; (void)user_data;
}

int main(void) {
    printf("=== DRM Display Unit Test ===\n\n");

    int drm_fd = open("/dev/dri/card0", O_RDWR);
    if (drm_fd < 0) {
        perror("open /dev/dri/card0");
        return 1;
    }
    printf("Opened /dev/dri/card0 fd=%d\n\n", drm_fd);

    /* Try to become DRM master */
    if (drmSetMaster(drm_fd)) {
        printf("drmSetMaster failed: %s (errno=%d)\n", strerror(errno), errno);
    } else {
        printf("Acquired DRM master\n");
    }

    /* Set universal planes capability */
    if (drmSetClientCap(drm_fd, DRM_CLIENT_CAP_UNIVERSAL_PLANES, 1))
        printf("UNIVERSAL_PLANES cap failed: %s\n", strerror(errno));
    else
        printf("Set UNIVERSAL_PLANES cap\n");

    if (drmSetClientCap(drm_fd, DRM_CLIENT_CAP_ATOMIC, 1))
        printf("ATOMIC cap failed: %s\n", strerror(errno));
    else
        printf("Set ATOMIC cap\n");

    /* --- Enumerate connectors --- */
    drmModeRes *res = drmModeGetResources(drm_fd);
    if (!res) {
        fprintf(stderr, "drmModeGetResources failed\n");
        return 1;
    }

    printf("Connectors: %d\n", res->count_connectors);
    uint32_t conn_id = 0, crtc_id = 0;
    for (int i = 0; i < res->count_connectors; i++) {
        drmModeConnector *conn = drmModeGetConnector(drm_fd, res->connectors[i]);
        if (!conn) continue;
        printf("  [%d] connector %u: %s, %d modes, %d encoders\n",
               i, conn->connector_id,
               conn->connection == DRM_MODE_CONNECTED ? "CONNECTED" : "disconnected",
               conn->count_modes, conn->count_encoders);
        if (conn->connection == DRM_MODE_CONNECTED && conn->count_modes > 0 && !conn_id) {
            conn_id = conn->connector_id;
            printf("  => Selected connector %u, mode: %ux%u\n",
                   conn_id, conn->modes[0].hdisplay, conn->modes[0].vdisplay);
        }
        drmModeFreeConnector(conn);
    }

    /* --- Enumerate CRTCs --- */
    printf("\nCRTCs: %d\n", res->count_crtcs);
    for (int i = 0; i < res->count_crtcs; i++) {
        drmModeCrtc *crtc = drmModeGetCrtc(drm_fd, res->crtcs[i]);
        if (!crtc) continue;
        printf("  [%d] crtc %u: fb=%u %ux%u\n",
               i, crtc->crtc_id, crtc->buffer_id,
               crtc->width, crtc->height);
        drmModeFreeCrtc(crtc);
    }

    /* --- Find encoder/CRTC for our connector --- */
    if (conn_id) {
        drmModeConnector *conn = drmModeGetConnector(drm_fd, conn_id);
        if (conn) {
            if (conn->encoder_id) {
                drmModeEncoder *enc = drmModeGetEncoder(drm_fd, conn->encoder_id);
                if (enc) {
                    crtc_id = enc->crtc_id;
                    printf("\nEncoder %u => crtc %u\n", enc->encoder_id, crtc_id);
                    drmModeFreeEncoder(enc);
                }
            }
            if (!crtc_id && conn->count_encoders > 0) {
                drmModeEncoder *enc = drmModeGetEncoder(drm_fd, conn->encoders[0]);
                if (enc) {
                    crtc_id = enc->crtc_id;
                    printf("\nFallback encoder %u => crtc %u\n", enc->encoder_id, crtc_id);
                    drmModeFreeEncoder(enc);
                }
            }
            drmModeFreeConnector(conn);
        }
    }
    if (!crtc_id && res->count_crtcs > 0) {
        crtc_id = res->crtcs[0];
        printf("\nFallback: using first crtc %u\n", crtc_id);
    }
    printf("Selected crtc: %u\n", crtc_id);

    /* --- Enumerate all planes --- */
    drmModePlaneRes *planes = drmModeGetPlaneResources(drm_fd);
    printf("\nPlanes: %d\n", planes ? planes->count_planes : 0);

    int crtc_index = -1;
    for (int i = 0; i < res->count_crtcs; i++) {
        if (res->crtcs[i] == crtc_id) { crtc_index = i; break; }
    }
    printf("CRTC %u is at index %d\n", crtc_id, crtc_index);

    uint32_t overlay_plane_id = 0;

    if (planes) {
        for (uint32_t i = 0; i < planes->count_planes; i++) {
            drmModePlane *p = drmModeGetPlane(drm_fd, planes->planes[i]);
            if (!p) continue;

            uint32_t plane_id = planes->planes[i];
            int compatible = (p->possible_crtcs & (1u << crtc_index)) != 0;

            drmModeObjectProperties *props =
                drmModeObjectGetProperties(drm_fd, plane_id, DRM_MODE_OBJECT_PLANE);

            uint64_t type_val = 0;
            uint64_t in_fmt_blob = 0;
            if (props) {
                for (uint32_t j = 0; j < props->count_props; j++) {
                    drmModePropertyRes *prop = drmModeGetProperty(drm_fd, props->props[j]);
                    if (!prop) continue;
                    if (strcmp(prop->name, "type") == 0)
                        type_val = props->prop_values[j];
                    if (strcmp(prop->name, "IN_FORMATS") == 0)
                        in_fmt_blob = props->prop_values[j];
                    drmModeFreeProperty(prop);
                }
            }

            int has_nv12 = 0;
            if (in_fmt_blob) {
                drmModePropertyBlobRes *blob = drmModeGetPropertyBlob(drm_fd, in_fmt_blob);
                if (blob && blob->length >= 4) {
                    uint32_t *data = (uint32_t *)blob->data;
                    uint32_t num_u32 = blob->length / 4;
                    /* Dump first 16 u32 values */
                    printf("  plane %u IN_FORMATS blob: %zu bytes (%u u32s)\n",
                           plane_id, (size_t)blob->length, num_u32);
                    for (uint32_t k = 0; k < (num_u32 < 16 ? num_u32 : 16); k++) {
                        uint32_t v = data[k];
                        char cc[5] = {(char)(v & 0xff), (char)((v >> 8) & 0xff),
                                      (char)((v >> 16) & 0xff), (char)((v >> 24) & 0xff), 0};
                        printf("    [%u] = 0x%08x (%s)\n", k, v, cc);
                    }
                    /* The blob format is: count_formats(u32), count_modifiers(u32),
                       formats[count_formats](u32), then modifier entries.
                       But we might have a different format. Try scanning all u32s. */
                    for (uint32_t k = 0; k < num_u32; k++) {
                        if (data[k] == DRM_FORMAT_NV12) { has_nv12 = 1; break; }
                    }
                }
                if (blob) drmModeFreePropertyBlob(blob);
            }

            const char *type_str = "Unknown";
            if (type_val == 0) type_str = "Overlay(0)";
            else if (type_val == 1) type_str = "Primary(1)";
            else if (type_val == 2) type_str = "Cursor(2)";

            printf("  plane %u: type=%s compat=%d nv12=%d possible_crtcs=0x%x\n",
                   plane_id, type_str, compatible, has_nv12, p->possible_crtcs);

            /* Select overlay plane with NV12 support */
            if (type_val == 2 && has_nv12 && compatible && !overlay_plane_id) {
                overlay_plane_id = plane_id;
                printf("  => Selected overlay plane %u (NV12, compatible)\n", plane_id);
            }

            if (props) drmModeFreeObjectProperties(props);
            drmModeFreePlane(p);
        }

        /* If no compatible overlay, try any NV12 overlay */
        if (!overlay_plane_id) {
            printf("\nNo compatible NV12 overlay. Trying all NV12 overlays:\n");
            for (uint32_t i = 0; i < planes->count_planes; i++) {
                uint32_t plane_id = planes->planes[i];
                drmModeObjectProperties *props =
                    drmModeObjectGetProperties(drm_fd, plane_id, DRM_MODE_OBJECT_PLANE);
                if (!props) continue;
                uint64_t type_val = 0, in_fmt_blob = 0;
                for (uint32_t j = 0; j < props->count_props; j++) {
                    drmModePropertyRes *prop = drmModeGetProperty(drm_fd, props->props[j]);
                    if (!prop) continue;
                    if (strcmp(prop->name, "type") == 0) type_val = props->prop_values[j];
                    if (strcmp(prop->name, "IN_FORMATS") == 0) in_fmt_blob = props->prop_values[j];
                    drmModeFreeProperty(prop);
                }
                int has_nv12 = 0;
                if (in_fmt_blob) {
                    drmModePropertyBlobRes *blob = drmModeGetPropertyBlob(drm_fd, in_fmt_blob);
                    if (blob && blob->length > 8) {
                        uint32_t *data = (uint32_t *)blob->data;
                        for (uint32_t f = 0; f < data[0]; f++) {
                            if (data[2 + f] == DRM_FORMAT_NV12) { has_nv12 = 1; break; }
                        }
                    }
                    if (blob) drmModeFreePropertyBlob(blob);
                }
                if (type_val == 2 && has_nv12 && !overlay_plane_id) {
                    overlay_plane_id = plane_id;
                    printf("  => Using plane %u (may not be compatible)\n", plane_id);
                }
                drmModeFreeObjectProperties(props);
            }
        }

        drmModeFreePlaneResources(planes);
    }

    printf("\nOverlay plane for display: %u\n", overlay_plane_id);

    /* --- Allocate DMA buffer and test FB creation --- */
    printf("\n--- FB creation test ---\n");
    int buf_fd = -1;
    size_t buf_sz = 1920 * 1080 * 3 / 2;
    if (dma_heap_alloc("reserved", buf_sz, &buf_fd) < 0) {
        if (dma_heap_alloc("system-uncached", buf_sz, &buf_fd) < 0) {
            fprintf(stderr, "Cannot allocate DMA buffer\n");
            return 1;
        }
    }
    void *map = mmap(NULL, buf_sz, PROT_READ | PROT_WRITE, MAP_SHARED, buf_fd, 0);
    if (map != MAP_FAILED) {
        memset(map, 0x10, 1920 * 1080);         /* Y = black */
        memset((uint8_t*)map + 1920 * 1080, 0x80, buf_sz - 1920 * 1080); /* UV */
    }

    uint32_t fb_id = 0;
    int ret;

    /* Test 1: drmModeSetCrtc on primary plane (expected to fail with NV12) */
    printf("\nTest 1: drmModeSetCrtc with NV12...\n");
    ret = create_fb(drm_fd, buf_fd, 1920, 1080, 1920, DRM_FORMAT_NV12, &fb_id);
    if (ret == 0) {
        drmModeConnector *tc = drmModeGetConnector(drm_fd, conn_id);
        drmModeModeInfo mode_info;
        if (tc && tc->count_modes > 0) mode_info = tc->modes[0];
        else memset(&mode_info, 0, sizeof(mode_info));
        if (tc) drmModeFreeConnector(tc);

        ret = drmModeSetCrtc(drm_fd, crtc_id, fb_id, 0, 0, &conn_id, 1, &mode_info);
        if (ret) {
            printf("  FAIL: drmModeSetCrtc: %s (errno=%d)\n", strerror(errno), errno);
        } else {
            printf("  OK: drmModeSetCrtc succeeded\n");
        }
        drmModeRmFB(drm_fd, fb_id);
        fb_id = 0;
    }

    /* Test 2: Try Esmart1 (plane 107) on CRTC 121 via drmModeSetPlane */
    printf("\nTest 2: Force Esmart1 plane 107 on CRTC %u (ignoring possible_crtcs)...\n", crtc_id);
    fb_id = 0;
    ret = create_fb(drm_fd, buf_fd, 1920, 1080, 1920, DRM_FORMAT_NV12, &fb_id);
    if (ret == 0) {
        ret = drmModeSetPlane(drm_fd, 107, crtc_id, fb_id,
                              0, 0, 0, 1920, 1080,
                              0, 0, 1920 << 16, 1080 << 16);
        if (ret) {
            printf("  FAIL: drmModeSetPlane(107): %s (errno=%d)\n", strerror(errno), errno);
        } else {
            printf("  OK: Esmart1 plane 107 on CRTC %u succeeded!\n", crtc_id);
            sleep(2);
        }
        drmModeRmFB(drm_fd, fb_id);
        fb_id = 0;
    }

    /* Test 2b: Try atomic modesetting to force plane 107 on CRTC 121 */
    printf("\nTest 2b: Atomic modesetting - force plane 107 on CRTC %u...\n", crtc_id);
    fb_id = 0;
    ret = create_fb(drm_fd, buf_fd, 1920, 1080, 1920, DRM_FORMAT_NV12, &fb_id);
    if (ret == 0) {
        drmModeAtomicReq *req = drmModeAtomicAlloc();
        if (req) {
            /* Set the FB_ID on the plane */
            drmModeAtomicAddProperty(req, 107,
                /* We need property IDs, not names for atomic. Find them first. */
                0, fb_id);

            /* Find property IDs for plane 107 */
            drmModeObjectProperties *props =
                drmModeObjectGetProperties(drm_fd, 107, DRM_MODE_OBJECT_PLANE);
            if (props) {
                uint32_t fb_prop_id = 0, crtc_prop_id = 0, src_prop_id = 0, dst_prop_id = 0;
                uint32_t src_x_prop = 0, src_y_prop = 0;
                for (uint32_t j = 0; j < props->count_props; j++) {
                    drmModePropertyRes *prop = drmModeGetProperty(drm_fd, props->props[j]);
                    if (!prop) continue;
                    if (strcmp(prop->name, "FB_ID") == 0) fb_prop_id = props->props[j];
                    if (strcmp(prop->name, "CRTC_ID") == 0) crtc_prop_id = props->props[j];
                    if (strcmp(prop->name, "SRC_X") == 0) src_x_prop = props->props[j];
                    if (strcmp(prop->name, "SRC_Y") == 0) src_y_prop = props->props[j];
                    if (strcmp(prop->name, "SRC_W") == 0) src_prop_id = props->props[j];
                    if (strcmp(prop->name, "SRC_H") == 0) dst_prop_id = props->props[j];
                    drmModeFreeProperty(prop);
                }
                printf("  plane 107 props: FB_ID=%u CRTC_ID=%u\n", fb_prop_id, crtc_prop_id);

                if (fb_prop_id && crtc_prop_id) {
                    drmModeAtomicFree(req);
                    req = drmModeAtomicAlloc();
                    drmModeAtomicAddProperty(req, 107, fb_prop_id, fb_id);
                    drmModeAtomicAddProperty(req, 107, crtc_prop_id, crtc_id);

                    /* Set source/dest rectangles via SRC_W/SRC_H/DST_W/DST_H properties */
                    uint32_t dst_w_prop = 0, dst_h_prop = 0;
                    for (uint32_t j = 0; j < props->count_props; j++) {
                        drmModePropertyRes *prop = drmModeGetProperty(drm_fd, props->props[j]);
                        if (!prop) continue;
                        if (strcmp(prop->name, "DST_W") == 0) dst_w_prop = props->props[j];
                        if (strcmp(prop->name, "DST_H") == 0) dst_h_prop = props->props[j];
                        drmModeFreeProperty(prop);
                    }
                    if (src_x_prop) drmModeAtomicAddProperty(req, 107, src_x_prop, 0);
                    if (src_y_prop) drmModeAtomicAddProperty(req, 107, src_y_prop, 0);
                    if (src_prop_id) drmModeAtomicAddProperty(req, 107, src_prop_id, 1920 << 16);
                    if (dst_prop_id) drmModeAtomicAddProperty(req, 107, dst_prop_id, 1080 << 16);
                    if (dst_w_prop) drmModeAtomicAddProperty(req, 107, dst_w_prop, 1920);
                    if (dst_h_prop) drmModeAtomicAddProperty(req, 107, dst_h_prop, 1080);

                    uint32_t cursor_x_prop = 0, cursor_y_prop = 0;
                    for (uint32_t j = 0; j < props->count_props; j++) {
                        drmModePropertyRes *prop = drmModeGetProperty(drm_fd, props->props[j]);
                        if (!prop) continue;
                        if (strcmp(prop->name, "CRTC_X") == 0) cursor_x_prop = props->props[j];
                        if (strcmp(prop->name, "CRTC_Y") == 0) cursor_y_prop = props->props[j];
                        drmModeFreeProperty(prop);
                    }
                    if (cursor_x_prop) drmModeAtomicAddProperty(req, 107, cursor_x_prop, 0);
                    if (cursor_y_prop) drmModeAtomicAddProperty(req, 107, cursor_y_prop, 0);

                    printf("  Submitting atomic commit...\n");
                    ret = drmModeAtomicCommit(drm_fd, req,
                        DRM_MODE_ATOMIC_ALLOW_MODESET, NULL);
                    if (ret) {
                        printf("  FAIL: atomic commit: %s (errno=%d)\n", strerror(errno), errno);
                    } else {
                        printf("  OK: atomic commit succeeded!\n");
                        sleep(2);
                    }
                }
                drmModeFreeObjectProperties(props);
            }
            drmModeAtomicFree(req);
        }
        drmModeRmFB(drm_fd, fb_id);
        fb_id = 0;
    }

    /* Test 3: Use primary plane with XRGB8888 (expected to work) */
    printf("\nTest 3: drmModeSetCrtc with XRGB8888...\n");
    int xrgb_buf_fd = -1;
    {
        uint32_t handles2[4] = {0}, pitches2[4] = {0}, offsets2[4] = {0};
        size_t xrgb_sz = 1920 * 1080 * 4;
        if (dma_heap_alloc("system-uncached", xrgb_sz, &xrgb_buf_fd) == 0) {
            void *xrgb_map = mmap(NULL, xrgb_sz, PROT_READ | PROT_WRITE, MAP_SHARED, xrgb_buf_fd, 0);
            if (xrgb_map != MAP_FAILED) {
                memset(xrgb_map, 0x00, xrgb_sz);
                munmap(xrgb_map, xrgb_sz);
            }
            drmPrimeFDToHandle(drm_fd, xrgb_buf_fd, &handles2[0]);
            pitches2[0] = 1920 * 4;
            offsets2[0] = 0;
            ret = drmModeAddFB2(drm_fd, 1920, 1080, DRM_FORMAT_XRGB8888,
                                handles2, pitches2, offsets2, &fb_id, 0);
            if (ret) {
                printf("  FAIL: drmModeAddFB2 XRGB8888: %s\n", strerror(errno));
            } else {
                drmModeConnector *tc2 = drmModeGetConnector(drm_fd, conn_id);
                drmModeModeInfo mode2;
                if (tc2 && tc2->count_modes > 0) mode2 = tc2->modes[0];
                else memset(&mode2, 0, sizeof(mode2));
                if (tc2) drmModeFreeConnector(tc2);

                ret = drmModeSetCrtc(drm_fd, crtc_id, fb_id, 0, 0, &conn_id, 1, &mode2);
                if (ret) {
                    printf("  FAIL: drmModeSetCrtc XRGB8888: %s\n", strerror(errno));
                } else {
                    printf("  OK: XRGB8888 modeset works!\n");
                }
                drmModeRmFB(drm_fd, fb_id);
                fb_id = 0;
            }
            close(xrgb_buf_fd);
        }
    }

    /* Test 4: Page flip test */
    if (overlay_plane_id) {
        printf("\nTest 4: Page flip via overlay plane...\n");
        fb_id = 0;
        ret = create_fb(drm_fd, buf_fd, 1920, 1080, 1920, DRM_FORMAT_NV12, &fb_id);
        if (ret == 0) {
            ret = drmModeSetPlane(drm_fd, overlay_plane_id, crtc_id, fb_id,
                                  0, 0, 0, 1920, 1080,
                                  0, 0, 1920 << 16, 1080 << 16);
            printf("  drmModeSetPlane: %s\n", ret ? "FAIL" : "OK");
            if (!ret) {
                sleep(1);
                drmEventContext ev = {.version = 2, .page_flip_handler = page_flip_handler};
                drmHandleEvent(drm_fd, &ev);
                printf("  OK: page flip handler returned\n");
            }
            drmModeRmFB(drm_fd, fb_id);
        }
    }

    if (map != MAP_FAILED) munmap(map, buf_sz);
    close(buf_fd);
    drmModeFreeResources(res);
    drmDropMaster(drm_fd);
    close(drm_fd);

    printf("\n=== Done ===\n");
    return 0;
}