/*
 * DWN - Desktop Window Manager
 * retoor <retoor@molodetz.nl>
 * Screenshot capture implementation using X11 and libpng
 */

#include "screenshot.h"
#include "dwn.h"
#include "client.h"
#include "workspace.h"
#include "util.h"

#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <png.h>

static bool initialized = false;
static pthread_mutex_t screenshot_mutex = PTHREAD_MUTEX_INITIALIZER;
static ScreenshotRequest *screenshot_queue = NULL;
static atomic_int screenshot_shutting_down = 0;

static const char base64_chars[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

bool screenshot_init(void)
{
    atomic_store(&screenshot_shutting_down, 0);
    initialized = true;
    LOG_INFO("Screenshot module initialized");
    return true;
}

void screenshot_cleanup(void)
{
    atomic_store(&screenshot_shutting_down, 1);

    pthread_mutex_lock(&screenshot_mutex);
    while (screenshot_queue != NULL) {
        ScreenshotRequest *req = screenshot_queue;
        screenshot_queue = req->next;
        pthread_mutex_unlock(&screenshot_mutex);

        pthread_join(req->thread, NULL);
        screenshot_result_free(&req->result);
        dwn_free(req);

        pthread_mutex_lock(&screenshot_mutex);
    }
    pthread_mutex_unlock(&screenshot_mutex);

    initialized = false;
}

static size_t base64_encoded_size(size_t input_size)
{
    return 4 * ((input_size + 2) / 3);
}

char *screenshot_to_base64(const ScreenshotResult *result, size_t *out_len)
{
    if (result == NULL || result->data == NULL || result->size == 0) {
        return NULL;
    }

    size_t encoded_size = base64_encoded_size(result->size);
    char *encoded = dwn_malloc(encoded_size + 1);
    if (encoded == NULL) {
        return NULL;
    }

    const unsigned char *input = result->data;
    size_t input_len = result->size;
    char *output = encoded;

    for (size_t i = 0; i < input_len; i += 3) {
        unsigned int n = ((unsigned int)input[i]) << 16;
        if (i + 1 < input_len) {
            n |= ((unsigned int)input[i + 1]) << 8;
        }
        if (i + 2 < input_len) {
            n |= (unsigned int)input[i + 2];
        }

        *output++ = base64_chars[(n >> 18) & 0x3F];
        *output++ = base64_chars[(n >> 12) & 0x3F];
        *output++ = (i + 1 < input_len) ? base64_chars[(n >> 6) & 0x3F] : '=';
        *output++ = (i + 2 < input_len) ? base64_chars[n & 0x3F] : '=';
    }

    *output = '\0';
    if (out_len != NULL) {
        *out_len = output - encoded;
    }

    return encoded;
}

void screenshot_result_free(ScreenshotResult *result)
{
    if (result == NULL) {
        return;
    }
    if (result->data != NULL) {
        dwn_free(result->data);
        result->data = NULL;
    }
    result->size = 0;
    result->width = 0;
    result->height = 0;
}

const char *screenshot_status_string(ScreenshotStatus status)
{
    switch (status) {
        case SCREENSHOT_OK:
            return "OK";
        case SCREENSHOT_ERROR_INVALID_ARG:
            return "Invalid argument";
        case SCREENSHOT_ERROR_X11:
            return "X11 error";
        case SCREENSHOT_ERROR_PNG:
            return "PNG encoding error";
        case SCREENSHOT_ERROR_NO_MEMORY:
            return "Out of memory";
        case SCREENSHOT_ERROR_WINDOW_NOT_FOUND:
            return "Window not found";
        default:
            return "Unknown error";
    }
}

typedef struct {
    unsigned char *data;
    size_t size;
    size_t capacity;
} PngBuffer;

static void png_write_callback(png_structp png, png_bytep data, png_size_t length)
{
    PngBuffer *buf = (PngBuffer *)png_get_io_ptr(png);
    if (buf == NULL) {
        return;
    }

    size_t new_size = buf->size + length;
    if (new_size > buf->capacity) {
        size_t new_capacity = buf->capacity * 2;
        if (new_capacity < new_size) {
            new_capacity = new_size + 1024;
        }
        unsigned char *new_data = dwn_realloc(buf->data, new_capacity);
        if (new_data == NULL) {
            png_error(png, "Memory allocation failed");
            return;
        }
        buf->data = new_data;
        buf->capacity = new_capacity;
    }

    memcpy(buf->data + buf->size, data, length);
    buf->size = new_size;
}

static void png_flush_callback(png_structp png)
{
    (void)png;
}

static ScreenshotStatus ximage_to_png(XImage *img, ScreenshotResult *result)
{
    if (img == NULL || result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (png == NULL) {
        return SCREENSHOT_ERROR_PNG;
    }

    png_infop info = png_create_info_struct(png);
    if (info == NULL) {
        png_destroy_write_struct(&png, NULL);
        return SCREENSHOT_ERROR_PNG;
    }

    if (setjmp(png_jmpbuf(png))) {
        png_destroy_write_struct(&png, &info);
        return SCREENSHOT_ERROR_PNG;
    }

    PngBuffer buf = {0};
    buf.capacity = (size_t)(img->width * img->height * 3 / 2 + 1024);
    buf.data = dwn_malloc(buf.capacity);
    if (buf.data == NULL) {
        png_destroy_write_struct(&png, &info);
        return SCREENSHOT_ERROR_NO_MEMORY;
    }

    png_set_write_fn(png, &buf, png_write_callback, png_flush_callback);

    png_set_IHDR(png, info, img->width, img->height, 8,
                 PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

    png_write_info(png, info);

    unsigned char *row = dwn_malloc(img->width * 3);
    if (row == NULL) {
        dwn_free(buf.data);
        png_destroy_write_struct(&png, &info);
        return SCREENSHOT_ERROR_NO_MEMORY;
    }

    unsigned long red_mask = img->red_mask;
    unsigned long green_mask = img->green_mask;
    unsigned long blue_mask = img->blue_mask;

    int red_shift = 0;
    int green_shift = 0;
    int blue_shift = 0;

    while ((red_mask & 1) == 0 && red_shift < 32) {
        red_shift++;
        red_mask >>= 1;
    }
    while ((green_mask & 1) == 0 && green_shift < 32) {
        green_shift++;
        green_mask >>= 1;
    }
    while ((blue_mask & 1) == 0 && blue_shift < 32) {
        blue_shift++;
        blue_mask >>= 1;
    }

    for (int y = 0; y < img->height; y++) {
        for (int x = 0; x < img->width; x++) {
            unsigned long pixel = XGetPixel(img, x, y);
            row[x * 3 + 0] = (pixel >> red_shift) & 0xFF;
            row[x * 3 + 1] = (pixel >> green_shift) & 0xFF;
            row[x * 3 + 2] = (pixel >> blue_shift) & 0xFF;
        }
        png_write_row(png, row);
    }

    png_write_end(png, info);

    dwn_free(row);
    png_destroy_write_struct(&png, &info);

    result->data = buf.data;
    result->size = buf.size;
    result->width = img->width;
    result->height = img->height;

    return SCREENSHOT_OK;
}

ScreenshotStatus screenshot_capture_fullscreen(ScreenshotResult *result)
{
    if (result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    if (dwn == NULL || dwn->display == NULL) {
        return SCREENSHOT_ERROR_X11;
    }

    memset(result, 0, sizeof(ScreenshotResult));

    XImage *img = XGetImage(dwn->display, dwn->root,
                            0, 0, dwn->screen_width, dwn->screen_height,
                            AllPlanes, ZPixmap);
    if (img == NULL) {
        LOG_ERROR("XGetImage failed for fullscreen capture");
        return SCREENSHOT_ERROR_X11;
    }

    ScreenshotStatus status = ximage_to_png(img, result);
    XDestroyImage(img);

    return status;
}

ScreenshotStatus screenshot_capture_window(Window window, ScreenshotResult *result)
{
    if (result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    if (dwn == NULL || dwn->display == NULL) {
        return SCREENSHOT_ERROR_X11;
    }

    if (window == None) {
        return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
    }

    memset(result, 0, sizeof(ScreenshotResult));

    XWindowAttributes wa;
    if (!XGetWindowAttributes(dwn->display, window, &wa)) {
        return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
    }

    int x, y;
    Window child;
    XTranslateCoordinates(dwn->display, window, dwn->root, 0, 0, &x, &y, &child);

    int width = wa.width;
    int height = wa.height;

    if (x < 0) {
        width += x;
        x = 0;
    }
    if (y < 0) {
        height += y;
        y = 0;
    }
    if (x + width > dwn->screen_width) {
        width = dwn->screen_width - x;
    }
    if (y + height > dwn->screen_height) {
        height = dwn->screen_height - y;
    }

    if (width <= 0 || height <= 0) {
        return SCREENSHOT_ERROR_X11;
    }

    XImage *img = XGetImage(dwn->display, dwn->root, x, y, width, height,
                            AllPlanes, ZPixmap);
    if (img == NULL) {
        LOG_ERROR("XGetImage failed for window capture");
        return SCREENSHOT_ERROR_X11;
    }

    ScreenshotStatus status = ximage_to_png(img, result);
    XDestroyImage(img);

    return status;
}

ScreenshotStatus screenshot_capture_active(ScreenshotResult *result)
{
    if (result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    if (dwn == NULL) {
        return SCREENSHOT_ERROR_X11;
    }

    Workspace *ws = workspace_get_current();
    if (ws == NULL || ws->focused == NULL) {
        return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
    }

    Client *c = ws->focused;
    Window target = (c->frame != None) ? c->frame : c->window;

    return screenshot_capture_window(target, result);
}

ScreenshotStatus screenshot_capture_area(int x, int y, int width, int height, ScreenshotResult *result)
{
    if (result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    if (width <= 0 || height <= 0) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    if (dwn == NULL || dwn->display == NULL) {
        return SCREENSHOT_ERROR_X11;
    }

    memset(result, 0, sizeof(ScreenshotResult));

    if (x < 0) {
        width += x;
        x = 0;
    }
    if (y < 0) {
        height += y;
        y = 0;
    }
    if (x + width > dwn->screen_width) {
        width = dwn->screen_width - x;
    }
    if (y + height > dwn->screen_height) {
        height = dwn->screen_height - y;
    }

    if (width <= 0 || height <= 0) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }

    XImage *img = XGetImage(dwn->display, dwn->root, x, y, width, height,
                            AllPlanes, ZPixmap);
    if (img == NULL) {
        LOG_ERROR("XGetImage failed for area capture");
        return SCREENSHOT_ERROR_X11;
    }

    ScreenshotStatus status = ximage_to_png(img, result);
    XDestroyImage(img);

    return status;
}

static void *screenshot_worker(void *arg)
{
    ScreenshotRequest *req = (ScreenshotRequest *)arg;

    if (atomic_load(&screenshot_shutting_down)) {
        req->state = SCREENSHOT_STATE_ERROR;
        req->status = SCREENSHOT_ERROR_X11;
        return NULL;
    }

    switch (req->mode) {
        case SCREENSHOT_MODE_FULLSCREEN:
            req->status = screenshot_capture_fullscreen(&req->result);
            break;
        case SCREENSHOT_MODE_WINDOW:
            req->status = screenshot_capture_window(req->window, &req->result);
            break;
        case SCREENSHOT_MODE_ACTIVE:
            req->status = screenshot_capture_active(&req->result);
            break;
        case SCREENSHOT_MODE_AREA:
            req->status = screenshot_capture_area(req->x, req->y,
                                                  req->width, req->height,
                                                  &req->result);
            break;
    }

    req->state = (req->status == SCREENSHOT_OK)
                 ? SCREENSHOT_STATE_COMPLETED
                 : SCREENSHOT_STATE_ERROR;

    return NULL;
}

ScreenshotRequest *screenshot_capture_async(ScreenshotMode mode,
                                            void (*callback)(ScreenshotRequest *))
{
    if (!initialized || atomic_load(&screenshot_shutting_down)) {
        return NULL;
    }

    ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
    if (req == NULL) {
        return NULL;
    }

    req->mode = mode;
    req->state = SCREENSHOT_STATE_PENDING;
    req->callback = callback;

    pthread_mutex_lock(&screenshot_mutex);
    req->next = screenshot_queue;
    screenshot_queue = req;
    pthread_mutex_unlock(&screenshot_mutex);

    if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
        pthread_mutex_lock(&screenshot_mutex);
        ScreenshotRequest **pp = &screenshot_queue;
        while (*pp != NULL) {
            if (*pp == req) {
                *pp = req->next;
                break;
            }
            pp = &(*pp)->next;
        }
        pthread_mutex_unlock(&screenshot_mutex);
        dwn_free(req);
        return NULL;
    }

    LOG_DEBUG("Screenshot async request started (mode=%d)", mode);
    return req;
}

ScreenshotRequest *screenshot_capture_window_async(Window window,
                                                   void (*callback)(ScreenshotRequest *))
{
    if (!initialized || atomic_load(&screenshot_shutting_down)) {
        return NULL;
    }

    ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
    if (req == NULL) {
        return NULL;
    }

    req->mode = SCREENSHOT_MODE_WINDOW;
    req->window = window;
    req->state = SCREENSHOT_STATE_PENDING;
    req->callback = callback;

    pthread_mutex_lock(&screenshot_mutex);
    req->next = screenshot_queue;
    screenshot_queue = req;
    pthread_mutex_unlock(&screenshot_mutex);

    if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
        pthread_mutex_lock(&screenshot_mutex);
        ScreenshotRequest **pp = &screenshot_queue;
        while (*pp != NULL) {
            if (*pp == req) {
                *pp = req->next;
                break;
            }
            pp = &(*pp)->next;
        }
        pthread_mutex_unlock(&screenshot_mutex);
        dwn_free(req);
        return NULL;
    }

    LOG_DEBUG("Screenshot async window request started");
    return req;
}

ScreenshotRequest *screenshot_capture_area_async(int x, int y, int w, int h,
                                                 void (*callback)(ScreenshotRequest *))
{
    if (!initialized || atomic_load(&screenshot_shutting_down)) {
        return NULL;
    }

    ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
    if (req == NULL) {
        return NULL;
    }

    req->mode = SCREENSHOT_MODE_AREA;
    req->x = x;
    req->y = y;
    req->width = w;
    req->height = h;
    req->state = SCREENSHOT_STATE_PENDING;
    req->callback = callback;

    pthread_mutex_lock(&screenshot_mutex);
    req->next = screenshot_queue;
    screenshot_queue = req;
    pthread_mutex_unlock(&screenshot_mutex);

    if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
        pthread_mutex_lock(&screenshot_mutex);
        ScreenshotRequest **pp = &screenshot_queue;
        while (*pp != NULL) {
            if (*pp == req) {
                *pp = req->next;
                break;
            }
            pp = &(*pp)->next;
        }
        pthread_mutex_unlock(&screenshot_mutex);
        dwn_free(req);
        return NULL;
    }

    LOG_DEBUG("Screenshot async area request started");
    return req;
}

void screenshot_process_pending(void)
{
    if (atomic_load(&screenshot_shutting_down)) {
        return;
    }

    pthread_mutex_lock(&screenshot_mutex);
    ScreenshotRequest **pp = &screenshot_queue;

    while (*pp != NULL) {
        ScreenshotRequest *req = *pp;

        if (req->state == SCREENSHOT_STATE_COMPLETED ||
            req->state == SCREENSHOT_STATE_ERROR) {

            pthread_join(req->thread, NULL);

            *pp = req->next;
            pthread_mutex_unlock(&screenshot_mutex);

            if (req->callback != NULL) {
                req->callback(req);
            }

            pthread_mutex_lock(&screenshot_mutex);
        } else {
            pp = &(*pp)->next;
        }
    }

    pthread_mutex_unlock(&screenshot_mutex);
}

void screenshot_cancel(ScreenshotRequest *req)
{
    if (req == NULL) {
        return;
    }

    pthread_mutex_lock(&screenshot_mutex);
    ScreenshotRequest **pp = &screenshot_queue;
    while (*pp != NULL) {
        if (*pp == req) {
            *pp = req->next;
            break;
        }
        pp = &(*pp)->next;
    }
    pthread_mutex_unlock(&screenshot_mutex);

    pthread_cancel(req->thread);
    pthread_join(req->thread, NULL);

    screenshot_result_free(&req->result);
    dwn_free(req);
}