// SPDX-FileCopyrightText: 2025 KylinSoft Co., Ltd.
//
// SPDX-License-Identifier: GPL-1.0-or-later

#include <dlfcn.h>
#include <stdlib.h>
#include <unistd.h>

#include <drm_fourcc.h>
#include <gbm.h>

#include <wlr/interfaces/wlr_buffer.h>
#include <wlr/render/drm_format_set.h>

#include <kywc/log.h>

#include "render/opengl.h"
#include "renderer_p.h"

struct gbm_buffer {
    struct wlr_buffer base;
    struct wl_list link;
    struct gbm_bo *gbm_bo;
    struct wlr_dmabuf_attributes dmabuf;
};

struct gbm_allocator {
    struct gbm_device *gbm_device;
    struct wl_list buffers;

    struct wlr_drm_format_set formats;

    int (*get_format_modifier_plane_count)(struct gbm_device *gbm, uint32_t format,
                                           uint64_t modifier);
    struct gbm_bo *(*create_with_modifiers2)(struct gbm_device *gbm, uint32_t width,
                                             uint32_t height, uint32_t format,
                                             const uint64_t *modifiers, const unsigned int count,
                                             uint32_t flags);
    int (*get_fd_for_plane)(struct gbm_bo *bo, int plane);

    struct wlr_renderer *renderer;
    struct wl_listener destroy;
};

static struct gbm_allocator *allocator = NULL;

const struct wlr_drm_format_set *
gbm_allocator_get_single_plane_formates(struct wlr_renderer *renderer)
{
    if (!allocator || allocator->renderer != renderer || allocator->formats.len == 0) {
        return NULL;
    }
    return &allocator->formats;
}

static struct gbm_buffer *get_gbm_buffer_from_buffer(struct wlr_buffer *wlr_buffer)
{
    struct gbm_buffer *buffer = wl_container_of(wlr_buffer, buffer, base);
    return buffer;
}

static void buffer_destroy(struct wlr_buffer *wlr_buffer)
{
    struct gbm_buffer *buffer = get_gbm_buffer_from_buffer(wlr_buffer);
    wlr_dmabuf_attributes_finish(&buffer->dmabuf);
    if (buffer->gbm_bo) {
        gbm_bo_destroy(buffer->gbm_bo);
    }
    wl_list_remove(&buffer->link);
    free(buffer);
}

static bool buffer_get_dmabuf(struct wlr_buffer *wlr_buffer, struct wlr_dmabuf_attributes *attribs)
{
    struct gbm_buffer *buffer = get_gbm_buffer_from_buffer(wlr_buffer);
    *attribs = buffer->dmabuf;
    return true;
}

static const struct wlr_buffer_impl buffer_impl = {
    .destroy = buffer_destroy,
    .get_dmabuf = buffer_get_dmabuf,
};

static bool drm_format_has(const struct wlr_drm_format *fmt, uint64_t modifier)
{
    for (size_t i = 0; i < fmt->len; ++i) {
        if (fmt->modifiers[i] == modifier) {
            return true;
        }
    }
    return false;
}

