/* BIOS-OS-APP note - a BIOS, an OS and an app to take a note */
typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;

#define GDT_CODE 0xB
#define GDT_DATA 0x2
#define GDT_TSS 0x9
#define GDT_CALL_GATE 0xC
#define FLAG_FLASH_ADDR 0xff800000
#define FLAG_ADDRESS 0x1400000
#define STACK_START 0x1000
#define USER_STACK 0x5000
#define GDT_BASE 0x1000000 
#define TSS_BASE 0x1000100
#define KERNEL_CODE_BASE 0x0
#define KERNEL_CODE_LIMIT 0x200000
#define KERNEL_DATA_BASE 0x1000000
#define KERNEL_DATA_LIMIT 0x200000
#define USER_CODE_BASE 0x200000
#define USER_CODE_LIMIT 0x100000
#define USER_DATA_BASE 0x1100000
#define USER_DATA_LIMIT 0x100000
#define USER_RODATA (USER_DATA_BASE + 0x1000)
#define KERNEL_CS 0x08
#define KERNEL_DS 0x10
#define USER_CS 0x18
#define USER_DS 0x20
#define TSS_SEG 0x28
#define READ_GATE 0x30
#define WRITE_GATE 0x38
#define EXIT_GATE 0x40
#define UART_BASE 0x3F8
#define UART_THR (UART_BASE + 0)
#define UART_RBR (UART_BASE + 0)
#define UART_LCR (UART_BASE + 3)
#define UART_DLL (UART_BASE + 0)
#define UART_DLH (UART_BASE + 1)
#define UART_LSR (UART_BASE + 5)
#define UART_FCR (UART_BASE + 2)
#define BOOT_CS 0x08
#define BOOT_DS 0x10
#define KERNEL_LOAD_ADDR 0x10000

#define SEGMENT_INIT(base, limit, _type, _s, _dpl)                             \
  {                                                                            \
    .limit_low = (limit) & 0xFFFF, .base_low = (base) & 0xFFFF,                \
    .base_mid = ((base) >> 16) & 0xFF, .type = (_type), .s = (_s),             \
    .dpl = (_dpl), .p = 1, .limit_high = ((limit) >> 16) & 0x0F, .avl = 0,     \
    .rsvd = 0, .db = 1, .g = 1, .base_high = ((base) >> 24) & 0xFF             \
  }

#define IO_WRITE(byte, port)                                                   \
  asm volatile("outb %b0, %1" : : "a"(byte), "d"(port))

#define IO_READ(port, byte)                                                    \
  asm volatile("inb %1, %0" : "=a"(byte) : "d"(port))

#define IO_WRITE16(port, word)                                                 \
  asm volatile("outw %w0, %w1" : : "a"(word), "d"(port))

#define APP_ADDR_TO_KERNEL_ADDR(X) ((u8*)X + 0x100000)

extern u8 _kernel_start[];
extern u8 _kernel_end[];
extern u8 _userspace_start[];
extern u8 _userspace_end[];
extern u8 _userspace_data_start[];
extern u8 _userspace_data_end[];

struct gdt_seg {
    u16 limit_low;
    u16 base_low;
    u8 base_mid;
    u8 type : 4;
    u8 s : 1;
    u8 dpl : 2;
    u8 p : 1;
    u8 limit_high : 4;
    u8 avl : 1;
    u8 rsvd : 1;
    u8 db : 1;
    u8 g : 1;
    u8 base_high;
} __attribute__((packed));

struct gdt_ptr {
    u16 limit;
    u32 base;
} __attribute__((packed));

struct tss_desc {
    u32 r1;
    u32 esp0;
    u32 ss0;
    u32 r2[2];
    u32 esp1;
    u32 ss1;
    u32 r3[2];
    u32 esp2;
    u32 ss2;
    u32 r4[2];
    u8 unused[76];
} __attribute__((packed));

struct callgate_desc {
    u16 offset_low;
    u16 selector;
    u8 param_count;
    u8 type : 4;
    u8 zero : 1;
    u8 dpl : 2;
    u8 p : 1;
    u16 offset_high;
} __attribute__((packed));


__attribute__((section(".bios.gdt"))) struct gdt_seg bios_gdt[] = {
    SEGMENT_INIT(0x0, 0x0, 0x0, 0, 0),
    SEGMENT_INIT(0x0, 0xFFFFFFFF, 0xB, 1, 0),
    SEGMENT_INIT(0x0, 0xFFFFFFFF, 0x2, 1, 0),
};

__attribute__((section(".bios.gdt_ptr")))
struct gdt_ptr gdtr = {
   .limit = sizeof(bios_gdt) - 1,
   .base = 0xffffff00,
};

