#include "multiboot.h"
#include <lib/print.h>
#include <lib/stream.h>

#include <idt.h>

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

#include "memory/mm.h"

#include "elf.h"

extern void init_idt();
extern void init_gdt();
extern void init_pic();
extern void init_pit();
extern int init_ps2();
extern void enable_interrupts();

void* page = NULL;

extern char __kernel_start;
extern char __kernel_end;

void* user_page_directories[10];

void load_origin_program(multiboot_info* mbi) {
  user_page_directories[0] = create_user_pd();
  kprintf("Created origin process page directory @ [phys] 0x%x\n", user_page_directories[0]);
  extern void load_pd(void* pd);
  extern void flush_tlb();
  load_pd(user_page_directories[0]);
  flush_tlb();
  kprintf("switched to user address space\n");

  uint32_t mods_virtual_addr = mbi->mods_addr + 0xC0000000;
  kprintf("parsing mutliboot modules @ 0x%x\n", mods_virtual_addr);

  if (mbi->mods_count == 0) {
    return;
  }

  multiboot_module_entry* modules = (multiboot_module_entry*)mods_virtual_addr;

  uint32_t elf_vbase = 0xE0000000;
  uint32_t elf_pbase = modules[0].mod_start;
  uint32_t elf_limit = modules[0].mod_end - modules[0].mod_start;
  uint32_t elf_pages_needed = (elf_limit / 4096) + 1;

  for (uint32_t i = 0; i < elf_pages_needed; i++) {
    const uint32_t vaddr = elf_vbase + (i * 4096);
    const uint32_t paddr = elf_pbase + (i * 4096);
    map_page(vaddr, paddr, PAGE_PRESENT | PAGE_WRITABLE);
  }

  elf_header_t* elf = (elf_header_t*)elf_vbase;
  elf_program_header_t* ph = (elf_program_header_t*)(elf_vbase + elf->e_phoff);

  kprintf("elf header @ 0x%x\n", elf);
  kprintf("program headers @ 0x%x\n", ph);

  // load program sections into physical memory and map the pages
  for (int i = 0; i < elf->e_phnum; i++) {
    if (ph[i].p_type == PT_LOAD) {
      uint32_t v_start = ph[i].p_vaddr & 0xFFFFF000;
      uint32_t v_end = (ph[i].p_vaddr + ph[i].p_memsz + 4095) & 0xFFFFF000;

      for (uint32_t page = v_start; page < v_end; page += 4096) {
        uint32_t frame_paddr = (uint32_t)p_alloc_frame();
        map_page(page, frame_paddr, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
      }

      for (uint32_t j = 0; j < ph[i].p_memsz; j++) {
        ((char*)ph[i].p_vaddr)[j] = 0;
      }

      for (uint32_t j = 0; j < ph[i].p_filesz; j++) {
        ((char*)ph[i].p_vaddr)[j] = ((char*)(elf_vbase + ph[i].p_offset))[j];
      }
    }
  }

  // allocate the user stack
  uint32_t stack_paddr = (uint32_t)p_alloc_frame();
  map_page(0xBFFFF000, stack_paddr, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);

  kprintf("program loaded into memory @ [virt] 0x%x\n", elf->e_entry);

  extern void liftoff(uint32_t entry, uint32_t stack_top);
  kprintf("launching origin process\n");
  uint32_t stack_top =  0xBFFFF000 + 4095;
  liftoff(elf->e_entry, stack_top);
}

void kmain(uint32_t magic, multiboot_info* mbi)  {
  extern print_stream_t vga_print_stream;
  extern print_stream_t* kout;
  kout = &vga_print_stream;
  kout->trunc();

  kprintf("kernel start: 0x%x\n", &__kernel_start);
  kprintf("kernel end: 0x%x\n", &__kernel_end);
  uint32_t size = (&__kernel_end - &__kernel_start) / 1024;
  kprintf("kernel size: %dKB\n", size);

  if (!(mbi->flags & (1 << 6))) {
    return;
  } 

  kprintf("multiboot virtual address: 0x%x\n", mbi);

  kprintf("parsing multiboot memory map\n");
  uint32_t mmap_virtual_addr = mbi->mmap_addr + 0xC0000000;
  kprintf("mmap virtual address: 0x%x\n", mmap_virtual_addr);
  init_pmm((void*)mmap_virtual_addr, mbi->mmap_length);
  uint32_t mem_size = total_free_memory();
  kprintf("Total available memory: %d MB\n", mem_size / 1024 / 1024);

  kprintf("RockOS booting...\n");
  init_gdt();
  extern uint64_t gdt[5];
  kprintf("GDT set @ 0x%x\n", &gdt);

  init_idt();
  extern idt_entry idt[256];
  kprintf("IDT set @ 0x%x\n", &idt);

  init_pic();
  kprintf("PIC initialized\n");

  init_pit();
  kprintf("PIT initialized\n"); 

  // int result = init_ps2();
  // if (result == -1) { 
  //   kprintf("PS2 init failed\n");
  //   return;
  // }
  // kprintf("PS2 initialized\n");  

  init_page_tables();
  kprintf("Page tables initialized\n");

  kprintf("scanning for drives\n");
  extern void scan_drives();
  scan_drives();

  while(1);
}