#include "wallpaper.h"
#define STBI_NO_STDIO
#include "userland/stb_image.h"
#include "graphics.h"
#include "fat32.h"
#include "memory_manager.h"
#include "wm.h"
#include "io.h"
#include <stddef.h>

// Static buffer for the current wallpaper (max 1920x1080)
#define MAX_WP_WIDTH 1920
#define MAX_WP_HEIGHT 1080
static uint32_t* wp_pixels = NULL;
static int wp_width = 0;
static int wp_height = 0;

// Deferred wallpaper action (set from interrupt context, processed in main loop)
static volatile const char *pending_wallpaper_path = NULL;
static char pending_path_buf[256];

// Simple nearest-neighbor scale from decoded RGBA to ARGB pixel buffer
static void scale_rgba_to_argb(const unsigned char *rgba, int src_w, int src_h,
                               uint32_t *dst, int dst_w, int dst_h) {
    for (int y = 0; y < dst_h; y++) {
        int src_y = y * src_h / dst_h;
        if (src_y >= src_h) src_y = src_h - 1;
        for (int x = 0; x < dst_w; x++) {
            int src_x = x * src_w / dst_w;
            if (src_x >= src_w) src_x = src_w - 1;
            
            size_t idx = ((size_t)src_y * (size_t)src_w + (size_t)src_x) * 4;
            unsigned char r = rgba[idx];
            unsigned char g = rgba[idx + 1];
            unsigned char b = rgba[idx + 2];
            unsigned char a = rgba[idx + 3];
            dst[y * dst_w + x] = ((uint32_t)a << 24) | ((uint32_t)r << 16) | ((uint32_t)g << 8) | b;
        }
    }
}

// Decode JPEG data from memory and set as wallpaper (MUST be called from non-interrupt context)
static int decode_and_set_wallpaper(const unsigned char *jpg_data, unsigned int jpg_size) {
    int img_w, img_h, channels;
    // We request 4 channels (RGBA) for better alignment and consistency with working userland code
    unsigned char *rgba = stbi_load_from_memory(jpg_data, (int)jpg_size, &img_w, &img_h, &channels, 4);
    
    if (!rgba || img_w <= 0 || img_h <= 0) {
        return 0;
    }

    // Scale to screen size
    int screen_w = get_screen_width();
    int screen_h = get_screen_height();
    if (screen_w > MAX_WP_WIDTH) screen_w = MAX_WP_WIDTH;
    if (screen_h > MAX_WP_HEIGHT) screen_h = MAX_WP_HEIGHT;

    if (!wp_pixels) {
        wp_pixels = (uint32_t*)kmalloc(MAX_WP_WIDTH * MAX_WP_HEIGHT * sizeof(uint32_t));
        if (!wp_pixels) {
            stbi_image_free(rgba);
            return 0;
        }
    }

    scale_rgba_to_argb(rgba, img_w, img_h, wp_pixels, screen_w, screen_h);
    wp_width = screen_w;
    wp_height = screen_h;

    stbi_image_free(rgba);

    graphics_set_bg_image(wp_pixels, wp_width, wp_height);
    return 1;
}

// Simple serial output for debugging (COM1 = 0x3F8)
static void serial_char(char c) {
    while (!(inb(0x3F8 + 5) & 0x20));  // Wait for transmit ready
    outb(0x3F8, c);
}

static void serial_str(const char *s) {
    while (*s) serial_char(*s++);
}

static void serial_num(int n) {
    if (n < 0) { serial_char('-'); n = -n; }
    if (n >= 10) serial_num(n / 10);
    serial_char('0' + (n % 10));
}

// Request wallpaper change by file path (safe to call from interrupt context)
void wallpaper_request_set_from_file(const char *path) {
    // Copy path to buffer
    int i = 0;
    while (path[i] && i < 255) {
        pending_path_buf[i] = path[i];
        i++;
    }
    pending_path_buf[i] = 0;
    pending_wallpaper_path = pending_path_buf;
}

// Process deferred wallpaper actions (called from main loop, NOT interrupt context)
void wallpaper_process_pending(void) {
    if (pending_wallpaper_path) {
        const char *path = (const char *)pending_wallpaper_path;
        pending_wallpaper_path = NULL;

        serial_str("[WP] Processing wallpaper: ");
        serial_str(path);
        serial_str("\n");

        // Read file from filesystem
        FAT32_FileHandle *fh = fat32_open(path, "r");
        if (fh) {
            uint32_t file_size = fh->size;
            if (file_size > 0 && file_size <= 4 * 1024 * 1024) {
                // Add padding to avoid stb_image reading past the end and potential corruption
                size_t padded_size = file_size + 128;
                unsigned char *buf = (unsigned char*)kmalloc(padded_size);
                if (buf) {
                    mem_memset(buf, 0, padded_size);
                    int total_read = 0;
                    while (total_read < (int)file_size) {
                        int chunk = fat32_read(fh, buf + total_read, (int)file_size - total_read);
                        if (chunk <= 0) break;
                        total_read += chunk;
                    }
                    fat32_close(fh);

                    if (total_read > 0) {
                        decode_and_set_wallpaper(buf, (unsigned int)total_read);
                        wm_refresh();
                    }
                    kfree(buf);
                } else {
                    fat32_close(fh);
                }
            } else {
                fat32_close(fh);
            }
        }
    }
}

uint32_t* wallpaper_get_pixels(void) { return wp_pixels; }
int wallpaper_get_width(void) { return wp_width; }
int wallpaper_get_height(void) { return wp_height; }

void wallpaper_init(void) {
    // We expect Limine modules to have been copied to /Library/images/Wallpapers/ by main.c
    // Set a default wallpaper if one exists
    if (fat32_exists("/Library/images/Wallpapers/mountain.jpg")) {
        wallpaper_request_set_from_file("/Library/images/Wallpapers/mountain.jpg");
    } else if (fat32_exists("/Library/images/Wallpapers/moon.jpg")) {
        wallpaper_request_set_from_file("/Library/images/Wallpapers/moon.jpg");
    }
}