#include "multiboot.h"
#include "interrupts.h"
#include "gdt.h"
#include "common.h"
#include "scheduler.h"
#include "vfs.h"
#include "elf.h"
#include "gdt.h"
#include "kbd.h"

#include <lib/print.h>
#include <lib/stream.h>
#include <lib/memory.h>
#include <lib/console.h>

#include <memory/mm.h>

#include <drivers/iso/iso.h>
#include <drivers/ide/ide.h>
#include <drivers/pic/pic.h>
#include <drivers/ps2/ps2.h>
#include <drivers/pit/pit.h>
#include <drivers/vga/vga.h>
#include <drivers/cpio/cpio.h>
#include <drivers/devfs/devfs.h>
#include <drivers/serial/serial.h>

#include <stddef.h>
#include <stdint.h>

extern char __kernel_start;
extern char __kernel_end;

#define SERIAL_DBG

extern uint32_t page_directory[1024];
void* kernel_cr3 = page_directory;

void parse_multiboot_modules(uint32_t mbi_vaddr) {
  // find where grub put the boot module containing initrd
  multiboot_info* mbi = (multiboot_info*)mbi_vaddr;

  uint32_t mod_count = mbi->mods_count;
  multiboot_module_t* mod = (multiboot_module_t*)((uint32_t)mbi->mods_addr + 0xC0000000);

  for (uint32_t i = 0; i < mod_count; i++) {
    uint32_t start = mod[i].mod_start;
    uint32_t end = mod[i].mod_end;
    uint32_t cmdline = mod[i].cmdline;

    multiboot_module_info_table[i].mod_start = start + 0xC0000000;
    multiboot_module_info_table[i].mod_end = end + 0xC0000000;
    multiboot_module_info_table[i].cmdline = cmdline + 0xC0000000;
  }
}

vfs_node_t* load_initrd() {
  multiboot_module_t* initrd_module = &multiboot_module_info_table[0];
  kprintf("initrd found @ [v] 0x%x\n", initrd_module->mod_start);
  return cpio_read_fs(initrd_module->mod_start);
}

void load_root_program(const char* path) {
  lock_scheduler();
  uint32_t kernel_cr3 = fetch_cr3();
  void* cr3 = create_new_pd();
  
  load_pd(cr3);
  flush_tlb();

  uint32_t eip = load_elf_program(path);
  if (eip == 0) {
    kprintf("Failed to load elf program file\n");
    return;
  }

  uint32_t stack_sz = PAGE_SIZE * 16;
  uint32_t stack_top = 0xBFFF0000;
  uint32_t stack_bottom = stack_top - stack_sz;
  for (uint32_t vaddr = stack_bottom; vaddr < stack_top; vaddr += PAGE_SIZE) {
      void* phys_frame = p_alloc_frame();
      map_page(vaddr, (uint32_t)phys_frame, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
  }

  uint32_t* u_esp = (uint32_t*)stack_top;
  *(--u_esp) = 0;          // envp = NULL
  *(--u_esp) = 0;          // argv = NULL
  *(--u_esp) = 0;          // argc = 0

  process_t* process = kalloc(sizeof(process_t));
  init_task(process, eip, (uint32_t)cr3, 1, (uint32_t)u_esp);
  add_task(process);

  unlock_scheduler();
  
  load_pd((void*)kernel_cr3);
  flush_tlb();  
}

void kmain(multiboot_info* mbi)  {
  disable_interrupts();

  init_serial();

  // need mmap and fb info
  if (!(mbi->flags & (1 << 6)) || !(mbi->flags & (1 << 12))) {
    return;
  } 

  parse_multiboot_modules((uint32_t)mbi);

  uint32_t vga_fb_address = mbi->framebuffer_addr;
  uint32_t vga_fb_width = mbi->framebuffer_width;
  uint32_t vga_fb_height = mbi->framebuffer_height;
  uint32_t vga_fb_pitch = mbi->framebuffer_pitch;
  uint32_t vga_fb_bpp = mbi->framebuffer_bpp;

  init_gdt();
  init_idt();
  init_pic();
  init_pit();

  uint32_t mmap_virtual_addr = mbi->mmap_addr + 0xC0000000;
  init_pmm((void*)mmap_virtual_addr, mbi->mmap_length);
  init_page_tables();

  disable_interrupts();
  init_scheduler();
  enable_interrupts();

  vfs_node_t* vfs_root = vfs_create_root();
  vfs_node_t* initrd_root = load_initrd();
  
  if (!initrd_root) {
    kprintf("failed to load initrd\n");
  }

  vfs_mount(vfs_root, initrd_root);
  kprintf("initrd read, mounted @ /\n");

  vfs_node_t* devfs = create_devfs();
  if (!devfs) {
    kprintf("failed to create devfs\n");
    return;
  }

  vfs_node_t* kbd_dev = devfs->create(devfs, "kbd");
  kbd_dev->read = kbd_read;

  vfs_node_t* vga_dev = devfs->create(devfs, "vga");
  vga_init(vga_dev,
    vga_fb_address,
    vga_fb_width,
    vga_fb_height,
    vga_fb_pitch,
    vga_fb_bpp);

  kprintf("vga initialized\n");

  vfs_node_t* dev = vfs_find("/dev");
  if (!dev) {
    kprintf("failed to find dev mnt point\n");
    return;
  }

  dev->mnt = devfs;
  kprintf("mounted devfs\n");
  
  int font_loaded = set_console_font("/usr/fonts/kernel_font.psf");
  if (!font_loaded) {
    kprintf("failed to load console font\n");
  }
  console_init();

  kprintf("kernel console initialized\n");
  kprintf("basic initializations complete. check serial for details\n");
  kprintf("starting rockos\n");

  load_root_program("/usr/bin/rockterm");
}