#include <lunaix/peripheral/ps2kbd.h>
#include <lunaix/clock.h>
#include <lunaix/timer.h>
#include <lunaix/common.h>
#include <lunaix/syslog.h>

#include <hal/cpu.h>
#include <arch/x86/interrupts.h>
#include <klibc/string.h>

#include <stdint.h>

#define PS2_DEV_CMD_MAX_ATTEMPTS 5

LOG_MODULE("PS2KBD");

static struct ps2_cmd_queue cmd_q;
static struct ps2_key_buffer key_buf;
static struct ps2_kbd_state kbd_state;

#define KEY_NUM(x)      (x + 0x30)
#define KEY_NPAD(x)      ON_KEYPAD(KEY_NUM(x))

// 我们使用 Scancode Set 2

// 大部分的扫描码（键码）
static kbd_keycode_t scancode_set2[] = {
    0, KEY_F9, 0, KEY_F5, KEY_F3, KEY_F1, KEY_F2, KEY_F12, 0, KEY_F10, KEY_F8, KEY_F6,
    KEY_F4, KEY_HTAB, '`', 0, 0, KEY_LALT, KEY_LSHIFT, 0, KEY_LCTRL, 'q', KEY_NUM(1), 
    0, 0, 0, 'z', 's', 'a', 'w', KEY_NUM(2), 0, 0, 'c', 'x', 'd', 'e', KEY_NUM(4), KEY_NUM(3), 
    0, 0, KEY_SPACE, 'v', 'f', 't', 'r', KEY_NUM(5),
    0, 0, 'n', 'b', 'h', 'g', 'y', KEY_NUM(6), 0, 0, 0, 'm', 'j', 'u', KEY_NUM(7), KEY_NUM(8),
    0, 0, ',', 'k', 'i', 'o', KEY_NUM(0), KEY_NUM(9), 0, 0, '.', '/', 'l', ';', 'p', '-', 0, 0,
    0, '\'', 0, '[', '=', 0, 0, KEY_CAPSLK, KEY_RSHIFT, KEY_LF, ']', 0, '\\', 0, 0, 0, 0, 0, 0, 0,
    0, KEY_BS, 0, 0, KEY_NPAD(1), 0, KEY_NPAD(4), KEY_NPAD(7), 0, 0, 0, KEY_NPAD(0), ON_KEYPAD('.'),
    KEY_NPAD(2), KEY_NPAD(5), KEY_NPAD(6), KEY_NPAD(8), KEY_ESC, KEY_NUMSLK, KEY_F11, ON_KEYPAD('+'),
    KEY_NPAD(3), ON_KEYPAD('-'), ON_KEYPAD('*'), KEY_NPAD(9), KEY_SCRLLK, 0, 0, 0, 0, KEY_F7
};

// 一些特殊的键码（以 0xe0 位前缀的）
static kbd_keycode_t scancode_set2_ex[] = {
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, KEY_RALT, 0, 0,
    KEY_RCTRL, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ON_KEYPAD('/'), 0, 0, 0, 0, 0, 0, 0, 0, 
    0, 0, 0, 0, 0, 0, 0, ON_KEYPAD(KEY_LF), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    KEY_END, 0, KEY_LEFT, KEY_HOME,
    0, 0, 0, KEY_INSERT, KEY_DELETE, KEY_DOWN, 0, KEY_RIGHT, KEY_UP, 0, 0,
    0, 0, KEY_PG_DOWN, 0, 0, KEY_PG_UP
};

