#include <mercury/kernel/storage/bd/memory.hpp>
#include <mercury/kernel/storage/fs/tarfs.hpp>
#include <mercury/kernel/framebuffer.hpp>
#include <mercury/kernel/terminal.hpp>
#include <mercury/kernel/limine.hpp>
#include <mercury/kernel/paging.hpp>
#include <mercury/kernel/vfile.hpp>

using namespace mercury::kernel;

LIMINE_BASE_REVISION(1)

static volatile limine_memmap_request memmap_request {
  .id = LIMINE_MEMMAP_REQUEST,
  .revision = 0,
  .response = 0
};

static volatile limine_kernel_address_request kernel_address_request {
  .id = LIMINE_KERNEL_ADDRESS_REQUEST,
  .revision = 0,
  .response = 0
};

static volatile limine_framebuffer_request framebuffer_request {
  .id = LIMINE_FRAMEBUFFER_REQUEST,
  .revision = 0,
  .response = 0
};

static volatile limine_hhdm_request hhdm_request {
  .id = LIMINE_HHDM_REQUEST,
  .revision = 0,
  .response = 0
};

static limine_internal_module initfs_module = {
  .path = "initfs.tgz",
  .cmdline = "initfs",
  .flags = LIMINE_INTERNAL_MODULE_REQUIRED | LIMINE_INTERNAL_MODULE_COMPRESSED
};

static limine_internal_module termfont_module = {
  .path = "termfont.psf",
  .cmdline = "termfont",
  .flags = LIMINE_INTERNAL_MODULE_REQUIRED
};

static limine_internal_module *internal_modules[] = {
  &initfs_module, &termfont_module
};

static volatile limine_module_request module_request = {
  .id = LIMINE_MODULE_REQUEST,
  .revision = 2,
  .response = 0,
  .internal_module_count = 2,
  .internal_modules = internal_modules
};

bool try_map_module_by_cmdline(
  const char *cmdline, void *&vaddr_out, uint64_t &len_out
) {
  auto response = module_request.response;
  for (uint64_t i = 0; i < response->module_count; ++i) {
    limine_file *file = response->modules[i];
    for (uint64_t j = 0; cmdline[j] == file->cmdline[j]; ++j)
      if (!cmdline[j]) {

        //module start is guaranteed to be page-aligned, end is not.
        uint64_t start_paddr =
          (uint64_t)file->address - hhdm_request.response->offset;
        uint64_t end_paddr =
          ((start_paddr + file->size - 1) / 4096 + 1) * 4096;

        uint64_t start_vaddr =
          paging::find_unmapped_vram_region((end_paddr - start_paddr) / 4096);
        for (uint64_t i = 0; i < end_paddr - start_paddr; i += 4096)
          paging::map_kernel_page(
            start_paddr + i, start_vaddr + i, true, false);

        vaddr_out = (void *)start_vaddr;
        len_out = file->size;

        return true;
      }
  }
  return false;
}

//defined in linker script, page-aligned:
extern uint8_t __kernel_start;
extern uint8_t __kernel_end;
extern uint8_t __kernel_rx_start;
extern uint8_t __kernel_rx_end;
extern uint8_t __kernel_ro_start;
extern uint8_t __kernel_ro_end;
extern uint8_t __kernel_rw_start;
extern uint8_t __kernel_rw_end;

uint8_t *initfs;
uint64_t initfs_len;

[[noreturn]] static void with_kernel_p4();