static bool export_gbm_bo(struct gbm_bo *bo, struct wlr_dmabuf_attributes *out)
{
    struct wlr_dmabuf_attributes attribs = { 0 };
    attribs.n_planes = gbm_bo_get_plane_count(bo);
    if (attribs.n_planes > WLR_DMABUF_MAX_PLANES) {
        kywc_log(KYWC_ERROR, "GBM BO contains too many planes (%d)", attribs.n_planes);
        return false;
    }

    attribs.width = gbm_bo_get_width(bo);
    attribs.height = gbm_bo_get_height(bo);
    attribs.format = gbm_bo_get_format(bo);
    attribs.modifier = gbm_bo_get_modifier(bo);

    int i;
    int32_t handle = -1;

    for (i = 0; i < attribs.n_planes; ++i) {
        if (allocator->get_fd_for_plane != NULL) {
            attribs.fd[i] = allocator->get_fd_for_plane(bo, i);
            if (attribs.fd[i] < 0) {
                kywc_log(KYWC_ERROR, "gbm_bo_get_fd_for_plane failed");
                goto error_fd;
            }
        } else {
            union gbm_bo_handle plane_handle = gbm_bo_get_handle_for_plane(bo, i);
            if (plane_handle.s32 < 0) {
                kywc_log(KYWC_ERROR, "gbm_bo_get_handle_for_plane failed");
                goto error_fd;
            }
            if (i == 0) {
                handle = plane_handle.s32;
            } else if (plane_handle.s32 != handle) {
                kywc_log(KYWC_ERROR, "Failed to export GBM BO: "
                                     "all planes don't have the same GEM handle");
                goto error_fd;
            }

            attribs.fd[i] = gbm_bo_get_fd(bo);
            if (attribs.fd[i] < 0) {
                kywc_log(KYWC_ERROR, "gbm_bo_get_fd failed");
                goto error_fd;
            }
        }

        attribs.offset[i] = gbm_bo_get_offset(bo, i);
        attribs.stride[i] = gbm_bo_get_stride_for_plane(bo, i);
    }

    *out = attribs;
    return true;

error_fd:
    for (int j = 0; j < i; ++j) {
        close(attribs.fd[j]);
    }
    return false;
}

struct wlr_buffer *gbm_allocator_create_buffer(int width, int height,
                                               const struct wlr_drm_format *format)
{
    if (!allocator) {
        return NULL;
    }

    struct gbm_device *gbm_device = allocator->gbm_device;
    uint64_t fallback_modifier = DRM_FORMAT_MOD_INVALID;
    uint32_t usage = GBM_BO_USE_RENDERING;
    bool has_modifier = true;
    struct gbm_bo *bo = NULL;

    if (allocator->create_with_modifiers2) {
        bo = allocator->create_with_modifiers2(gbm_device, width, height, format->format,
                                               format->modifiers, format->len, usage);
    }

    if (!bo) {
        if (format->len == 1 && format->modifiers[0] == DRM_FORMAT_MOD_LINEAR) {
            usage |= GBM_BO_USE_LINEAR;
            fallback_modifier = DRM_FORMAT_MOD_LINEAR;
        } else if (!drm_format_has(format, DRM_FORMAT_MOD_INVALID)) {
            // If the format doesn't accept an implicit modifier, bail out.
            kywc_log(KYWC_ERROR, "gbm_bo_create_with_modifiers failed");
            return NULL;
        }
        bo = gbm_bo_create(gbm_device, width, height, format->format, usage);
        has_modifier = false;
    }

    if (!bo) {
        kywc_log(KYWC_ERROR, "gbm_bo_create failed");
        return NULL;
    }

    struct gbm_buffer *buffer = calloc(1, sizeof(*buffer));
    if (!buffer) {
        gbm_bo_destroy(bo);
        return NULL;
    }

    wlr_buffer_init(&buffer->base, &buffer_impl, width, height);
    buffer->gbm_bo = bo;

    if (!export_gbm_bo(bo, &buffer->dmabuf)) {
        free(buffer);
        gbm_bo_destroy(bo);
        return NULL;
    }

    if (!has_modifier) {
        buffer->dmabuf.modifier = fallback_modifier;
    }

    wl_list_insert(&allocator->buffers, &buffer->link);

    kywc_log(KYWC_DEBUG,
             "Allocated %dx%d GBM buffer "
             "with format (0x%08" PRIX32 "), modifier (0x%016" PRIX64 ")",
             buffer->base.width, buffer->base.height, buffer->dmabuf.format,
             buffer->dmabuf.modifier);

    return &buffer->base;
}