__attribute__((section(".bios.text")))
void bios_uart_init(void) {
    /* setup the UART0 with baud 115200 8N1 */
    IO_WRITE(0x80, UART_LCR);
    IO_WRITE(0x01, UART_DLL);
    IO_WRITE(0x00, UART_DLH);
    IO_WRITE(0x03, UART_LCR);
    IO_WRITE(0xC7, UART_FCR);
}

void kernel_run(void);

__attribute__((section(".bios.text")))
void bios_start(void) {
    u8 *os, *os_origin;
    u32 os_size;
    u32 i;

    bios_uart_init();
    os = (u8*)KERNEL_LOAD_ADDR;
    os_origin = _kernel_start;
    os_size = _kernel_end - _kernel_start;
    for (i = 0; i < os_size; i++) {
        os[i] = os_origin[i];
    }
    kernel_run();
}

__attribute__((section(".bios.jmp_pm"), naked))
void bios_enable_protected(void) {
  asm volatile(
      ".code16\n"
      "cli\n"
      "mov %[GDT], %%bx\n"
      "lgdtl %%cs:(%%bx)\n"
      "mov %%cr0, %%eax\n"
      "or $0x01, %%eax\n"
      "mov %%eax, %%cr0\n"
      "ljmpl %[CODE_SEL],$1f\n"
      "1:\n"
      ".code32\n"
      "mov %[DATA_SEL], %%ax\n"
      "movw %%ax, %%ds\n"
      "movw %%ax, %%ss\n"
      "mov %[STACK], %%esp\n"
      "jmp bios_start\n"
        :
        : [GDT] "i"(&gdtr), [STACK] "i"(STACK_START), [DATA_SEL] "i"(BOOT_DS), [CODE_SEL] "i"(BOOT_CS)
        : "eax");
}

/* where everything starts */
__attribute__((section(".bios.reset_vector"), naked))
void reset_vector(void) {
  asm volatile(".code16\n"
               "jmp bios_enable_protected\n"
               ".code32\n");
}

__attribute__((section(".kernel.text")))
void kernel_uart_tx(u8 byte) {
    u8 status;
    do {
        IO_READ(UART_LSR, status);
    } while ((status & 0x20) == 0);
    IO_WRITE(byte, UART_THR);
}

__attribute__((section(".kernel.text")))
u8 kernel_uart_rx(void) {
    u8 status;
    do {
        IO_READ(UART_LSR, status);
    } while ((status & 0x01) == 0);
    IO_READ(UART_RBR, status);
    return status;
}

__attribute__((section(".kernel.text")))
void kernel_uart_write(const char *buf, u32 len) {
    u8 *ubuf;
    ubuf = APP_ADDR_TO_KERNEL_ADDR(buf);
    while (len--) {
        kernel_uart_tx(*ubuf++);
    }
}

__attribute__((section(".kernel.text")))
int kernel_uart_read(char *buf, u32 buf_size) {
    u8 *ubuf;
    char c;
    int i = 0;
    ubuf = APP_ADDR_TO_KERNEL_ADDR(buf);
    while (i < buf_size) {
        c = kernel_uart_rx();
        if (c == '\n' || c == '\r')
            break;
        ubuf[i] = c;
        i++;
    }
    return i;
}

__attribute__((section(".kernel.text")))
void kernel_write_gdt_segment(struct gdt_seg *s, u32 base, u32 limit, u8 type, u8 dpl) {
    limit = limit >> 12;
    s->limit_low = limit & 0xFFFF;
    s->limit_high = (limit >> 16) & 0x0F;
    s->base_low = base & 0xFFFF;
    s->base_mid = (base >> 16) & 0xFF;
    s->base_high = (base >> 24) & 0xFF;
    s->type = type;
    s->dpl = dpl;
    s->p = s->db = s->g = 1;
    s->s = (type == GDT_TSS) ? 0 : 1;
    s->avl = s->rsvd = 0;
}

__attribute__((section(".kernel.text")))
void kernel_memset(void *dest, u8 byte, int len) {
  u8 *p = (u8 *)dest;
  while (len--) {
    *p++ = byte;
  }
}

__attribute__((section(".kernel.text"), naked))
void sys_read() {
  asm volatile("pushl %%ebx\n"
               "pushl %%eax\n"
               "mov %[KDATA_SEL], %%ax\n"
               "mov %%ax, %%ds\n"
               "mov %%ax, %%es\n"
               "call kernel_uart_read\n"
               "mov %%eax, %%ebx\n"
               "mov %[UDATA_SEL], %%ax\n"
               "mov %%ax, %%ds\n"
               "mov %%ax, %%es\n"
               "mov %%ebx, %%eax\n"
               "add $8, %%esp\n"
               "retfl\n"
               :
               : [KDATA_SEL] "i" (KERNEL_DS), [UDATA_SEL] "i" (USER_DS)
               : "eax", "ebx");
}

