/*
 * DWN - Desktop Window Manager
 * retoor <retoor@molodetz.nl>
 * Screenshot capture implementation - Production Hardened
 */

#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 <setjmp.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+/";

/* Thread state tracking to prevent double-join */
typedef enum {
    THREAD_STATE_RUNNING,
    THREAD_STATE_COMPLETED,
    THREAD_STATE_JOINED
} ThreadState;

typedef struct ScreenshotRequestInternal {
    ScreenshotRequest base;
    ThreadState thread_state;
    struct ScreenshotRequestInternal *next;
} ScreenshotRequestInternal;

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);
    
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)screenshot_queue;
    while (req != NULL) {
        ScreenshotRequestInternal *next = req->next;
        
        /* Only join if not already joined */
        if (req->thread_state != THREAD_STATE_JOINED) {
            pthread_mutex_unlock(&screenshot_mutex);
            pthread_join(req->base.thread, NULL);
            pthread_mutex_lock(&screenshot_mutex);
        }
        
        /* Free resources */
        screenshot_result_free(&req->base.result);
        dwn_free(req);
        
        req = next;
    }
    screenshot_queue = NULL;
    
    pthread_mutex_unlock(&screenshot_mutex);
    
    initialized = false;
}

static size_t base64_encoded_size(size_t input_size)
{
    /* Check for overflow */
    if (input_size > ((size_t)-1 - 2) / 4 * 3) {
        return 0;
    }
    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) {
        if (out_len) *out_len = 0;
        return NULL;
    }
    
    size_t encoded_size = base64_encoded_size(result->size);
    if (encoded_size == 0) {
        if (out_len) *out_len = 0;
        return NULL;
    }
    
    char *encoded = dwn_malloc(encoded_size + 1);
    if (encoded == NULL) {
        if (out_len) *out_len = 0;
        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 = (size_t)(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;

typedef struct {
    PngBuffer *buf;
    unsigned char *row;
    jmp_buf jmpbuf;
    const char *error_msg;
} PngWriteContext;

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) {
        png_error(png, "Invalid buffer pointer");
        return;
    }
    
    /* Check for addition overflow */
    if (length > (size_t)-1 - buf->size) {
        png_error(png, "Buffer size overflow");
        return;
    }
    
    size_t new_size = buf->size + length;
    if (new_size > buf->capacity) {
        /* Grow capacity with overflow check */
        size_t new_capacity = buf->capacity * 2;
        if (new_capacity < new_size) {
            new_capacity = new_size + 1024;
        }
        if (new_capacity < new_size || new_capacity > (size_t)-1 / 2) {
            png_error(png, "Buffer capacity overflow");
            return;
        }
        
        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 += length;
}

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

static void png_error_handler(png_structp png, png_const_charp msg)
{
    PngWriteContext *ctx = (PngWriteContext *)png_get_error_ptr(png);
    if (ctx != NULL) {
        ctx->error_msg = msg;
        longjmp(ctx->jmpbuf, 1);
    }
}

static ScreenshotStatus ximage_to_png(XImage *img, ScreenshotResult *result)
{
    if (img == NULL || result == NULL) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }
    
    /* Check for overflow in image dimensions */
    if (img->width <= 0 || img->height <= 0) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }
    
    if ((size_t)img->width > (size_t)-1 / (size_t)img->height / 3) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }
    
    /* Initialize result */
    memset(result, 0, sizeof(ScreenshotResult));
    
    /* Calculate initial buffer capacity with overflow check */
    size_t initial_capacity = (size_t)(img->width * img->height * 3 / 2 + 1024);
    if (initial_capacity < 1024) {
        return SCREENSHOT_ERROR_INVALID_ARG;
    }
    
    PngBuffer buf = {0};
    buf.capacity = initial_capacity;
    buf.data = dwn_malloc(buf.capacity);
    if (buf.data == NULL) {
        return SCREENSHOT_ERROR_NO_MEMORY;
    }
    
    /* Allocate row buffer */
    size_t row_size = (size_t)img->width * 3;
    if (row_size < (size_t)img->width) {  /* Overflow check */
        dwn_free(buf.data);
        return SCREENSHOT_ERROR_NO_MEMORY;
    }
    
    unsigned char *row = dwn_malloc(row_size);
    if (row == NULL) {
        dwn_free(buf.data);
        return SCREENSHOT_ERROR_NO_MEMORY;
    }
    
    /* Set up context for error handling */
    PngWriteContext ctx;
    ctx.buf = &buf;
    ctx.row = row;
    ctx.error_msg = NULL;
    
    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, &ctx, png_error_handler, NULL);
    if (png == NULL) {
        dwn_free(buf.data);
        dwn_free(row);
        return SCREENSHOT_ERROR_PNG;
    }
    
    png_infop info = png_create_info_struct(png);
    if (info == NULL) {
        png_destroy_write_struct(&png, NULL);
        dwn_free(buf.data);
        dwn_free(row);
        return SCREENSHOT_ERROR_PNG;
    }
    
    /* Set up error handling */
    if (setjmp(ctx.jmpbuf)) {
        /* Error occurred - cleanup all resources */
        png_destroy_write_struct(&png, &info);
        dwn_free(buf.data);
        dwn_free(row);
        LOG_ERROR("PNG encoding error: %s", ctx.error_msg ? ctx.error_msg : "unknown");
        return SCREENSHOT_ERROR_PNG;
    }
    
    png_set_write_fn(png, &buf, png_write_callback, png_flush_callback);
    
    png_set_IHDR(png, info, (png_uint_32)img->width, (png_uint_32)img->height, 8,
                 PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
    
    png_write_info(png, info);
    
    /* Calculate bit shifts for color extraction */
    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;
    }
    
    /* Write image data */
    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);
    png_destroy_write_struct(&png, &info);
    
    dwn_free(row);
    
    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;
    }
    
    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;
    }
    
    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;
    
    /* Clamp to screen bounds */
    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;
    }
    
    /* Clamp to screen bounds */
    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)
{
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)arg;
    
    if (atomic_load(&screenshot_shutting_down)) {
        req->base.status = SCREENSHOT_ERROR_X11;
        req->base.state = SCREENSHOT_STATE_ERROR;
        return NULL;
    }
    
    switch (req->base.mode) {
        case SCREENSHOT_MODE_FULLSCREEN:
            req->base.status = screenshot_capture_fullscreen(&req->base.result);
            break;
        case SCREENSHOT_MODE_WINDOW:
            req->base.status = screenshot_capture_window(req->base.window, &req->base.result);
            break;
        case SCREENSHOT_MODE_ACTIVE:
            req->base.status = screenshot_capture_active(&req->base.result);
            break;
        case SCREENSHOT_MODE_AREA:
            req->base.status = screenshot_capture_area(req->base.x, req->base.y,
                                                        req->base.width, req->base.height,
                                                        &req->base.result);
            break;
        default:
            req->base.status = SCREENSHOT_ERROR_INVALID_ARG;
            break;
    }
    
    req->base.state = (req->base.status == SCREENSHOT_OK)
                      ? SCREENSHOT_STATE_COMPLETED
                      : SCREENSHOT_STATE_ERROR;
    
    return NULL;
}

