#include <graphics/layer.h>
#include <graphics/draw.h>
#include <graphics/context.h>
#include <memory/heap.h>
#include <memory/pmm.h>
#include <scheduler.h>
#include <serial.h>
#include <idt.h>
#include <pic.h>
#include <common.h>

// --- Layer list (sorted by z, lowest first) ---

static layer_t g_layers[LAYER_MAX];
static UINT32  g_layer_count = 0;
static layer_t* g_layer_list = NULL;

// Compositor back buffer
static EFI_GRAPHICS_OUTPUT_BLT_PIXEL* g_back_buffer = NULL;

// Focus tracking
static layer_t* g_focused = NULL;

// Shift+F10 state (set by IRQ handler, consumed by compositor)
static volatile bool     g_shift_held = false;
static volatile bool     g_switch_pending = false;
static volatile layer_t* g_switch_target = NULL;

// PS/2 scan code set 1
#define PS2_F10        0x44
#define PS2_LSHIFT     0x2A
#define PS2_RSHIFT     0x36
#define PS2_BREAK_BIT  0x80

// Forward declare
static void layer_insert_sorted(layer_t* layer);
static void layer_remove(layer_t* layer);
static layer_t* find_next_window(layer_t* from);

// --- PS/2 keyboard IRQ handler ---

static void ps2_irq_handler(trap_frame* frame) {
    (void)frame;
    pic_send_eoi(1);

    UINT8 sc;
    ASM("inb %1, %0" : "=a"(sc) : "Nd"((UINT16)0x60));

    bool is_break = (sc & PS2_BREAK_BIT) != 0;
    UINT8 code = sc & 0x7F;

    if (code == PS2_LSHIFT || code == PS2_RSHIFT) {
        g_shift_held = !is_break;
        return;
    }

    if (!is_break && g_shift_held && code == PS2_F10) {
        layer_t* next = find_next_window((layer_t*)g_focused);
        if (next) {
            g_switch_target = next;
            g_switch_pending = true;
        }
    }
}

// --- Layer management ---

layer_t* layer_create(const char* name, layer_type_t type, UINT32 w, UINT32 h) {
    if (g_layer_count >= LAYER_MAX) {
        serial_write("LAYER: limit reached\n");
        return NULL;
    }

    UINT32 id = g_layer_count++;
    layer_t* layer = &g_layers[id];

    layer->id = id;
    layer->type = type;
    layer->x = 0;
    layer->y = 0;
    layer->w = w;
    layer->h = h;
    layer->z = 0;
    layer->visible = true;
    layer->next = NULL;

    const char* s = name;
    char* d = layer->name;
    for (SSINT32 i = 0; i < LAYER_NAME_LEN - 1 && *s; i++) {
        *d++ = *s++;
    }
    *d = '\0';

    UINTN buf_size = (UINTN)w * h * sizeof(EFI_GRAPHICS_OUTPUT_BLT_PIXEL);
    layer->buffer = (EFI_GRAPHICS_OUTPUT_BLT_PIXEL*)kmalloc(buf_size);
    if (!layer->buffer) {
        serial_write("LAYER: buffer alloc failed for ");
        serial_write(name);
        serial_write("\n");
        g_layer_count--;
        return NULL;
    }

    EFI_GRAPHICS_OUTPUT_BLT_PIXEL* p = layer->buffer;
    for (UINTN i = 0; i < (UINTN)w * h; i++) {
        p->Blue = 0; p->Green = 0; p->Red = 0; p->Reserved = 0;
        p++;
    }

    layer_insert_sorted(layer);

    serial_write("LAYER: created '");
    serial_write(name);
    serial_write("' id=");
    serial_write_hex(id);
    serial_write("\n");

    return layer;
}

void layer_destroy(layer_t* layer) {
    if (!layer) return;

    layer_remove(layer);

    if (layer->buffer) {
        kfree(layer->buffer);
        layer->buffer = NULL;
    }

    if (g_focused == layer) g_focused = NULL;
    if (g_switch_target == layer) {
        g_switch_target = NULL;
        g_switch_pending = false;
    }

    serial_write("LAYER: destroyed '");
    serial_write(layer->name);
    serial_write("'\n");
}

layer_t* layer_get_by_id(UINT32 id) {
    if (id >= LAYER_MAX) return NULL;
    layer_t* l = &g_layers[id];
    return l->buffer ? l : NULL;
}

layer_t* layer_get_focused(void) {
    return g_focused;
}

void layer_set_z(layer_t* layer, SSINT32 z) {
    if (!layer) return;
    layer_remove(layer);
    layer->z = z;
    layer_insert_sorted(layer);
}

void layer_set_pos(layer_t* layer, SSINT32 x, SSINT32 y) {
    if (!layer) return;
    layer->x = x;
    layer->y = y;
}

void layer_set_visible(layer_t* layer, bool visible) {
    if (!layer) return;
    layer->visible = visible;
}

// --- Sorted insert/remove ---

