#include <devices/mouse.h>
#include <interrupt/idt.h>
#include <interrupt/pic.h>
#include <serial.h>

static inline void outb(UINT16 port, UINT8 val) {
    ASM("outb %0, %1" : : "a"(val), "Nd"(port));
}

static inline UINT8 inb(UINT16 port) {
    UINT8 ret;
    ASM("inb %1, %0" : "=a"(ret) : "Nd"(port));
    return ret;
}

// 等待 PS/2 控制器输入缓冲区清空
static void ps2_wait_input(void) {
    for (int i = 0; i < 100000; i++) {
        if (!(inb(PS2_STATUS_PORT) & PS2_STAT_IBF)) return;
    }
}

// 等待 PS/2 控制器输出缓冲区有数据
static void ps2_wait_output(void) {
    for (int i = 0; i < 100000; i++) {
        if (inb(PS2_STATUS_PORT) & PS2_STAT_OBF) return;
    }
}

// 向 PS/2 控制器发送命令
static void ps2_write_cmd(UINT8 cmd) {
    ps2_wait_input();
    outb(PS2_CMD_PORT, cmd);
}

// 读取 PS/2 控制器返回的数据
static UINT8 ps2_read_data(void) {
    ps2_wait_output();
    return inb(PS2_DATA_PORT);
}

// 向辅助端口（鼠标）发送数据
static void ps2_write_mouse(UINT8 data) {
    ps2_write_cmd(0xD4);  // Write to Auxiliary Device
    ps2_wait_input();
    outb(PS2_DATA_PORT, data);
}

// 鼠标 3 字节数据包状态
static int    g_pkt_phase = 0;     // 0=byte0, 1=byte1, 2=byte2
static UINT8  g_pkt_byte0 = 0;
static UINT8  g_pkt_byte1 = 0;

// 鼠标光标位置（相对于屏幕中心）
static volatile SSINT32 g_mouse_x = 0;
static volatile SSINT32 g_mouse_y = 0;
static volatile bool    g_mouse_btn_l = false;
static volatile bool    g_mouse_btn_r = false;
static volatile bool    g_mouse_btn_m = false;

// IRQ12 处理函数
static void mouse_irq_handler(trap_frame* frame) {
    (void)frame;
    pic_send_eoi(PS2_MOUSE_IRQ);

    // 读取状态寄存器，检查输出缓冲区是否有数据
    UINT8 status = inb(PS2_STATUS_PORT);
    if (!(status & PS2_STAT_OBF)) return;

    UINT8 data = inb(PS2_DATA_PORT);

    // 解析 3 字节数据包
    switch (g_pkt_phase) {
        case 0:  // byte 0: 按钮 + 标志
            if (!(data & PS2_MOUSE_ALWAYS1)) return;  // 无效字节，重新同步
            g_pkt_byte0 = data;
            g_pkt_phase = 1;
            break;
        case 1:  // byte 1: X delta
            g_pkt_byte1 = data;
            g_pkt_phase = 2;
            break;
        case 2:  // byte 2: Y delta
            g_pkt_byte0 = g_pkt_byte0;  // 确保编译器不优化
            g_pkt_byte1 = g_pkt_byte1;

            // 解析按钮状态
            g_mouse_btn_l = (g_pkt_byte0 & PS2_MOUSE_LEFT)   != 0;
            g_mouse_btn_r = (g_pkt_byte0 & PS2_MOUSE_RIGHT)  != 0;
            g_mouse_btn_m = (g_pkt_byte0 & PS2_MOUSE_MIDDLE) != 0;

            // 解析 X delta
            SSINT32 dx = (SSINT32)(SSINT8)g_pkt_byte1;
            if (g_pkt_byte0 & PS2_MOUSE_X_SIGN) dx |= 0xFFFFFF00;

            // 解析 Y delta
            SSINT32 dy = (SSINT32)(SSINT8)data;
            if (g_pkt_byte0 & PS2_MOUSE_Y_SIGN) dy |= 0xFFFFFF00;

            // 更新光标位置（Y 轴在屏幕上向下为正）
            g_mouse_x += dx;
            g_mouse_y -= dy;

            g_pkt_phase = 0;
            break;
    }
}

// 丢弃输出缓冲区中的残留数据
static void ps2_flush_output(void) {
    for (int i = 0; i < 100; i++) {
        inb(PS2_DATA_PORT);
    }
}

void mouse_init(void) {
    // 1. 控制器自检（必须在命令字节写入之前，因为自检会重置命令字节）
    ps2_write_cmd(PS2_CMD_SELF_TEST);
    UINT8 self_test = ps2_read_data();
    if (self_test != 0x55) {
        serial_write("MOUSE: self-test FAILED\n");
    }

    // 2. 丢弃自检可能产生的残留数据
    ps2_flush_output();

    // 3. 启用辅助端口（鼠标）
    ps2_write_cmd(PS2_CMD_ENABLE_A2);

    // 4. 读取当前命令字节
    ps2_write_cmd(PS2_CMD_READ_MB);
    UINT8 cmd = ps2_read_data();

    // 5. 启用 IRQ12 + IRQ1
    cmd |= PS2_CMD_IRQ12 | PS2_CMD_IRQ1;

    // 6. 写入修改后的命令字节（在自检之后，不会被重置）
    ps2_write_cmd(PS2_CMD_WRITE_MB);
    ps2_wait_input();
    outb(PS2_DATA_PORT, cmd);

    // 7. 验证命令字节已写入
    ps2_write_cmd(PS2_CMD_READ_MB);
    UINT8 verify = ps2_read_data();
    if (verify != cmd) {
        serial_write("MOUSE: cmd verify mismatch\n");
    }

    // 重置鼠标
    ps2_write_mouse(PS2_MOUSE_RESET);
    UINT8 ack = ps2_read_data();
    if (ack == 0xFA) {
        ps2_read_data(); // 丢弃第二字节
    }

    // 启用鼠标数据上报
    ps2_write_mouse(PS2_MOUSE_ENABLE);
    ps2_read_data();

    // 注册 IRQ12 处理函数并取消屏蔽
    idt_set_handler(PS2_MOUSE_VECTOR, mouse_irq_handler);
    pic_unmask_irq(2);   // 取消主片级联屏蔽
    pic_unmask_irq(PS2_MOUSE_IRQ);
}

SSINT32 mouse_get_x(void) { return g_mouse_x; }
SSINT32 mouse_get_y(void) { return g_mouse_y; }
bool    mouse_get_btn_left(void)   { return g_mouse_btn_l; }
bool    mouse_get_btn_right(void)  { return g_mouse_btn_r; }
bool    mouse_get_btn_middle(void) { return g_mouse_btn_m; }