// 用于处理 Shift+<key> 的情况
static kbd_keycode_t scancode_set2_shift[] = {
    0, KEY_F9, 0, KEY_F5, KEY_F3, KEY_F1, KEY_F2, KEY_F12, 0, KEY_F10, KEY_F8, KEY_F6,
    KEY_F4, KEY_HTAB, '~', 0, 0, KEY_LALT, KEY_LSHIFT, 0, KEY_LCTRL, 'Q', '!', 
    0, 0, 0, 'Z', 'S', 'A', 'W', '@', 0, 0, 'C', 'X', 'D', 'E', '$', '#', 
    0, 0, KEY_SPACE, 'V', 'F', 'T', 'R', '%',
    0, 0, 'N', 'B', 'H', 'G', 'Y', '^', 0, 0, 0, 'M', 'J', 'U', '&', '*',
    0, 0, '<', 'K', 'I', 'O', ')', '(', 0, 0, '>', '?', 'L', ':', 'P', '_', 0, 0,
    0, '"', 0, '{', '+', 0, 0, KEY_CAPSLK, KEY_RSHIFT, KEY_LF, '}', 0, '|', 0, 0, 0, 0, 0, 0, 0,
    0, KEY_BS, 0, 0, KEY_NPAD(1), 0, KEY_NPAD(4), KEY_NPAD(7), 0, 0, 0, KEY_NPAD(0), ON_KEYPAD('.'),
    KEY_NPAD(2), KEY_NPAD(5), KEY_NPAD(6), KEY_NPAD(8), KEY_ESC, KEY_NUMSLK, KEY_F11, ON_KEYPAD('+'),
    KEY_NPAD(3), ON_KEYPAD('-'), ON_KEYPAD('*'), KEY_SCRLLK, 0, 0, 0, 0, KEY_F7
};


#define KBD_STATE_KWAIT         0x00
#define KBD_STATE_KSPECIAL      0x01
#define KBD_STATE_KRELEASED     0x02
#define KBD_STATE_CMDPROCS      0x40

void intr_ps2_kbd_handler(const isr_param* param);
static struct kdb_keyinfo_pkt* ps2_keybuffer_next_write();

// TODO: Abstract the bounded buffer out.
void ps2_device_post_cmd(char cmd, char arg) {
    mutex_lock(&cmd_q.mutex);
    int index = (cmd_q.queue_ptr + cmd_q.queue_len) % PS2_CMD_QUEUE_SIZE;
    if (index == cmd_q.queue_ptr && cmd_q.queue_len) {
        // 队列已满！
        mutex_unlock(&cmd_q.mutex);
        return;
    }

    struct ps2_cmd *container = &cmd_q.cmd_queue[index];
    container->cmd = cmd;
    container->arg = arg;
    cmd_q.queue_len++;

    // 释放锁，同理。
    mutex_unlock(&cmd_q.mutex);
}

void ps2_kbd_init() {

    memset(&cmd_q, 0, sizeof(cmd_q));
    memset(&key_buf, 0, sizeof(key_buf));
    memset(&kbd_state, 0, sizeof(kbd_state));
    
    mutex_init(&cmd_q.mutex);
    mutex_init(&key_buf.mutex);


    kbd_state.translation_table = scancode_set2;
    kbd_state.state = KBD_STATE_KWAIT;

    cpu_disable_interrupt();

    // XXX: 是否需要使用FADT探测PS/2控制器的存在？

    // 1、禁用任何的PS/2设备
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, PS2_CMD_PORT1_DISABLE, PS2_NO_ARG);
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, PS2_CMD_PORT2_DISABLE, PS2_NO_ARG);
    
    // 2、清空控制器缓冲区
    io_inb(PS2_PORT_ENC_DATA);

    char result;

    // 3、屏蔽所有PS/2设备（端口1&2）IRQ，并且禁用键盘键码转换功能
    result = ps2_issue_cmd(PS2_CMD_READ_CFG, PS2_NO_ARG);
    result = result & ~(PS2_CFG_P1INT | PS2_CFG_P2INT | PS2_CFG_TRANSLATION);
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, PS2_CMD_WRITE_CFG, result);

    // 4、控制器自检
    result = ps2_issue_cmd(PS2_CMD_SELFTEST, PS2_NO_ARG);
    if (result != PS2_RESULT_TEST_OK) {
        kprintf(KERROR "Controller self-test failed.");
        goto done;
    }

    // 5、设备自检（端口1自检，通常是我们的键盘）
    result = ps2_issue_cmd(PS2_CMD_SELFTEST_PORT1, PS2_NO_ARG);
    if (result != 0) {
        kprintf(KERROR "Interface test on port 1 failed.");
        goto done;
    }

    // 6、开启位于端口1的 IRQ，并启用端口1。不用理会端口2，那儿一般是鼠标。
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, PS2_CMD_PORT1_ENABLE, PS2_NO_ARG);
    result = ps2_issue_cmd(PS2_CMD_READ_CFG, PS2_NO_ARG);
    result = result | PS2_CFG_P1INT;
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, PS2_CMD_WRITE_CFG, result);

    // 至此，PS/2控制器和设备已完成初始化，可以正常使用。

    // 将我们的键盘驱动挂载到第204号中断上（已由IOAPIC映射至IRQ#1），
    intr_subscribe(PC_KBD_IV, intr_ps2_kbd_handler);

    // 搞一个计时器，将我们的 ps2_process_cmd 挂上去。每隔5毫秒执行排在队头的命令。
    //  为什么只执行队头的命令，而不是全部的命令？
    //      因为我们需要保证isr尽量的简短，运行起来快速。而发送这些命令非常的耗时。
    timer_run_ms(5, ps2_process_cmd, NULL, TIMER_MODE_PERIODIC);