__attribute__((section(".kernel.text"), naked))
void sys_write() {
  asm volatile("pushl %%ebx\n"
               "pushl %%eax\n"
               "mov %[KDATA_SEL], %%ax\n"
               "mov %%ax, %%ds\n"
               "mov %%ax, %%es\n"
               "call kernel_uart_write\n"
               "mov %%eax, %%ebx\n"
               "mov %[UDATA_SEL], %%ax\n"
               "mov %%ax, %%ds\n"
               "mov %%ax, %%es\n"
               "mov %%ebx, %%eax\n"
               "add $8, %%esp\n"
               "retfl\n"
               :
               : [KDATA_SEL] "i" (KERNEL_DS), [UDATA_SEL] "i" (USER_DS)
               : "eax", "ebx");
}

__attribute__((section(".kernel.text"), naked))
void sys_exit() {
    asm volatile("hlt\n");
}

__attribute__((section(".kernel.text")))
void kernel_load_application(void)
{
    u8 *dst, *start;
    u32 size;
    u32 i;

    dst = (u8*)USER_CODE_BASE;
    start = _userspace_start;
    size = _userspace_end - _userspace_start;
    for (i = 0; i < size; i++) {
        dst[i] = start[i];
    }

    dst = (u8*)USER_RODATA;
    start = _userspace_data_start;
    size = _userspace_data_end - _userspace_start;
    for (i = 0; i < size; i++) {
        dst[i] = start[i];
    }
}

/* we need isolation! let's kernel and user application have different memory segments */
__attribute__((section(".kernel.text")))
struct gdt_seg *kernel_setup_segmentation(struct gdt_seg *gdt)
{

    /* null segment */
    kernel_memset(gdt, 0, sizeof(struct gdt_seg));
    gdt++;
    kernel_write_gdt_segment(gdt, KERNEL_CODE_BASE, KERNEL_CODE_LIMIT, GDT_CODE, 0);
    gdt++;
    kernel_write_gdt_segment(gdt,KERNEL_DATA_BASE, KERNEL_DATA_LIMIT, GDT_DATA, 0);
    gdt++;
    kernel_write_gdt_segment(gdt, USER_CODE_BASE, USER_CODE_LIMIT, GDT_CODE, 3);
    gdt++;
    kernel_write_gdt_segment(gdt, USER_DATA_BASE, USER_DATA_LIMIT, GDT_DATA, 3);
    gdt++;

    return gdt;
}

__attribute__((section(".kernel.text")))
struct gdt_seg *kernel_setup_tss(struct gdt_seg *gdt)
{
    struct tss_desc *tss;

    tss = (struct tss_desc *)TSS_BASE;
    kernel_memset((void *)tss, 0, sizeof(struct tss_desc));
    tss->ss0 = KERNEL_DS;
    tss->esp0 = STACK_START;
    tss->ss2 = USER_DS;
    tss->esp2 = USER_STACK;

    kernel_write_gdt_segment(gdt, TSS_BASE, sizeof(struct tss_desc), GDT_TSS, 0);
    gdt++;

    return gdt;
}

__attribute__((section(".kernel.text")))
void kernel_write_callgate_segment(struct gdt_seg *gdt, u32 syscall)
{
    struct callgate_desc *call;
    call = (struct callgate_desc *)gdt;
    call->offset_low = syscall & 0xFFFF;
    call->offset_high = syscall >> 16;
    call->selector = KERNEL_CS;
    call->param_count = 0x0;
    call->type = GDT_CALL_GATE;
    call->dpl = 3;
    call->p = 1;
}

__attribute__((section(".kernel.text")))
struct gdt_seg *kernel_setup_syscalls(struct gdt_seg *gdt)
{
    struct callgate_desc *call;

    /* call gate: read 0x30 */
    kernel_write_callgate_segment(gdt, (u32)sys_read);
    gdt++;

    /* call gate: write 0x38 */
    kernel_write_callgate_segment(gdt, (u32)sys_write);
    gdt++;

    /* call gate: write 0x40 */
    kernel_write_callgate_segment(gdt, (u32)sys_exit);
    gdt++;

    return gdt;
}

