#include "mouse.h"
#include "fb.h"
#include "arch.h"
#include "usbkbd.h"
#include "dbg.h"

static void delay_us(uint64_t us)
{
    uint64_t start = rdtsc();
    uint64_t end = start + us * 4000;
    while(rdtsc() < end);
}

/* ── PS/2 aux (mouse/touchpad) ─────────────── */

#define PS2_DATA    0x60
#define PS2_STATUS  0x64
#define PS2_CMD     0x64

static int  ps2_mouse_byte;
static int  ps2_mouse_idx;
static uint8_t ps2_mouse_buf[4];

static void ps2_wait_input(void)
{
    uint64_t deadline = rdtsc() + 200000000ULL;
    while((inb(PS2_STATUS) & 1) == 0)
        if(rdtsc() > deadline) return;
}

static void ps2_wait_output(void)
{
    uint64_t deadline = rdtsc() + 200000000ULL;
    while(inb(PS2_STATUS) & 2)
        if(rdtsc() > deadline) return;
}

static void ps2_cmd(uint8_t cmd)
{
    ps2_wait_output();
    outb(PS2_CMD, cmd);
}

static void ps2_write(uint8_t reg, uint8_t data)
{
    ps2_cmd(reg);
    ps2_wait_output();
    outb(PS2_DATA, data);
}

static uint8_t ps2_read(void)
{
    ps2_wait_input();
    return inb(PS2_DATA);
}

static void ps2_to_mouse(uint8_t data)
{
    ps2_write(0xD4, data);
}

void ps2_mouse_init(void)
{
    ps2_mouse_byte = 0;
    ps2_mouse_idx = 0;

    /* Enable auxiliary device (mouse) */
    ps2_cmd(0xA8);
    delay_us(100);

    /* Read controller config byte */
    ps2_cmd(0x20);
    uint8_t cfg = ps2_read();
    dbg_printf("[ps2mouse] cfg=0x%02x\n", cfg);

    /* Enable IRQ12 (bit1) */
    ps2_write(0x60, cfg | (1 << 1));
    delay_us(100);

    /* Drain any stale data */
    while(inb(PS2_STATUS) & 1) inb(PS2_DATA);

    /* Reset mouse */
    ps2_to_mouse(0xFF);
    delay_us(1000);

    uint8_t ack = ps2_read();
    if(ack != 0xFA) {
        dbg_printf("[ps2mouse] no ack (0x%02x), no PS/2 mouse\n", ack);
        return;
    }
    uint8_t selftest = ps2_read();
    dbg_printf("[ps2mouse] ack=0x%02x selftest=0x%02x\n", ack, selftest);

    /* Detect IntelliMouse (4-byte with wheel) via device ID query */
    /* Set sample rate 100 */
    ps2_to_mouse(0xF3); ps2_read();
    ps2_to_mouse(0x64); ps2_read();
    /* Set resolution max */
    ps2_to_mouse(0xE8); ps2_read();
    ps2_to_mouse(0x03); ps2_read();
    /* Set sample rate 200 */
    ps2_to_mouse(0xF3); ps2_read();
    ps2_to_mouse(0xC8); ps2_read();
    /* Set sample rate 100 */
    ps2_to_mouse(0xF3); ps2_read();
    ps2_to_mouse(0x64); ps2_read();
    /* Read device ID */
    ps2_to_mouse(0xF2); ps2_read();
    uint8_t devid = ps2_read();
    dbg_printf("[ps2mouse] device ID=0x%02x\n", devid);

    ps2_mouse_byte = (devid == 0x03) ? 4 : 3;
    ps2_mouse_idx = 0;
}

int ps2_mouse_poll(int *buttons, int *dx, int *dy, int *wheel)
{
    if(!ps2_mouse_byte) return 0;

    /* Check aux data available (bit5 of status) */
    if(!(inb(PS2_STATUS) & 0x20)) return 0;

    uint8_t b = inb(PS2_DATA);
    ps2_mouse_buf[ps2_mouse_idx++] = b;

    if(ps2_mouse_idx >= ps2_mouse_byte) {
        ps2_mouse_idx = 0;
        uint8_t hdr = ps2_mouse_buf[0];
        int8_t ddx = (int8_t)ps2_mouse_buf[1];
        int8_t ddy = (int8_t)ps2_mouse_buf[2];
        int wh = (ps2_mouse_byte >= 4) ? (int8_t)ps2_mouse_buf[3] : 0;

        if(buttons) *buttons = (hdr & 1) | ((hdr >> 1) & 2) | ((hdr >> 2) & 4);
        if(dx) *dx = ddx;
        if(dy) *dy = -ddy;
        if(wheel) *wheel = -wh; /* PS/2: +1 = forward, invert for natural scroll */
        return 1;
    }
    return 0;
}

#define CURSOR_W 12
#define CURSOR_H 16

/* 12x16 cursor with per-pixel 2-bit colors
   Each row = 3 bytes (12 pixels × 2 bits)
   00 = transparent, 01 = fill (color 0), 10 = outline (color 1) */
static const uint8_t cursor_data[CURSOR_H][3] = {
    {0x80, 0x00, 0x00},
    {0xA0, 0x00, 0x00},
    {0x98, 0x00, 0x00},
    {0x96, 0x00, 0x00},
    {0x95, 0x80, 0x00},
    {0x95, 0x60, 0x00},
    {0x95, 0x58, 0x00},
    {0x95, 0x56, 0x00},
    {0x95, 0x55, 0x80},
    {0x95, 0x6A, 0xA0},
    {0x96, 0x58, 0x00},
    {0x98, 0x26, 0x00},
    {0xA0, 0x09, 0x80},
    {0x80, 0x02, 0x60},
    {0x00, 0x00, 0xA0},
    {0x00, 0x00, 0x00},
};