done:
    cpu_enable_interrupt();
}

void ps2_process_cmd(void* arg) {
    // 检查锁是否已被启用，如果启用，则表明该timer中断发生时，某个指令正在入队。
    // 如果是这种情况则跳过，留到下一轮再尝试处理。
    // 注意，这里其实是ISR的一部分（timer中断），对于单核CPU来说，ISR等同于单个的原子操作。
    // （因为EFLAGS.IF=0，所有可屏蔽中断被屏蔽。对于NMI的情况，那么就直接算是triple fault了，所以也没有讨论的意义）
    // 所以，假若我们遵从互斥锁的严格定义（即这里需要阻塞），那么中断将会被阻塞，进而造成死锁。
    // 因此，我们这里仅仅进行判断。
    // 会不会产生指令堆积？不会，因为指令发送的频率远远低于指令队列清空的频率。在目前，我们发送的唯一指令
    // 就只是用来开关键盘上的LED灯（如CAPSLOCK）。
    if (mutex_on_hold(&cmd_q.mutex) || !cmd_q.queue_len) {
        return;
    }

    // 处理队列排头的指令
    struct ps2_cmd *pending_cmd = &cmd_q.cmd_queue[cmd_q.queue_ptr];
    char result;
    int attempts = 0;

    // 尝试将命令发送至PS/2键盘（通过PS/2控制器）
    // 如果不成功（0x60 IO口返回 0xfe，即 NAK 或 Resend）
    // 则尝试最多五次
    do {
        result = ps2_issue_dev_cmd(pending_cmd->cmd, pending_cmd->arg);
        kbd_state.state += KBD_STATE_CMDPROCS;
        attempts++;
    } while(result == PS2_RESULT_NAK && attempts < PS2_DEV_CMD_MAX_ATTEMPTS);
    
    // XXX: 是否需要处理不成功的指令？

    cmd_q.queue_ptr = (cmd_q.queue_ptr + 1) % PS2_CMD_QUEUE_SIZE;
    cmd_q.queue_len--;
}

