#include "forth.h"
#include <stdint.h>
#include <stdarg.h>

// VGA text-mode buffer
static uint16_t* vga_buffer = (uint16_t*)0xB8000;
static int vga_row = 0;
static int vga_col = 0;

static void vga_putchar(char c) {
    if (c == '\n') {
        vga_col = 0;
        vga_row++;
    } else {
        int idx = vga_row * 80 + vga_col;
        vga_buffer[idx] = (uint16_t)((0x0F << 8) | (unsigned char)c);
        vga_col++;
        if (vga_col >= 80) {
            vga_col = 0;
            vga_row++;
        }
    }
    if (vga_row >= 25) {
        vga_row = 0;
        vga_col = 0;
    }
}

static inline void outb(uint16_t port, uint8_t val) {
    __asm__ volatile ("outb %0, %1" : : "a"(val), "Nd"(port));
}

static inline uint8_t inb(uint16_t port) {
    uint8_t ret;
    __asm__ volatile ("inb %1, %0" : "=a"(ret) : "Nd"(port));
    return ret;
}

static inline void io_wait(void) {
    __asm__ volatile ("outb %0, %1" : : "a"((uint8_t)0), "Nd"((uint16_t)0x80));
}

static void serial_init(void) {
    outb(0x3F8 + 1, 0x00);
    outb(0x3F8 + 3, 0x80);
    outb(0x3F8 + 0, 0x01);
    outb(0x3F8 + 1, 0x00);
    outb(0x3F8 + 3, 0x03);
    outb(0x3F8 + 2, 0xC7);
    outb(0x3F8 + 4, 0x0B);
}

static uint64_t idt_entries[256][2];

static void idt_set_gate(int vector, uint64_t offset, uint16_t selector, uint8_t type_attr) {
    uint64_t low = (offset & 0xFFFFULL) |
                   ((uint64_t)selector << 16) |
                   ((uint64_t)(type_attr & 0xFF) << 40) |
                   ((offset & 0xFFFF0000ULL) << 32);
    uint64_t high = offset >> 32;
    idt_entries[vector][0] = low;
    idt_entries[vector][1] = high;
}

static struct {
    uint16_t limit;
    uint64_t base;
} __attribute__((packed)) idt_ptr;

static void pic_remap(void) {
    uint8_t a1 = inb(0x21);
    uint8_t a2 = inb(0xA1);

    outb(0x20, 0x11);
    io_wait();
    outb(0xA0, 0x11);
    io_wait();
    outb(0x21, 0x20);
    io_wait();
    outb(0xA1, 0x28);
    io_wait();
    outb(0x21, 0x04);
    io_wait();
    outb(0xA1, 0x02);
    io_wait();
    outb(0x21, 0x01);
    io_wait();
    outb(0xA1, 0x01);
    io_wait();

    outb(0x21, a1);
    outb(0xA1, a2);
}

static void pic_unmask_keyboard(void) {
    uint8_t mask = inb(0x21);
    mask &= ~(1 << 1);
    outb(0x21, mask);
}

#define KEYBUF_SIZE 256
static volatile char keybuf[KEYBUF_SIZE];
static volatile int keybuf_head = 0;
static volatile int keybuf_tail = 0;

static const char scancode_map[] = {
    0,  27, '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', '-', '=', '\b',
    '\t', 'q', 'w', 'e', 'r', 't', 'y', 'u', 'i', 'o', 'p', '[', ']', '\n',
    0, 'a', 's', 'd', 'f', 'g', 'h', 'j', 'k', 'l', ';', '\'', '`',
    0, '\\', 'z', 'x', 'c', 'v', 'b', 'n', 'm', ',', '.', '/', 0,
    '*', 0, ' ', 0
};

void keyboard_handler(void) {
    uint8_t scancode = inb(0x60);
    if (scancode < sizeof(scancode_map)) {
        char c = scancode_map[scancode];
        if (c) {
            int next = (keybuf_tail + 1) % KEYBUF_SIZE;
            if (next != keybuf_head) {
                keybuf[keybuf_tail] = c;
                keybuf_tail = next;
                forth_putchar(c);
            }
        }
    }
    outb(0x20, 0x20);
}

void forth_putchar(char c) {
    vga_putchar(c);
    while ((inb(0x3F8 + 5) & 0x20) == 0) { }
    outb(0x3F8, (uint8_t)c);
}

int forth_getchar(void) {
    while (keybuf_head == keybuf_tail) {
        __asm__ volatile ("pause");
    }
    char c = keybuf[keybuf_head];
    keybuf_head = (keybuf_head + 1) % KEYBUF_SIZE;
    return (int)c;
}

void forth_fflush(void) {
}

void forth_panic(void) {
    __asm__ volatile ("cli; hlt");
    __builtin_unreachable();
}

int64_t forth_strtoll(const char* str, char** endptr, int base) {
    const char* s = str;
    int neg = 0;
    int64_t val = 0;
    while (forth_isspace(*s)) s++;
    if (*s == '-') { neg = 1; s++; }
    else if (*s == '+') { s++; }
    while (*s >= '0' && *s <= '9') {
        val = val * base + (*s - '0');
        s++;
    }
    if (endptr) *endptr = (char*)s;
    return neg ? -val : val;
}

static void print_str(const char* s) {
    while (*s) forth_putchar(*s++);
}

static void print_int(long long v) {
    char buf[32];
    int i = 30;
    int neg = 0;
    if (v < 0) { neg = 1; v = -v; }
    buf[31] = '\0';
    if (v == 0) {
        buf[i--] = '0';
    } else {
        while (v > 0 && i >= 0) {
            buf[i--] = '0' + (v % 10);
            v /= 10;
        }
    }
    if (neg) buf[i--] = '-';
    print_str(&buf[i + 1]);
}

void forth_printf(const char* fmt, ...) {
    va_list ap;
    va_start(ap, fmt);
    while (*fmt) {
        if (*fmt == '%' && *(fmt + 1)) {
            fmt++;
            if (*fmt == 's') {
                print_str(va_arg(ap, const char*));
            } else if (*fmt == 'c') {
                forth_putchar((char)va_arg(ap, int));
            } else if (*fmt == '%') {
                forth_putchar('%');
            } else if (*fmt == 'l') {
                if (*(fmt + 1) == 'l' && (*(fmt + 2) == 'd' || *(fmt + 2) == 'i' || *(fmt + 2) == 'u')) {
                    fmt += 2;
                    print_int(va_arg(ap, long long));
                }
            } else if (*fmt == 'd' || *fmt == 'i' || *fmt == 'u') {
                print_int(va_arg(ap, int));
            } else {
                forth_putchar('%');
                forth_putchar(*fmt);
            }
        } else {
            forth_putchar(*fmt);
        }
        fmt++;
    }
    va_end(ap);
}

void kernel_main(void) {
    extern void isr_keyboard(void);

    serial_init();

    idt_set_gate(0x21, (uint64_t)isr_keyboard, 0x08, 0x8E);
    idt_ptr.limit = sizeof(idt_entries) - 1;
    idt_ptr.base = (uint64_t)idt_entries;
    __asm__ volatile ("lidt %0" : : "m"(idt_ptr));

    pic_remap();
    outb(0x21, 0xFF); // mask all IRQs on master
    outb(0xA1, 0xFF); // mask all IRQs on slave
    pic_unmask_keyboard();

    __asm__ volatile ("sti");

    forth_printf("Kernel started\n");
    forth_run();
    forth_panic();
}