extern "C" [[noreturn]] void entry() {

  //TODO?: maybe we should check if the limine requests were
  //       fulfilled and display some error message if not

  //set up the physical memory usage bitmap:

  paging::mark_all_pram_used();
  auto memmap = memmap_request.response;
  for (uint64_t i = 0; i < memmap->entry_count; ++i) {
    //we don't start allocating physical pages until after we are done using
    //limine structures (specifically at the call to paging::map_kernel_stack),
    //so we consider bootloader reclaimable to be free. usable and bootloader
    //reclaimable are guaranteed by limine spec to be page-aligned.
    auto entry = memmap->entries[i];
    if (entry->type == LIMINE_MEMMAP_USABLE ||
        entry->type == LIMINE_MEMMAP_BOOTLOADER_RECLAIMABLE)
      paging::mark_pram_region_free(entry->base, entry->base + entry->length);
  }

  //set up page mappings:

  auto kernel_address = kernel_address_request.response;
  uint64_t kernel_offset = kernel_address->virtual_base
                         - kernel_address->physical_base;

  uint64_t hhdm = hhdm_request.response->offset;

  //framebuffer might not be page-aligned
  auto framebuffer = framebuffer_request.response->framebuffers[0];
  uint64_t fb_start = ((uint64_t)framebuffer->address / 4096) * 4096 - hhdm;
  uint64_t fb_end = (uint64_t)framebuffer->address
                  + framebuffer->pitch * framebuffer->height - hhdm;
  fb_end = ((fb_end - 1) / 4096 + 1) * 4096;

  paging::init_kernel_page_tables(kernel_offset);

  //kernel image rx
  for (uint64_t vaddr = (uint64_t)&__kernel_rx_start;
       vaddr < (uint64_t)&__kernel_rx_end; vaddr += 4096)
    paging::map_kernel_page(vaddr - kernel_offset, vaddr, false, true);

  //kernel image ro
  for (uint64_t vaddr = (uint64_t)&__kernel_ro_start;
       vaddr < (uint64_t)&__kernel_ro_end; vaddr += 4096)
    paging::map_kernel_page(vaddr - kernel_offset, vaddr, false, false);

  //kernel image rw
  for (uint64_t vaddr = (uint64_t)&__kernel_rw_start;
       vaddr < (uint64_t)&__kernel_rw_end; vaddr += 4096)
    paging::map_kernel_page(vaddr - kernel_offset, vaddr, true, false);

  //framebuffer
  uint64_t fb_vaddr =
    paging::find_unmapped_vram_region((fb_end - fb_start) / 4096);
  for (uint64_t i = 0; i < fb_end - fb_start; i += 4096)
    paging::map_kernel_page(fb_start + i, fb_vaddr + i, true, false);

  //initfs and termfont - these are required modules
  //so there is no worry about them not being present.
  try_map_module_by_cmdline("initfs", (void *&)initfs, initfs_len);
  try_map_module_by_cmdline(
    "termfont", (void *&)terminal::termfont, terminal::termfont_len);

  //set up framebuffer and terminal:
  //TODO: assumes framebuffer is 32-bpp rgb

  framebuffer::init_framebuffer(
    fb_vaddr, framebuffer->width, framebuffer->height, framebuffer->pitch);

  //switch to kernel p4

  paging::map_kernel_stack();
  switch_to_kernel_p4(&with_kernel_p4);

}

[[noreturn]] static void print_and_halt(const char *msg) {
  terminal::put_string_sz(msg);
  while (1)
    asm ("hlt");
}

static void draw_indent(unsigned indent) {
  for (unsigned i = 0; i < indent; ++i)
    terminal::put_char(' ');
}

static void dir_tree(
  const vfile::vfile &root, const vfile::vfile &f, unsigned indent
) {

  draw_indent(indent);
  terminal::put_string(f.dir_entry.name);
  terminal::put_char('\n');

  std::optional<vfile::vfile> followed;
  utility::vector<vfile::vfile> children;

  switch (f.dir_entry.type) {

    case storage::file_type::regular_file:
      draw_indent(indent + 2);
      terminal::put_string_sz("type: regular file\n");
      draw_indent(indent + 2);
      terminal::put_string_sz("length: ");
      terminal::put_int_decimal(f.dir_entry.length);
      terminal::put_string_sz(" bytes\n");
      break;

    case storage::file_type::symlink:
      draw_indent(indent + 2);
      terminal::put_string_sz("type: symlink\n");
      draw_indent(indent + 2);
      terminal::put_string_sz("target: ");
      terminal::put_string(f.dir_entry.target);
      terminal::put_char('\n');
      if (f.follow_symlinks(root, followed) != storage::fs_result::success)
        print_and_halt("failed to follow symlink.");
      if (!followed) {
        draw_indent(indent + 4);
        terminal::put_string_sz("[broken]\n");
      }
      dir_tree(root, *followed, indent + 4);
      break;

    case storage::file_type::directory:
      if (f.get_children(children) != storage::fs_result::success)
        print_and_halt("failed to get children.");
      draw_indent(indent + 2);
      terminal::put_string_sz("type: directory\n");
      draw_indent(indent + 2);
      terminal::put_string_sz("children:\n");
      if (children.count == 0) {
        draw_indent(indent + 4);
        terminal::put_string_sz("[empty]\n");
      }
      else
        for (unsigned i = 0; i < children.count; ++i)
          dir_tree(root, children.buffer[i], indent + 4);
      break;

  }

}

[[noreturn]] static void with_kernel_p4() {

  terminal::init_terminal();

  storage::bd::memory initfs_bd(initfs, initfs_len);
  storage::fs::tarfs_instance initfs_fs(&initfs_bd);
  initfs_bd.mounted_as = &initfs_fs;

  terminal::put_string_sz("kernel initialization complete.\n");

  int used_vram_kib = paging::get_used_vram_page_count() * 4;
  int free_pram_kib = paging::get_free_pram_page_count() * 4;

  terminal::put_int_decimal(used_vram_kib);
  terminal::put_string_sz(" kiB kernel memory mapped.\n");
  terminal::put_int_decimal(free_pram_kib);
  terminal::put_string_sz(" kiB physical memory free.\n");

  vfile::vfile root;
  root.bd = &initfs_bd;
  root.dir_entry.type = storage::file_type::directory;
  root.path.absolute = true;

  if (initfs_fs.get_root_node(root.dir_entry.node) !=
      storage::fs_result::success)
    print_and_halt("failed to get root node.");

  terminal::put_string_sz("initfs");
  dir_tree(root, root, 0);
  print_and_halt("halting.");

}