static void allocator_query_single_plane_formats(struct gbm_allocator *allocator)
{
    if (wlr_renderer_is_opengl(allocator->renderer)) {
        struct ky_opengl_renderer *r = ky_opengl_renderer_from_wlr_renderer(allocator->renderer);
        // modifiers is disabled, no need to query the formats
        if (!r->egl->has_modifiers) {
            return;
        }
    }

    if (!allocator->renderer->impl->get_render_formats) {
        return;
    }

    const struct wlr_drm_format_set *render_formats =
        allocator->renderer->impl->get_render_formats(allocator->renderer);
    if (!render_formats) {
        return;
    }

    allocator->get_format_modifier_plane_count =
        dlsym(RTLD_DEFAULT, "gbm_device_get_format_modifier_plane_count");
    if (allocator->get_format_modifier_plane_count == NULL) {
        kywc_log(KYWC_WARN, "No gbm_device_get_format_modifier_plane_count support!");
    }

    const struct wlr_drm_format *fmt;
    for (size_t i = 0; i < render_formats->len; i++) {
        fmt = &render_formats->formats[i];
        for (size_t j = 0; j < fmt->len; ++j) {
            if (fmt->modifiers[j] == DRM_FORMAT_MOD_INVALID ||
                fmt->modifiers[j] == DRM_FORMAT_MOD_LINEAR) {
                wlr_drm_format_set_add(&allocator->formats, fmt->format, fmt->modifiers[j]);
                continue;
            }
            if (allocator->get_format_modifier_plane_count &&
                allocator->get_format_modifier_plane_count(allocator->gbm_device, fmt->format,
                                                           fmt->modifiers[j]) != 1) {
                continue;
            }
            // treat as single plane if no get_format_modifier_plane_count
            wlr_drm_format_set_add(&allocator->formats, fmt->format, fmt->modifiers[j]);
        }
    }
}

static void handle_renderer_destroy(struct wl_listener *listener, void *data)
{
    struct gbm_allocator *allocator = wl_container_of(listener, allocator, destroy);
    wl_list_remove(&allocator->destroy.link);

    struct gbm_buffer *buf, *buf_tmp;
    wl_list_for_each_safe(buf, buf_tmp, &allocator->buffers, link) {
        gbm_bo_destroy(buf->gbm_bo);
        buf->gbm_bo = NULL;
        wl_list_remove(&buf->link);
        wl_list_init(&buf->link);
    }

    gbm_device_destroy(allocator->gbm_device);
    wlr_drm_format_set_finish(&allocator->formats);

    free(allocator);
    allocator = NULL;
}

bool gbm_allocator_create(struct wlr_renderer *renderer)
{
    int drm_fd = wlr_renderer_get_drm_fd(renderer);
    if (drm_fd < 0) {
        return false;
    }

    allocator = calloc(1, sizeof(*allocator));
    if (!allocator) {
        return false;
    }

    struct gbm_device *gbm_device = gbm_create_device(drm_fd);
    if (!gbm_device) {
        free(allocator);
        allocator = NULL;
        return false;
    }

    wl_list_init(&allocator->buffers);
    allocator->gbm_device = gbm_device;

    allocator->renderer = renderer;
    allocator->destroy.notify = handle_renderer_destroy;
    wl_signal_add(&renderer->events.destroy, &allocator->destroy);

    allocator_query_single_plane_formats(allocator);

    allocator->create_with_modifiers2 = dlsym(RTLD_DEFAULT, "gbm_bo_create_with_modifiers2");
    if (allocator->create_with_modifiers2 == NULL) {
        kywc_log(KYWC_WARN, "No gbm_bo_create_with_modifiers2 support!");
    }
    allocator->get_fd_for_plane = dlsym(RTLD_DEFAULT, "gbm_bo_get_fd_for_plane");
    if (allocator->get_fd_for_plane == NULL) {
        kywc_log(KYWC_WARN, "No gbm_bo_get_fd_for_plane support!");
    }

    return true;
}
