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

extern const uint8_t font8x8[96][8];

struct mboot_tag {
    uint32_t type;
    uint32_t size;
};

struct mboot_tag_framebuffer {
    uint32_t type;
    uint32_t size;
    uint64_t framebuffer_addr;
    uint32_t framebuffer_pitch;
    uint32_t framebuffer_width;
    uint32_t framebuffer_height;
    uint8_t framebuffer_bpp;
    uint8_t framebuffer_type;
    uint16_t reserved;
};

static uint64_t fb_addr = 0;
static uint32_t fb_pitch = 0;
static uint32_t fb_width = 0;
static uint32_t fb_height = 0;
static uint8_t fb_bpp = 0;
static int fb_cols = 0;
static int fb_rows = 0;
static int fb_col = 0;
static int fb_row = 0;

static void fb_draw_pixel(int x, int y, uint32_t color) {
    if (x < 0 || y < 0 || (uint32_t)x >= fb_width || (uint32_t)y >= fb_height) return;
    if (!fb_addr) return;
    if (fb_bpp == 32) {
        uint32_t* p = (uint32_t*)(fb_addr + y * fb_pitch + x * 4);
        *p = color;
    } else if (fb_bpp == 24) {
        uint8_t* p = (uint8_t*)(fb_addr + y * fb_pitch + x * 3);
        p[0] = color & 0xFF;
        p[1] = (color >> 8) & 0xFF;
        p[2] = (color >> 16) & 0xFF;
    }
}

static void fb_scroll(void) {
    int line_height = 8;
    uint32_t bytes_per_pixel = (fb_bpp + 7) / 8;
    uint32_t copy_bytes = fb_width * bytes_per_pixel;
    for (uint32_t y = 0; y < fb_height - line_height; y++) {
        uint8_t* dst = (uint8_t*)(fb_addr + y * fb_pitch);
        uint8_t* src = (uint8_t*)(fb_addr + (y + line_height) * fb_pitch);
        for (uint32_t i = 0; i < copy_bytes; i++) {
            dst[i] = src[i];
        }
    }
    for (uint32_t y = fb_height - line_height; y < fb_height; y++) {
        uint8_t* row = (uint8_t*)(fb_addr + y * fb_pitch);
        for (uint32_t i = 0; i < copy_bytes; i++) {
            row[i] = 0;
        }
    }
}

static void fb_draw_char(char c) {
    if (!fb_addr || fb_cols == 0) return;
    if (c == '\n') {
        fb_col = 0;
        fb_row++;
    } else if (c == '\r') {
        fb_col = 0;
    } else if (c == '\b') {
        if (fb_col > 0) fb_col--;
    } else {
        unsigned char ch = (unsigned char)c;
        if (ch < 32 || ch > 126) ch = '?';
        const uint8_t* glyph = font8x8[ch - 32];
        int start_x = fb_col * 8;
        int start_y = fb_row * 8;
        for (int dy = 0; dy < 8; dy++) {
            uint8_t row = glyph[dy];
            for (int dx = 0; dx < 8; dx++) {
                if (row & (0x80 >> dx)) {
                    fb_draw_pixel(start_x + dx, start_y + dy, 0xFFFFFF);
                } else {
                    fb_draw_pixel(start_x + dx, start_y + dy, 0x000000);
                }
            }
        }
        fb_col++;
    }
    if (fb_col >= fb_cols) {
        fb_col = 0;
        fb_row++;
    }
    if (fb_row >= fb_rows) {
        fb_scroll();
        fb_row = fb_rows - 1;
    }
}

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);
}

/* Each 64-bit IDT entry is 16 bytes = two uint64_t words.
 *
 * Word 0 layout:
 *   [15: 0] offset[15:0]
 *   [31:16] segment selector
 *   [34:32] IST (0 = no IST)
 *   [39:35] reserved (0)
 *   [47:40] type+attr  (0x8E = present, DPL=0, 64-bit interrupt gate)
 *   [63:48] offset[31:16]
 *
 * Word 1 layout:
 *   [31: 0] offset[63:32]
 *   [63:32] reserved (0)
 */
typedef struct {
    uint64_t low;
    uint64_t high;
} __attribute__((packed)) IdtEntry;

static IdtEntry idt_entries[256];

static void idt_set_gate(int vector, uint64_t offset, uint16_t selector, uint8_t type_attr) {
    uint64_t low =
        (offset & 0x000000000000FFFFULL)          |   /* offset[15:0]  */
        ((uint64_t)selector          << 16)        |   /* selector      */
        /* bits 34:32 = IST=0, bits 39:35 = 0 */
        ((uint64_t)type_attr         << 40)        |   /* type+attr     */
        ((offset  & 0x00000000FFFF0000ULL) << 32); /* offset[31:16] */
    uint64_t high = offset >> 32;                      /* offset[63:32] */
    idt_entries[vector].low  = low;
    idt_entries[vector].high = 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) {
    fb_draw_char(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(uint64_t mboot_ptr) {
    extern void isr_keyboard(void);

    /* Multiboot2 info tag types:
     *   5 = BIOS boot device, 6 = memory map, 8 = framebuffer info, ...
     * Note: in the *header* (request side), type 5 means "framebuffer
     * request", but the bootloader's *response* tag uses type 8. */
    struct mboot_tag* tag = (struct mboot_tag*)(mboot_ptr + 8);
    while (tag->type != 0) {
        if (tag->type == 8) {
            struct mboot_tag_framebuffer* fb = (struct mboot_tag_framebuffer*)tag;
            fb_addr = fb->framebuffer_addr;
            fb_pitch = fb->framebuffer_pitch;
            fb_width = fb->framebuffer_width;
            fb_height = fb->framebuffer_height;
            fb_bpp = fb->framebuffer_bpp;
        }
        tag = (struct mboot_tag*)((uint8_t*)tag + ((tag->size + 7) & ~7));
    }

    if (fb_addr && fb_bpp >= 24) {
        fb_cols = fb_width / 8;
        fb_rows = fb_height / 8;
        for (uint32_t y = 0; y < fb_height; y++) {
            for (uint32_t x = 0; x < fb_width; x++) {
                fb_draw_pixel(x, y, 0);
            }
        }
    }

    serial_init();

    idt_set_gate(0x21, (uint64_t)isr_keyboard, 0x08, 0x8E);
    idt_ptr.limit = (uint16_t)(sizeof(idt_entries) - 1);
    idt_ptr.base  = (uint64_t)(uintptr_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();
}

void do_fb_addr(Word* w) {
    (void)w;
    data_push((int64_t)fb_addr);
}

void do_fb_width(Word* w) {
    (void)w;
    data_push((int64_t)fb_width);
}

void do_fb_height(Word* w) {
    (void)w;
    data_push((int64_t)fb_height);
}

void do_fb_pitch(Word* w) {
    (void)w;
    data_push((int64_t)fb_pitch);
}

void do_fb_bpp(Word* w) {
    (void)w;
    data_push((int64_t)fb_bpp);
}

void do_fb_plot(Word* w) {
    (void)w;
    int64_t color = data_pop();
    int64_t y = data_pop();
    int64_t x = data_pop();
    fb_draw_pixel((int)x, (int)y, (uint32_t)color);
}