#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include "gdt.h"

#define IDT_ENTRIES 256

struct idt_entry {
    uint16_t base_lo;
    uint16_t sel;
    uint8_t always0;
    uint8_t flags;
    uint16_t base_hi;
} __attribute__((packed));

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

struct idt_entry idt[IDT_ENTRIES];
struct idt_ptr idtp;

extern void idt_load();
extern void isr_handler();
extern void irq_handler();
extern void pit_handler_asm();

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

void idt_set_gate(uint8_t num, uint32_t base, uint16_t sel, uint8_t flags) {
    idt[num].base_lo = base & 0xFFFF;
    idt[num].base_hi = (base >> 16) & 0xFFFF;
    idt[num].sel = sel;
    idt[num].always0 = 0;
    idt[num].flags = flags | 0x60; // Set the interrupt gate flag
}

void idt_install() {
    idtp.limit = (sizeof(struct idt_entry) * IDT_ENTRIES) - 1;
    idtp.base = (uint32_t)&idt;
    idt_load();
}

void pic_remap(int offset1, int offset2) {
    unsigned char a1, a2;

    a1 = inb(0x21);
    a2 = inb(0xA1);

    outb(0x20, 0x11);
    outb(0xA0, 0x11);
    outb(0x21, offset1);
    outb(0xA1, offset2);
    outb(0x21, 0x04);
    outb(0xA1, 0x02);

    outb(0x21, 0x01);
    outb(0xA1, 0x01);

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

enum vga_color {
    VGA_COLOR_BLACK = 0,
    VGA_COLOR_BLUE = 1,
    VGA_COLOR_GREEN = 2,
    VGA_COLOR_CYAN = 3,
    VGA_COLOR_RED = 4,
    VGA_COLOR_MAGENTA = 5,
    VGA_COLOR_BROWN = 6,
    VGA_COLOR_LIGHT_GREY = 7,
    VGA_COLOR_DARK_GREY = 8,
    VGA_COLOR_LIGHT_BLUE = 9,
    VGA_COLOR_LIGHT_GREEN = 10,
    VGA_COLOR_LIGHT_CYAN = 11,
    VGA_COLOR_LIGHT_RED = 12,
    VGA_COLOR_LIGHT_MAGENTA = 13,
    VGA_COLOR_LIGHT_BROWN = 14,
    VGA_COLOR_WHITE = 15,
};

static inline uint8_t vga_entry_color(enum vga_color fg, enum vga_color bg) {
    return fg | bg << 4;
}

static inline uint16_t vga_entry(unsigned char uc, uint8_t color) {
    return (uint16_t) uc | (uint16_t) color << 8;
}

size_t strlen(const char* str) {
    size_t len = 0;
    while (str[len])
        len++;
    return len;
}

static const size_t VGA_WIDTH = 80;
static const size_t VGA_HEIGHT = 25;

size_t terminal_row;
size_t terminal_column;
uint8_t terminal_color;
uint16_t* terminal_buffer;

void terminal_initialize(void) {
    terminal_row = 0;
    terminal_column = 0;
    terminal_color = vga_entry_color(VGA_COLOR_LIGHT_GREY, VGA_COLOR_BLACK);
    terminal_buffer = (uint16_t*) 0xB8000;
    for (size_t y = 0; y < VGA_HEIGHT; y++) {
        for (size_t x = 0; x < VGA_WIDTH; x++) {
            const size_t index = y * VGA_WIDTH + x;
            terminal_buffer[index] = vga_entry(' ', terminal_color);
        }
    }
}

void terminal_setcolor(uint8_t color) {
    terminal_color = color;
}

void terminal_putentryat(char c, uint8_t color, size_t x, size_t y) {
    const size_t index = y * VGA_WIDTH + x;
    terminal_buffer[index] = vga_entry(c, color);
}

void terminal_scroll(void) {
    for (size_t y = 0; y < VGA_HEIGHT - 1; y++) {
        for (size_t x = 0; x < VGA_WIDTH; x++) {
            terminal_buffer[y * VGA_WIDTH + x] = terminal_buffer[(y + 1) * VGA_WIDTH + x];
        }
    }
    for (size_t x = 0; x < VGA_WIDTH; x++) {
        terminal_buffer[(VGA_HEIGHT - 1) * VGA_WIDTH + x] = vga_entry(' ', terminal_color);
    }
    terminal_row = VGA_HEIGHT - 1;
}

void terminal_putchar(char c) {
    if (c == '\n') {
        terminal_column = 0;
        if (++terminal_row == VGA_HEIGHT) {
            terminal_scroll();
        }
    } else {
        terminal_putentryat(c, terminal_color, terminal_column, terminal_row);
        if (++terminal_column == VGA_WIDTH) {
            terminal_column = 0;
            if (++terminal_row == VGA_HEIGHT) {
                terminal_scroll();
            }
        }
    }
}

void terminal_write(const char* data, size_t size) {
    for (size_t i = 0; i < size; i++)
        terminal_putchar(data[i]);
}

void terminal_writestring(const char* data) {
    terminal_write(data, strlen(data));
}

volatile uint32_t tick = 0;

void pit_handler() {
    tick++;
    outb(0x20, 0x20);
}

void delay(uint32_t milliseconds) {
    uint32_t end = tick + milliseconds;
    while (tick < end) {
        asm volatile("hlt");
    }
}

uint8_t vga_color_from_char(char color_code) {
    switch (color_code) {
        case '0': return VGA_COLOR_BLACK;
        case '1': return VGA_COLOR_BLUE;
        case '2': return VGA_COLOR_GREEN;
        case '3': return VGA_COLOR_CYAN;
        case '4': return VGA_COLOR_RED;
        case '5': return VGA_COLOR_MAGENTA;
        case '6': return VGA_COLOR_BROWN;
        case '7': return VGA_COLOR_LIGHT_GREY;
        case '8': return VGA_COLOR_DARK_GREY;
        case '9': return VGA_COLOR_LIGHT_BLUE;
        case 'a': return VGA_COLOR_LIGHT_GREEN;
        case 'b': return VGA_COLOR_LIGHT_CYAN;
        case 'c': return VGA_COLOR_LIGHT_RED;
        case 'd': return VGA_COLOR_LIGHT_MAGENTA;
        case 'e': return VGA_COLOR_LIGHT_BROWN;
        case 'f': return VGA_COLOR_WHITE;
        default: return VGA_COLOR_LIGHT_GREY;
    }
}

void terminal_writestring_with_colors(const char* data) {
    while (*data) {
        if (*data == '&' && (*(data + 1) >= '0' && *(data + 1) <= '9' || *(data + 1) >= 'a' && *(data + 1) <= 'f')) {
            terminal_setcolor(vga_entry_color(vga_color_from_char(*(data + 1)), VGA_COLOR_BLACK));
            data += 2;
        } else {
            terminal_putchar(*data);
            data++;
        }
    }
}

#define PIT_FREQUENCY 1193182
#define PIT_CHANNEL_0 0x40
#define PIT_COMMAND 0x43

void pit_set_frequency(uint32_t frequency) {
    uint32_t divisor = PIT_FREQUENCY / frequency;
    outb(PIT_COMMAND, 0x36);
    outb(PIT_CHANNEL_0, divisor & 0xFF);
    outb(PIT_CHANNEL_0, (divisor >> 8) & 0xFF);
}



void kernel_main(void) {
    terminal_initialize();
    gdt_install();
    idt_install();
    pic_remap(0x20, 0x28);
    pit_set_frequency(1000);
    idt_set_gate(32, (uint32_t)pit_handler_asm, 0x08, 0x8E);
    asm volatile("sti");
    terminal_writestring_with_colors("Hello, kernel World!\n&5Magenta String Example\nThis is a &4red string&7.\n");
    for (int i = 0; i < 30; i++) {
        terminal_writestring_with_colors("Line ");
        terminal_putchar('0' + (i / 10));
        terminal_putchar('0' + (i % 10));
        terminal_writestring_with_colors("\n");
        delay(1000);
    }
}