/**
 * @file keyboard.c
 * @brief PS/2 keyboard driver implementation.
 *
 * Reads scancodes from I/O port 0x60 on IRQ1, translates scancode set 1
 * to ASCII using a simple lookup table, stores characters in a ring buffer,
 * and wakes any process blocked waiting for keyboard input.
 */

#include "keyboard.h"
#include "port_io.h"
#include "pic.h"
#include "process.h"
#include <stddef.h>
#include <string.h>

/* Debug print helpers defined in kernel.c */
extern void offset_print(const char *str);
extern void print_hex(uint32_t val);

/** PS/2 keyboard data port. */
#define KB_DATA_PORT 0x60

/** Ring buffer for keyboard input. */
static char kb_buffer[KB_BUFFER_SIZE];
static volatile uint32_t kb_head = 0;
static volatile uint32_t kb_tail = 0;

/** Process waiting for keyboard input (PID, or 0 if none). */
static volatile uint32_t kb_waiting_pid = 0;

/** Shift key state. */
static int shift_pressed = 0;

/**
 * Scancode set 1 to ASCII lookup table (unshifted).
 * Index = scancode, value = ASCII character (0 = unmapped).
 */
static const char scancode_ascii[128] = {
    0,   27, '1', '2', '3', '4', '5', '6',  /* 0x00 - 0x07 */
    '7', '8', '9', '0', '-', '=', '\b', '\t', /* 0x08 - 0x0F */
    'q', 'w', 'e', 'r', 't', 'y', 'u', 'i',  /* 0x10 - 0x17 */
    'o', 'p', '[', ']', '\n', 0,   'a', 's',  /* 0x18 - 0x1F */
    'd', 'f', 'g', 'h', 'j', 'k', 'l', ';',  /* 0x20 - 0x27 */
    '\'', '`', 0,  '\\', 'z', 'x', 'c', 'v', /* 0x28 - 0x2F */
    'b', 'n', 'm', ',', '.', '/', 0,   '*',   /* 0x30 - 0x37 */
    0,   ' ', 0,   0,   0,   0,   0,   0,     /* 0x38 - 0x3F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x40 - 0x47 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x48 - 0x4F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x50 - 0x57 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x58 - 0x5F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x60 - 0x67 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x68 - 0x6F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x70 - 0x77 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x78 - 0x7F */
};

/**
 * Scancode set 1 to ASCII lookup table (shifted).
 */
static const char scancode_ascii_shift[128] = {
    0,   27, '!', '@', '#', '$', '%', '^',     /* 0x00 - 0x07 */
    '&', '*', '(', ')', '_', '+', '\b', '\t',  /* 0x08 - 0x0F */
    'Q', 'W', 'E', 'R', 'T', 'Y', 'U', 'I',  /* 0x10 - 0x17 */
    'O', 'P', '{', '}', '\n', 0,   'A', 'S',  /* 0x18 - 0x1F */
    'D', 'F', 'G', 'H', 'J', 'K', 'L', ':',  /* 0x20 - 0x27 */
    '"', '~', 0,   '|', 'Z', 'X', 'C', 'V',  /* 0x28 - 0x2F */
    'B', 'N', 'M', '<', '>', '?', 0,   '*',   /* 0x30 - 0x37 */
    0,   ' ', 0,   0,   0,   0,   0,   0,     /* 0x38 - 0x3F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x40 - 0x47 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x48 - 0x4F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x50 - 0x57 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x58 - 0x5F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x60 - 0x67 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x68 - 0x6F */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x70 - 0x77 */
    0,   0,   0,   0,   0,   0,   0,   0,     /* 0x78 - 0x7F */
};

/** Left/right shift scancodes. */
#define SC_LSHIFT_PRESS   0x2A
#define SC_LSHIFT_RELEASE 0xAA
#define SC_RSHIFT_PRESS   0x36
#define SC_RSHIFT_RELEASE 0xB6

/**
 * Put a character into the ring buffer.
 */
static void kb_buffer_put(char c) {
    uint32_t next = (kb_head + 1) % KB_BUFFER_SIZE;
    if (next == kb_tail) {
        return; /* Buffer full, drop character */
    }
    kb_buffer[kb_head] = c;
    kb_head = next;
}

void keyboard_init(void) {
    kb_head = 0;
    kb_tail = 0;
    kb_waiting_pid = 0;
    shift_pressed = 0;

    offset_print("  KEYBOARD: flushing controller...\n");

    /* Flush any pending data from the keyboard controller.
     * Use a timeout to avoid hanging if the controller keeps reporting data
     * (some emulators/VMs behave differently). */
    int flush_count = 0;
    while ((inb(0x64) & 0x01) && flush_count < 1024) {
        inb(KB_DATA_PORT);
        flush_count++;
    }

    offset_print("  KEYBOARD: flushed ");
    print_hex((uint32_t)flush_count);
    offset_print("  KEYBOARD: bytes, unmasking IRQ1...\n");

    /* Unmask IRQ1 (keyboard) in the PIC */
    pic_clear_mask(1);

    offset_print("  KEYBOARD: initialized\n");
}

void keyboard_irq(registers_t *regs) {
    (void)regs;
    uint8_t scancode = inb(KB_DATA_PORT);

    /* Handle shift keys */
    if (scancode == SC_LSHIFT_PRESS || scancode == SC_RSHIFT_PRESS) {
        shift_pressed = 1;
        return;
    }
    if (scancode == SC_LSHIFT_RELEASE || scancode == SC_RSHIFT_RELEASE) {
        shift_pressed = 0;
        return;
    }

    /* Ignore key releases (bit 7 set) */
    if (scancode & 0x80) {
        return;
    }

    /* Translate scancode to ASCII */
    char c;
    if (shift_pressed) {
        c = scancode_ascii_shift[scancode];
    } else {
        c = scancode_ascii[scancode];
    }

    if (c == 0) {
        return; /* Unmapped key */
    }

    /* Put character in buffer */
    kb_buffer_put(c);

    /* Wake any process waiting for keyboard input */
    if (kb_waiting_pid != 0) {
        process_t *waiter = process_get(kb_waiting_pid);
        if (waiter && waiter->state == PROCESS_BLOCKED) {
            waiter->state = PROCESS_READY;
        }
        kb_waiting_pid = 0;
    }
}

uint32_t keyboard_read(char *buf, uint32_t count) {
    uint32_t n = 0;
    while (n < count && kb_tail != kb_head) {
        buf[n++] = kb_buffer[kb_tail];
        kb_tail = (kb_tail + 1) % KB_BUFFER_SIZE;
    }
    return n;
}

int keyboard_has_data(void) {
    return kb_head != kb_tail;
}

void keyboard_block_for_input(registers_t *regs) {
    process_t *cur = process_current();
    if (!cur) return;

    cur->state = PROCESS_BLOCKED;
    cur->saved_regs = *regs;

    /* Rewind EIP by 2 bytes so that when the process is unblocked and
     * scheduled, the CPU re-executes the INT 0x80 instruction. At that
     * point keyboard_has_data() will return true and the read succeeds. */
    cur->saved_regs.eip -= 2;

    kb_waiting_pid = cur->pid;

    /* Schedule next process */
    schedule_tick(regs);
}