void kbd_buffer_key_event(kbd_keycode_t key, uint8_t scancode, kbd_kstate_t state) {
    // forgive me on these ugly bit-level tricks, 
    // I really hate doing branching on these "fliping switch" things
    if (key == KEY_CAPSLK) {
        kbd_state.key_state ^= KBD_KEY_FCAPSLKED & -state;
    } else if (key == KEY_NUMSLK) {
        kbd_state.key_state ^= KBD_KEY_FNUMBLKED & -state;
    } else if (key == KEY_SCRLLK) {
        kbd_state.key_state ^= KBD_KEY_FSCRLLKED & -state;
    } else {
        if ((key & MODIFR)) {
            kbd_kstate_t tmp = (KBD_KEY_FLSHIFT_HELD << (key & 0x00ff));
            kbd_state.key_state = (kbd_state.key_state & ~tmp) | (tmp & -state);
        }
        else if (!(key & 0xff00) && (kbd_state.key_state & (KBD_KEY_FLSHIFT_HELD | KBD_KEY_FRSHIFT_HELD))) {
            key = scancode_set2_shift[scancode];
        }
        state = state | kbd_state.key_state;
        key = key & (0xffdf | -('a' > key || key > 'z' || !(state & KBD_KEY_FCAPSLKED)));

        if (!mutex_on_hold(&key_buf.mutex)) {
            struct kdb_keyinfo_pkt* keyevent_pkt = ps2_keybuffer_next_write();
            *keyevent_pkt = (struct kdb_keyinfo_pkt) {
                .keycode = key,
                .scancode = scancode,
                .state = state,
                .timestamp = clock_systime()
            };
        }

        // kprintf(KDEBUG "%c (t=%d, s=%x, c=%d)\n", key & 0x00ff, timestamp, state, key >> 8);
        return; // do not delete this return
    }

    // Ooops, this guy generates irq!
    ps2_device_post_cmd(PS2_KBD_CMD_SETLED, (kbd_state.key_state >> 1) & 0x00ff);
}

void intr_ps2_kbd_handler(const isr_param* param) {

    // Do not move this line. It is in the right place and right order.
    // This is to ensure we've cleared the output buffer everytime, so it won't pile up across irqs.
    uint8_t scancode = io_inb(PS2_PORT_ENC_DATA);
    kbd_keycode_t key;

    /*  
     *    判断键盘是否处在指令发送状态，防止误触发。（伪输入中断）
     * 这是因为我们需要向ps/2设备发送指令（比如控制led灯），而指令会有返回码。
     * 这就会有可能导致ps/2控制器在受到我们的命令后（在ps2_process_cmd中），
     * 产生IRQ#1中断（虽然说这种情况取决于底层BIOS实现，但还是会发生，比如QEMU和bochs）。
     * 所以这就是说，当IRQ#1中断产生时，我们的CPU正处在另一个ISR中。这样就会导致所有的外部中断被缓存在APIC内部的
     * FIFO队列里，进行排队等待（APIC长度为二的队列 {IRR, TMR}；参考 Intel Manual Vol.3A 10.8.4）
     * 那么当ps2_process_cmd执行完后（内嵌在#APIC_TIMER_IV），CPU返回EOI给APIC，APIC紧接着将排在队里的IRQ#1发送给CPU
     * 造成误触发。也就是说，我们此时读入的scancode实则上是上一个指令的返回代码。
     * 
     * Problem 1 (Fixed):
     *      但是这种方法有个问题，那就是，假若我们的某一个命令失败了一次，ps/2给出0xfe，我们重传，ps/2收到指令并给出0xfa。
     *  那么这样一来，将会由两个连续的IRQ#1产生。而APIC是最多可以缓存两个IRQ，于是我们就会漏掉一个IRQ，依然会误触发。
     * Solution:
     *      累加掩码 ;)
     */
    if ((kbd_state.state & 0xc0)) {
        kbd_state.state -= KBD_STATE_CMDPROCS;
        return;
    }
    
    //kprintf(KINFO "%x\n", scancode & 0xff);
    
    switch (kbd_state.state)
    {
    case KBD_STATE_KWAIT:
        if (scancode == 0xf0) { // release code
            kbd_state.state = KBD_STATE_KRELEASED;       
        } else if (scancode == 0xe0) {
            kbd_state.state = KBD_STATE_KSPECIAL;
            kbd_state.translation_table = scancode_set2_ex;
        } else {
            key = kbd_state.translation_table[scancode];
            kbd_buffer_key_event(key, scancode, KBD_KEY_FPRESSED);
        }
        break;
    case KBD_STATE_KSPECIAL:
        if (scancode == 0xf0) { //release code
            kbd_state.state = KBD_STATE_KRELEASED;       
        } else {
            key = kbd_state.translation_table[scancode];
            kbd_buffer_key_event(key, scancode, KBD_KEY_FPRESSED);

            kbd_state.state = KBD_STATE_KWAIT;
            kbd_state.translation_table = scancode_set2;
        }
        break;
    case KBD_STATE_KRELEASED:
        key = kbd_state.translation_table[scancode];
        kbd_buffer_key_event(key, scancode, KBD_KEY_FRELEASED);
        
        // reset the translation table to scancode_set2
        kbd_state.state = KBD_STATE_KWAIT;   
        kbd_state.translation_table = scancode_set2;
        break;
    
    default:
        break;
    }
}