static void layer_insert_sorted(layer_t* layer) {
    layer->next = NULL;

    if (!g_layer_list || layer->z < g_layer_list->z) {
        layer->next = g_layer_list;
        g_layer_list = layer;
        return;
    }

    layer_t* cur = g_layer_list;
    while (cur->next && cur->next->z <= layer->z) {
        cur = cur->next;
    }
    layer->next = cur->next;
    cur->next = layer;
}

static void layer_remove(layer_t* layer) {
    if (!g_layer_list) return;

    if (g_layer_list == layer) {
        g_layer_list = layer->next;
        layer->next = NULL;
        return;
    }

    layer_t* cur = g_layer_list;
    while (cur->next && cur->next != layer) {
        cur = cur->next;
    }
    if (cur->next == layer) {
        cur->next = layer->next;
        layer->next = NULL;
    }
}

static layer_t* find_next_window(layer_t* from) {
    layer_t* start = from ? from->next : g_layer_list;
    if (!start) start = g_layer_list;

    layer_t* cur = start;
    do {
        if (cur->type == LAYER_TYPE_WINDOW && cur->visible) {
            return cur;
        }
        cur = cur->next;
        if (!cur) cur = g_layer_list;
    } while (cur != start);

    return NULL;
}

// --- Initialization ---

void layer_init(void) {
    UINT32 hr = g_gfx.hr;
    UINT32 vr = g_gfx.vr;

    UINTN buf_size = (UINTN)hr * vr * sizeof(EFI_GRAPHICS_OUTPUT_BLT_PIXEL);
    g_back_buffer = (EFI_GRAPHICS_OUTPUT_BLT_PIXEL*)kmalloc(buf_size);
    if (!g_back_buffer) {
        serial_write("LAYER: back buffer alloc FAILED\n");
        return;
    }

    EFI_GRAPHICS_OUTPUT_BLT_PIXEL* p = g_back_buffer;
    for (UINTN i = 0; i < (UINTN)hr * vr; i++) {
        p->Blue = 0; p->Green = 0; p->Red = 0; p->Reserved = 0;
        p++;
    }

    // Register keyboard IRQ and unmask
    idt_set_handler(PIC_IRQ_BASE + 1, ps2_irq_handler);
    pic_unmask_irq(1);

    serial_write("LAYER: compositor init OK (back buffer = ");
    serial_write_hex(buf_size);
    serial_write(" bytes)\n");
}

// --- Compositor task ---

void layer_compositor_task(void) {
    serial_write("LAYER: compositor task running\n");

    UINT32 hr = g_gfx.hr;
    UINT32 vr = g_gfx.vr;

    while (1) {
        // Process deferred Shift+F10 switch (safe: not inside IRQ)
        if (g_switch_pending) {
            g_switch_pending = false;
            layer_t* target = (layer_t*)g_switch_target;
            if (target && target->visible && target->type == LAYER_TYPE_WINDOW) {
                g_focused = target;
                layer_set_z(target, 99);
                serial_write("LAYER: Shift+F10 -> '");
                serial_write(target->name);
                serial_write("'\n");
            }
        }

        // Clear back buffer
        EFI_GRAPHICS_OUTPUT_BLT_PIXEL black = {0, 0, 0, 0};
        draw_set_target(g_back_buffer, hr, vr);
        draw_rect(0, 0, hr, vr, black);
        draw_set_default_target();

        // Composite layers from lowest z to highest
        layer_t* cur = g_layer_list;
        while (cur) {
            if (cur->visible && cur->buffer) {
                SSINT32 sx = 0, sy = 0;
                SSINT32 dx = cur->x, dy = cur->y;
                UINT32 sw = cur->w, sh = cur->h;

                if (dx < 0) { sx = -dx; sw -= sx; dx = 0; }
                if (dy < 0) { sy = -dy; sh -= sy; dy = 0; }
                if (dx + (SSINT32)sw > (SSINT32)hr) sw = hr - dx;
                if (dy + (SSINT32)sh > (SSINT32)vr) sh = vr - dy;
                if (sw == 0 || sh == 0) { cur = cur->next; continue; }

                for (UINT32 row = 0; row < sh; row++) {
                    EFI_GRAPHICS_OUTPUT_BLT_PIXEL* src = cur->buffer + ((UINTN)cur->w * (sy + row)) + sx;
                    EFI_GRAPHICS_OUTPUT_BLT_PIXEL* dst = g_back_buffer + ((UINTN)hr * (dy + row)) + dx;
                    for (UINT32 col = 0; col < sw; col++) {
                        dst->Blue  = src->Blue;
                        dst->Green = src->Green;
                        dst->Red   = src->Red;
                        dst->Reserved = src->Reserved;
                        src++;
                        dst++;
                    }
                }
            }
            cur = cur->next;
        }

        // Blit to screen
        g_gfx.GOP->Blt(
            g_gfx.GOP,
            g_back_buffer,
            EfiBltBufferToVideo,
            0, 0, 0, 0,
            hr, vr, 0
        );

        yield();
    }
}