#include <hilbert/kernel/application.hpp>
#include <hilbert/kernel/paging.hpp>

//TODO - scheduling.

namespace hilbert::kernel::application {

  app_instance::app_instance()
    : state(app_state::paused), framebuffer_vaddr(0) {

    uint64_t p4_vaddr;
    paging::map_new_kernel_page(p4_vaddr, p4_paddr);
    p4 = (uint64_t *)p4_vaddr;

    uint64_t p3_paddr;
    uint64_t p3_vaddr;
    paging::map_new_kernel_page(p3_vaddr, p3_paddr);
    p3 = (uint64_t *)p3_vaddr;

    for (int i = 1; i < 511; ++i)
      p4[i] = 0;
    p4[0] = paging::encode_pte(p3_paddr, true, true, true);
    p4[511] = paging::kernel_p4e;

    for (int i = 0; i < 512; ++i) {
      p3[i] = 0;
      p2s[i] = 0;
      p1s[i] = 0;
      p1es_to_free_on_exit[i] = 0;
    }

  }

  void app_instance::map_page(uint64_t vaddr, uint64_t paddr,
    bool write, bool execute, bool free_pram_on_exit
  ) {

    uint64_t i = ((vaddr / 4096) / 512) / 512;
    uint64_t j = ((vaddr / 4096) / 512) % 512;
    uint64_t k = (vaddr / 4096) % 512;

    if (p2s[i] == 0) {
      uint64_t p2_paddr;
      uint64_t p2_vaddr;
      paging::map_new_kernel_page(p2_vaddr, p2_paddr);
      p3[i] = paging::encode_pte(p2_paddr, true, true, true);
      p2s[i] = (uint64_t *)p2_vaddr;
      p1s[i] = new uint64_t *[512];
      p1es_to_free_on_exit[i] = new bool *[512];
      for (int u = 0; u < 512; ++u) {
        p2s[i][u] = 0;
        p1s[i][u] = 0;
        p1es_to_free_on_exit[i][u] = 0;
      }
    }

    if (p2s[i][j] == 0) {
      uint64_t p1_paddr;
      uint64_t p1_vaddr;
      paging::map_new_kernel_page(p1_vaddr, p1_paddr);
      p2s[i][j] = paging::encode_pte(p1_paddr, true, true, true);
      p1s[i][j] = (uint64_t *)p1_vaddr;
      p1es_to_free_on_exit[i][j] = new bool[512];
      for (int u = 0; u < 512; ++u) {
        p1s[i][j][u] = 0;
        p1es_to_free_on_exit[i][j][u] = false;
      }
    }

    p1s[i][j][k] = paging::encode_pte(paddr, true, write, execute);
    p1es_to_free_on_exit[i][j][k] = free_pram_on_exit;

  }

  bool app_instance::is_page_owned(uint64_t vaddr) {
    uint64_t i = ((vaddr / 4096) / 512) / 512;
    uint64_t j = ((vaddr / 4096) / 512) % 512;
    uint64_t k = (vaddr / 4096) % 512;
    return
      i < 512 && p1s[i] != 0 && p1s[i][j] != 0 &&
      p1s[i][j][k] != 0 && p1es_to_free_on_exit[i][j][k];
  }

  uint64_t app_instance::get_free_vaddr_pages(uint64_t count) {
    uint64_t start = 0x200000 / 4096;
    uint64_t length = 0;
    while (start + length <= 0x8000000000 / 4096) {
      if (length == count)
        return start * 4096;
      int i = ((start + length) / 512) / 512;
      int j = ((start + length) / 512) % 512;
      int k = (start + length) % 512;
      if (p1s[i] == 0 || p1s[i][j] == 0 || p1s[i][j][k] == 0)
        ++length;
      else {
        start += length + 1;
        length = 0;
      }
    }
    //TODO: handle out of memory
    return 0;
  }

  uint64_t app_instance::count_mapped_vram_pages() {
    uint64_t count = 0;
    for (int i = 0; i < 512; ++i)
      if (p1s[i] != 0)
        for (int j = 0; j < 512; ++j)
          if (p1s[i][j] != 0)
            for (int k = 0; k < 512; ++k)
              if (p1s[i][j][k] != 0)
                ++count;
    return count;
  }

  app_instance *running_app;

  static uint8_t correct_magic[16] = {
    0x7f, 0x45, 0x4c, 0x46, 0x02, 0x01, 0x01, 0x00,
    0x02, 0x00, 0x3e, 0x00, 0x01, 0x00, 0x00, 0x00
  };

#define READ(a, b, c) \
  { \
    storage::fs_result _result = file.read_file(a, b, c); \
    if (_result == storage::fs_result::device_error) \
      return create_app_result::device_error; \
    if (_result == storage::fs_result::fs_corrupt) \
      return create_app_result::fs_corrupt; \
  }