static uint32_t cursor_palette[2];

static inline int cursor_pixel_at(int r, int c) {
    if(c < 0 || c >= CURSOR_W || r < 0 || r >= CURSOR_H) return 0;
    int bi = c / 4;
    int sh = (3 - (c % 4)) * 2;
    return (cursor_data[r][bi] >> sh) & 3;
}

static int cursor_x, cursor_y;
static int cursor_visible;
static uint32_t saved_bg[CURSOR_W * CURSOR_H];
static int saved_x, saved_y;

void mouse_cursor_init(void)
{
    cursor_palette[0] = make_color(0xFF, 0xFF, 0xFF);
    cursor_palette[1] = make_color(0x00, 0x00, 0x00);
    cursor_x = fb.width / 2;
    cursor_y = fb.height / 2;
    cursor_visible = 1;
    saved_x = -1;
    saved_y = -1;
    /* Draw cursor at initial position immediately */
    mouse_cursor_update(0, 0);
}

static inline uint32_t *fb_pixel(int x, int y)
{
    uint32_t *ptr = (uint32_t *)(uintptr_t)fb.base;
    int stride = fb.pitch / 4;
    return &ptr[y * stride + x];
}

static void save_background(int x, int y)
{
    int stride = fb.pitch / 4;
    uint32_t *ptr = (uint32_t *)(uintptr_t)fb.base;
    for(int row = 0; row < CURSOR_H; row++) {
        int py = y + row;
        if(py < 0 || py >= (int)fb.height) continue;
        for(int col = 0; col < CURSOR_W; col++) {
            int px = x + col;
            if(px < 0 || px >= (int)fb.width) continue;
            saved_bg[row * CURSOR_W + col] = ptr[py * stride + px];
        }
    }
}

static void restore_background(int x, int y)
{
    int stride = fb.pitch / 4;
    uint32_t *ptr = (uint32_t *)(uintptr_t)fb.base;
    for(int row = 0; row < CURSOR_H; row++) {
        int py = y + row;
        if(py < 0 || py >= (int)fb.height) continue;
        for(int col = 0; col < CURSOR_W; col++) {
            int px = x + col;
            if(px < 0 || px >= (int)fb.width) continue;
            ptr[py * stride + px] = saved_bg[row * CURSOR_W + col];
        }
    }
}

static void draw_cursor(int x, int y)
{
    int stride = fb.pitch / 4;
    uint32_t *ptr = (uint32_t *)(uintptr_t)fb.base;

    for(int row = 0; row < CURSOR_H; row++) {
        int py = y + row;
        if(py < 0 || py >= (int)fb.height) continue;
        for(int col = 0; col < CURSOR_W; col++) {
            int px = x + col;
            if(px < 0 || px >= (int)fb.width) continue;
            int pix = cursor_pixel_at(row, col);
            if(pix == 0) continue;
            ptr[py * stride + px] = cursor_palette[pix - 1];
        }
    }
}

void mouse_cursor_update(int dx, int dy)
{
    if(!cursor_visible) return;

    if(saved_x >= 0 && saved_y >= 0)
        restore_background(saved_x, saved_y);

    cursor_x += dx;
    cursor_y += dy;
    if(cursor_x < 0) cursor_x = 0;
    if(cursor_y < 0) cursor_y = 0;
    if(cursor_x >= (int)fb.width - CURSOR_W) cursor_x = fb.width - CURSOR_W;
    if(cursor_y >= (int)fb.height - CURSOR_H) cursor_y = fb.height - CURSOR_H;

    save_background(cursor_x, cursor_y);
    draw_cursor(cursor_x, cursor_y);
    saved_x = cursor_x;
    saved_y = cursor_y;
}

void mouse_cursor_render(void)
{
    if(!cursor_visible) return;
    if(saved_x >= 0 && saved_y >= 0)
        restore_background(saved_x, saved_y);
    save_background(cursor_x, cursor_y);
    draw_cursor(cursor_x, cursor_y);
    saved_x = cursor_x;
    saved_y = cursor_y;
}

/* ── Click / wheel tracking ─────────────────── */

static int  click_pending;
static int  click_x, click_y, click_btn;
static int  prev_buttons;
static int  wheel_accum;

void mouse_poll(void)
{
    int btns, dx, dy, wh;
    /* PS/2 mouse first (always works on real hardware) */
    if(ps2_mouse_poll(&btns, &dx, &dy, &wh)) {
        mouse_cursor_update(dx, dy);
        wheel_accum += wh;
        if(btns != prev_buttons) {
            int pressed = btns & ~prev_buttons;
            if(pressed) {
                click_pending = 1;
                click_x = cursor_x;
                click_y = cursor_y;
                click_btn = pressed;
            }
            prev_buttons = btns;
        }
        return;
    }
    /* USB mouse fallback (QEMU) */
    if(usb_mouse_get_event(&btns, &dx, &dy, &wh)) {
        mouse_cursor_update(dx, dy);
        wheel_accum += wh;
        if(btns != prev_buttons) {
            int pressed = btns & ~prev_buttons;
            if(pressed) {
                click_pending = 1;
                click_x = cursor_x;
                click_y = cursor_y;
                click_btn = pressed;
            }
            prev_buttons = btns;
        }
    }
}

int mouse_get_click(int *x, int *y, int *btn)
{
    if(!click_pending) return 0;
    if(x) *x = click_x;
    if(y) *y = click_y;
    if(btn) *btn = click_btn;
    click_pending = 0;
    return click_btn; /* return which button (1=left, 2=right, 4=middle) */
}

int mouse_get_wheel(void)
{
    int w = wheel_accum;
    wheel_accum = 0;
    return w;
}