static ScreenshotRequest *screenshot_create_request(void)
{
    if (!initialized || atomic_load(&screenshot_shutting_down)) {
        return NULL;
    }
    
    ScreenshotRequestInternal *req = dwn_calloc(1, sizeof(ScreenshotRequestInternal));
    if (req == NULL) {
        return NULL;
    }
    
    req->base.state = SCREENSHOT_STATE_PENDING;
    req->thread_state = THREAD_STATE_RUNNING;
    
    return (ScreenshotRequest *)req;
}

static void screenshot_queue_add(ScreenshotRequestInternal *req)
{
    pthread_mutex_lock(&screenshot_mutex);
    
    req->next = (ScreenshotRequestInternal *)screenshot_queue;
    screenshot_queue = (ScreenshotRequest *)req;
    
    pthread_mutex_unlock(&screenshot_mutex);
}

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

ScreenshotRequest *screenshot_capture_async(ScreenshotMode mode,
                                            void (*callback)(ScreenshotRequest *))
{
    ScreenshotRequest *base_req = screenshot_create_request();
    if (base_req == NULL) {
        return NULL;
    }
    
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)base_req;
    req->base.mode = mode;
    req->base.callback = callback;
    
    screenshot_queue_add(req);
    
    if (pthread_create(&req->base.thread, NULL, screenshot_worker, req) != 0) {
        screenshot_queue_remove(req);
        dwn_free(req);
        return NULL;
    }
    
    LOG_DEBUG("Screenshot async request started (mode=%d)", mode);
    return (ScreenshotRequest *)req;
}