  struct load_info {
    uint64_t foffset;
    uint64_t fsize;
    uint64_t vaddr;
    uint64_t vpages;
    bool writable;
    bool executable;
  };

  create_app_result create_app(
    const vfile::vfile &file, app_instance *&out,
    const vfile::vfile &working_dir
  ) {

    uint8_t magic[16];
    if (file.dir_entry.length < 64)
      return create_app_result::app_corrupt;
    READ(0, 8, magic)
    READ(16, 8, magic + 8)
    for (int i = 0; i < 16; ++i)
      if (magic[i] != correct_magic[i])
        return create_app_result::app_corrupt;

    uint64_t entry_point;
    uint64_t phead_start;
    uint16_t phead_entry_size;
    uint16_t phead_entry_count;

    READ(24, 8, &entry_point)
    READ(32, 8, &phead_start)
    READ(54, 2, &phead_entry_size)
    READ(56, 2, &phead_entry_count)

    if (file.dir_entry.length <
        phead_start + phead_entry_size * phead_entry_count)
      return create_app_result::app_corrupt;

    utility::vector<load_info> load_infos;

    for (uint16_t i = 0; i < phead_entry_count; ++i) {

      uint64_t entry_start = phead_start + phead_entry_size * i;

      uint32_t seg_type;
      READ(entry_start, 4, &seg_type)
      if (seg_type != 1)
        continue;

      uint64_t foffset;
      uint64_t vaddr;
      uint64_t fsize;
      uint64_t vsize;
      uint32_t flags;

      READ(entry_start + 8, 8, &foffset)
      READ(entry_start + 16, 8, &vaddr)
      READ(entry_start + 32, 8, &fsize)
      READ(entry_start + 40, 8, &vsize)
      READ(entry_start + 4, 4, &flags)

      if (vaddr & 4095)
        return create_app_result::app_corrupt;
      if (file.dir_entry.length < foffset + fsize)
        return create_app_result::app_corrupt;
      if (fsize > vsize)
        return create_app_result::app_corrupt;

      if (vaddr < 0x200000)
        return create_app_result::app_corrupt;

      uint64_t vpages = (vsize - 1) / 4096 + 1;

      if (vaddr + vpages * 4096 > 0x8000000000)
        return create_app_result::app_corrupt;

      load_info info = {
        .foffset = foffset,
        .fsize = fsize,
        .vaddr = vaddr,
        .vpages = vpages,
        .writable = (flags & 2) == 2,
        .executable = (flags & 1) == 1
      };
      load_infos.add_end(info);

    }

    out = new app_instance();

    for (unsigned i = 0; i < load_infos.count; ++i) {
      const auto &info = load_infos.buffer[i];
      for (uint64_t j = 0; j < info.vpages; ++j) {
        uint64_t paddr = paging::take_pram_page();
        out->map_page(info.vaddr + j * 4096, paddr,
          info.writable, info.executable, true);
        uint64_t kvaddr = paging::find_unmapped_vram_region(1);
        paging::map_kernel_page(paddr, kvaddr, true, false);
        storage::fs_result result = storage::fs_result::success;
        if (info.fsize > j * 4096) {
          if (info.fsize >= j * 4096 + 4096)
            result = file.read_file(
              info.foffset + j * 4096, 4096, (void *)kvaddr);
          else {
            int to_read = info.fsize - j * 4096;
            result = file.read_file(
              info.foffset + j * 4096, to_read, (void *)kvaddr);
            uint8_t *blank = (uint8_t *)(kvaddr + to_read);
            for (int i = 0; i < 4096 - to_read; ++i)
              blank[i] = 0;
          }
        }
        else {
          uint8_t *blank = (uint8_t *)kvaddr;
          for (int i = 0; i < 4096; ++i)
            blank[i] = 0;
        }
        paging::unmap_kernel_page(kvaddr);
        if (result == storage::fs_result::device_error) {
          delete out;
          return create_app_result::device_error;
        }
        if (result == storage::fs_result::fs_corrupt) {
          delete out;
          return create_app_result::fs_corrupt;
        }
      }
    }

    for (uint64_t vaddr = 0x1000; vaddr < 0x1ff000; vaddr += 4096) {
      uint64_t paddr = paging::take_pram_page();
      uint64_t kvaddr = paging::find_unmapped_vram_region(1);
      paging::map_kernel_page(paddr, kvaddr, true, false);
      uint8_t *p = (uint8_t *)kvaddr;
      for (int i = 0; i < 4096; ++i)
        p[i] = 0;
      paging::unmap_kernel_page(kvaddr);
      out->map_page(vaddr, paddr, true, false, true);
    }

    out->saved_regs.rsp = 0x1ff000;
    out->saved_regs.rip = entry_point;

    out->working_dir = working_dir;

    return create_app_result::success;

  }

}