__attribute__((section(".kernel.text")))
void kernel_load_gdt_and_run_application(struct gdt_seg *gdt)
{
    struct gdt_ptr gdtr;
    gdtr.limit = (u32)gdt - GDT_BASE - 1;
    gdtr.base = GDT_BASE;
    asm volatile("lgdt %[GDT]\n"
                "mov %[TSS], %%ax\n"
                "ltr %%ax\n"
                "mov %[DATA_SEL], %%ax\n"
                "movw %%ax, %%ds\n"
                "movw %%ax, %%es\n"
                "pushl %[STACK_SEL]\n"
                "pushl %[U_STACK]\n"
                "pushl %[CODE_SEL]\n"
                "pushl $user_app\n"
                "retfl\n"
                :
                : [GDT] "m"(gdtr), [U_STACK] "i"(USER_STACK), [STACK_SEL] "i"(USER_DS | 0x3), [CODE_SEL] "i"(USER_CS | 0x3), [DATA_SEL] "i"(USER_DS), [TSS] "i"(TSS_SEG)
                : "eax");
}

__attribute__((section(".kernel.text")))
void copy_flag(void) {
    char *flag_dst = (char*)(FLAG_ADDRESS);
    char *flag_src = (char*)FLAG_FLASH_ADDR;
    while (*flag_src) {
        *flag_dst++ = *flag_src++;
    }
}
__attribute__((section(".kernel.text")))
void kernel_run(void) {
    struct gdt_seg *gdt;

    copy_flag();
    kernel_load_application();
    gdt = (struct gdt_seg *)GDT_BASE;
    gdt = kernel_setup_segmentation(gdt);
    gdt = kernel_setup_tss(gdt);
    gdt = kernel_setup_syscalls(gdt);
    kernel_load_gdt_and_run_application(gdt);
}

__attribute__((section(".userspace.text")))
int userspace_read(char *buf, int nbytes) {
  int len;
  asm volatile("mov %[_buf], %%eax\n"
               "mov %[_nbyte], %%ebx\n"
               "call %[GATE],$0x0\n"
               "mov %%eax, %[_len]\n"
               : [_len] "=m"(len)
               : [_buf] "m"(buf), [_nbyte] "m"(nbytes), [GATE] "i"(READ_GATE)
               : "eax", "ebx");
  return len;
}
__attribute__((section(".userspace.text")))
void userspace_exit() {
  asm volatile("call %[GATE],$0x0\n"
               :
               :[GATE] "i"(EXIT_GATE)
               :);
}

__attribute__((section(".userspace.text")))
void userspace_write(const char *buf, int nbytes) {
  /* my stack my call convention */
  asm volatile("mov %[_buf], %%eax\n"
               "mov %[_nbyte], %%ebx\n"
               "call %[GATE],$0x0\n"
               :
               : [_buf] "m"(buf), [_nbyte] "m"(nbytes), [GATE] "i"(WRITE_GATE)
               : "eax", "ebx");
}

__attribute__((section(".userspace.text")))
int strlen(const char *str) {
  int len = 0;
  while (*str++) {
    len++;
  }
  return len;
}

__attribute__((section(".userspace.text")))
void puts(const char *str) {
  userspace_write(str, strlen(str));
}

__attribute__((section(".userspace.text")))
void memcpy(char *dst, const char *src, int len) {
  while (len--) {
    *dst++ = *src++;
  }
}

__attribute__((section(".userspace.text")))
void print_menu(void)
{
    puts("1. Write note\n");
    puts("2. Read note\n");
    puts("3. Write final remarks\n");
    puts(">\n");
}

__attribute__((section(".userspace.text")))
void user_app(void) {
    struct my_notebook {
        char data[200];
        char *ptr;
        u8 used;
    } notebook;
    notebook.ptr = notebook.data;
    notebook.used = 0;
    char buf[128];
    u8 bytes_read;
    u8 new_size;

    puts("HoD Notebook Taker!\n");
    while (1) {
        print_menu();
        bytes_read = userspace_read(buf, 2);
        if (bytes_read == 0) {
            continue;
        }
        switch (buf[0]) {
        case '1':
            puts("what to note?\n");
            bytes_read = userspace_read(buf, sizeof(buf));
            new_size = notebook.used + bytes_read;
            if (new_size >= sizeof(notebook.data)) {
                puts("No space left\n");
                break;
            }
            memcpy(notebook.ptr, buf, bytes_read);
            notebook.ptr += bytes_read;
            notebook.used += bytes_read;
            break;
        case '2':
            puts("Read note\n");
            userspace_write(notebook.data, notebook.used);
            puts("\n");
            break;
        case '3':
            puts("Write final remarks\n");
            bytes_read = userspace_read(notebook.ptr, sizeof(notebook.data) - notebook.used);
            notebook.used += bytes_read;
            puts("Here is your note:\n");
            userspace_write(notebook.data, notebook.used);
            puts("\nexit\n");
            userspace_exit();
        }
    }
}