ScreenshotRequest *screenshot_capture_window_async(Window window,
                                                   void (*callback)(ScreenshotRequest *))
{
    ScreenshotRequest *base_req = screenshot_create_request();
    if (base_req == NULL) {
        return NULL;
    }
    
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)base_req;
    req->base.mode = SCREENSHOT_MODE_WINDOW;
    req->base.window = window;
    req->base.callback = callback;
    
    screenshot_queue_add(req);
    
    if (pthread_create(&req->base.thread, NULL, screenshot_worker, req) != 0) {
        screenshot_queue_remove(req);
        dwn_free(req);
        return NULL;
    }
    
    LOG_DEBUG("Screenshot async window request started");
    return (ScreenshotRequest *)req;
}

ScreenshotRequest *screenshot_capture_area_async(int x, int y, int w, int h,
                                                 void (*callback)(ScreenshotRequest *))
{
    ScreenshotRequest *base_req = screenshot_create_request();
    if (base_req == NULL) {
        return NULL;
    }
    
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)base_req;
    req->base.mode = SCREENSHOT_MODE_AREA;
    req->base.x = x;
    req->base.y = y;
    req->base.width = w;
    req->base.height = h;
    req->base.callback = callback;
    
    screenshot_queue_add(req);
    
    if (pthread_create(&req->base.thread, NULL, screenshot_worker, req) != 0) {
        screenshot_queue_remove(req);
        dwn_free(req);
        return NULL;
    }
    
    LOG_DEBUG("Screenshot async area request started");
    return (ScreenshotRequest *)req;
}

void screenshot_process_pending(void)
{
    if (atomic_load(&screenshot_shutting_down)) {
        return;
    }
    
    pthread_mutex_lock(&screenshot_mutex);
    
    ScreenshotRequestInternal **pp = (ScreenshotRequestInternal **)&screenshot_queue;
    
    while (*pp != NULL) {
        ScreenshotRequestInternal *req = *pp;
        
        if (req->base.state == SCREENSHOT_STATE_COMPLETED ||
            req->base.state == SCREENSHOT_STATE_ERROR) {
            
            /* Remove from queue first */
            *pp = req->next;
            pthread_mutex_unlock(&screenshot_mutex);
            
            /* Join thread if not already joined */
            if (req->thread_state != THREAD_STATE_JOINED) {
                pthread_join(req->base.thread, NULL);
                req->thread_state = THREAD_STATE_JOINED;
            }
            
            /* Invoke callback */
            if (req->base.callback != NULL) {
                req->base.callback((ScreenshotRequest *)req);
            }
            
            pthread_mutex_lock(&screenshot_mutex);
        } else {
            pp = (ScreenshotRequestInternal **)&(*pp)->next;
        }
    }
    
    pthread_mutex_unlock(&screenshot_mutex);
}

void screenshot_cancel(ScreenshotRequest *base_req)
{
    if (base_req == NULL) {
        return;
    }
    
    ScreenshotRequestInternal *req = (ScreenshotRequestInternal *)base_req;
    
    /* Remove from queue */
    screenshot_queue_remove(req);
    
    /* Note: We don't cancel the thread here as it could be in the middle of X11 operations.
     * Instead, we let it complete and ignore the result. This is safer than pthread_cancel
     * which could leave resources in an inconsistent state. */
    
    /* Wait for completion */
    if (req->thread_state != THREAD_STATE_JOINED) {
        pthread_join(req->base.thread, NULL);
        req->thread_state = THREAD_STATE_JOINED;
    }
    
    /* Free resources */
    screenshot_result_free(&req->base.result);
    dwn_free(req);
}