static uint8_t ps2_issue_cmd(char cmd, uint16_t arg) {
    ps2_post_cmd(PS2_PORT_CTRL_CMDREG, cmd, arg);

    char result;
    
    // 等待PS/2控制器返回。通过轮询（polling）状态寄存器的 bit 0
    // 如置位，则表明返回代码此时就在 0x60 IO口上等待读取。
    while(!((result = io_inb(PS2_PORT_CTRL_STATUS)) & PS2_STATUS_OFULL));

    return io_inb(PS2_PORT_ENC_CMDREG);
}

static uint8_t ps2_issue_dev_cmd(char cmd, uint16_t arg) {
    ps2_post_cmd(PS2_PORT_ENC_CMDREG, cmd, arg);

    char result;
    
    // 等待PS/2控制器返回。通过轮询（polling）状态寄存器的 bit 0
    // 如置位，则表明返回代码此时就在 0x60 IO口上等待读取。
    while(!((result = io_inb(PS2_PORT_CTRL_STATUS)) & PS2_STATUS_OFULL));

    return io_inb(PS2_PORT_ENC_CMDREG);
}

static void ps2_post_cmd(uint8_t port, char cmd, uint16_t arg) {
    char result;
    // 等待PS/2输入缓冲区清空，这样我们才可以写入命令
    while((result = io_inb(PS2_PORT_CTRL_STATUS)) & PS2_STATUS_IFULL);

    io_outb(port, cmd);
    if (!(arg & PS2_NO_ARG)) {
        // 所有参数一律通过0x60传入。
        io_outb(PS2_PORT_ENC_CMDREG, (uint8_t)(arg & 0x00ff));
    }
}

int kbd_recv_key(struct kdb_keyinfo_pkt* key_event) {
    if (!key_buf.buffered_len) {
        return 0;
    }
    mutex_lock(&key_buf.mutex);

    struct kdb_keyinfo_pkt* pkt_current = &key_buf.buffer[key_buf.read_ptr];

    *key_event = *pkt_current;
    key_buf.buffered_len--;
    key_buf.read_ptr = (key_buf.read_ptr + 1) % PS2_KBD_RECV_BUFFER_SIZE;

    mutex_unlock(&key_buf.mutex);
    return 1;
}

static struct kdb_keyinfo_pkt* ps2_keybuffer_next_write() {
    int index = (key_buf.read_ptr + key_buf.buffered_len) % PS2_KBD_RECV_BUFFER_SIZE;
    if (index == key_buf.read_ptr && key_buf.buffered_len) {
        // the reader is lagged so much such that the buffer is full.
        // It is suggested to read from beginning for nearly up-to-date readings.
        key_buf.read_ptr = 0;
        key_buf.buffered_len = index;
    }
    else {
        key_buf.buffered_len++;
    }
    return &key_buf.buffer[index];
}