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

//TODO - scheduling.

namespace hilbert::kernel::application {

  process::process() : 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 process::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 process::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 process::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 process::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;
  }

  utility::id_allocator<process *> *processes;
  utility::queue<thread *> *paused_threads;
  utility::queue<thread *> *threads_waiting_for_input;
  thread *running_thread;
  utility::list<socket_listener *> *all_socket_listeners;

  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::success) \
      return stream_result::io_error; \
  }

#define TRY_MAR(expr) \
  { \
    storage::fs_result _result = expr; \
    if (_result != storage::fs_result::success) { \
      delete process_out; \
      return stream_result::io_error; \
    } \
  }

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

  storage::fs_result map_and_read(
    const vfile::vfile &file, process *process, uint64_t vaddr, uint64_t faddr,
    uint64_t len, bool writable, bool executable
  ) {

    uint64_t page_vaddr = vaddr & ~4095;
    int at_start = vaddr & 4095;
    int at_end = 4096 - len - at_start;

    uint64_t page_paddr = paging::take_pram_page();
    process->map_page(page_vaddr, page_paddr, writable, executable, true);
    uint64_t page_kvaddr = paging::find_unmapped_vram_region(1);
    paging::map_kernel_page(page_paddr, page_kvaddr, true, false);

    storage::fs_result result = storage::fs_result::success;

    if (at_start) {
      uint8_t *blank = (uint8_t *)page_kvaddr;
      for (int i = 0; i < at_start; ++i)
        blank[i] = 0;
    }

    if (len != 0)
      result = file.read_file(faddr, len, (void *)(page_kvaddr + at_start));

    if (at_end) {
      uint8_t *blank = (uint8_t *)(page_kvaddr + at_start + len);
      for (int i = 0; i < at_end; ++i)
        blank[i] = 0;
    }

    paging::unmap_kernel_page(page_kvaddr);
    return result;

  }

  stream_result create_application(
    const vfile::vfile &file, process *&process_out, thread *&thread_out
  ) {

    uint8_t magic[16];
    if (file.dir_entry.type != storage::file_type::regular_file)
      return stream_result::not_a_regular_file;
    if (file.dir_entry.length < 64)
      return stream_result::not_an_executable;
    READ(0, 8, magic)
    READ(16, 8, magic + 8)
    for (int i = 0; i < 16; ++i)
      if (magic[i] != correct_magic[i])
        return stream_result::not_an_executable;

    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 stream_result::not_an_executable;

    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 voffset;
      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)

      voffset = vaddr % 4096;
      vaddr -= voffset;

      if (vsize == 0)
        continue;

      if (file.dir_entry.length < foffset + fsize)
        return stream_result::not_an_executable;
      if (fsize > vsize)
        return stream_result::not_an_executable;

      if (vaddr < 0x200000)
        return stream_result::not_an_executable;

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

      if (vaddr + vpages * 4096 > 0x8000000000)
        return stream_result::not_an_executable;

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

    }

    process_out = new process();

    for (unsigned i = 0; i < load_infos.count; ++i) {
      const auto &info = load_infos.buffer[i];

      uint64_t vaddr = info.vaddr + info.voffset;
      uint64_t faddr = info.foffset;
      uint64_t v_remaining = info.vpages * 4096 - info.voffset;
      uint64_t f_remaining = info.fsize;

      if (info.voffset != 0) {
        int to_read = info.fsize < 4096 - info.voffset
                    ? info.fsize : 4096 - info.voffset;
        if (to_read > 0) {
          TRY_MAR(map_and_read(file, process_out, vaddr, faddr, to_read,
            info.writable, info.executable))
          vaddr += to_read;
          faddr += to_read;
          v_remaining -= to_read;
          f_remaining -= to_read;
        }
      }

      while (f_remaining > 0) {
        int to_read = f_remaining < 4096 ? f_remaining : 4096;
        TRY_MAR(map_and_read(file, process_out, vaddr, faddr, to_read,
          info.writable, info.executable))
        vaddr += to_read;
        faddr += to_read;
        v_remaining -= to_read;
        f_remaining -= to_read;
      }

      if (vaddr & 4095) {
        v_remaining -= 4096 - (vaddr & 4095);
        vaddr += 4096 - (vaddr & 4095);
      }

      while (v_remaining > 0) {
        map_and_read(
          file, process_out, vaddr, 0, 0, info.writable, info.executable);
        vaddr += 4096;
        v_remaining -= 4096;
      }

    }

    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);
      process_out->map_page(vaddr, paddr, true, false, true);
    }

    thread_out = new thread();
    process_out->threads.insert_end(thread_out);
    thread_out->the_process = process_out;

    thread_out->state = thread_state::paused;

    thread_out->cpu.rax = 0;
    thread_out->cpu.rbx = 0;
    thread_out->cpu.rcx = 0;
    thread_out->cpu.rdx = 0;
    thread_out->cpu.rdi = 0;
    thread_out->cpu.rsi = 0;
    thread_out->cpu.rbp = 0;
    thread_out->cpu.rsp = 0x1ff000;
    thread_out->cpu.r8 = 0;
    thread_out->cpu.r9 = 0;
    thread_out->cpu.r10 = 0;
    thread_out->cpu.r11 = 0;
    thread_out->cpu.r12 = 0;
    thread_out->cpu.r13 = 0;
    thread_out->cpu.r14 = 0;
    thread_out->cpu.r15 = 0;

    thread_out->cpu.rflags = 0x200;
    thread_out->cpu.rip = entry_point;
    thread_out->cpu.cr3 = process_out->p4_paddr;
    thread_out->cpu.in_syscall = false;

    return stream_result::success;

  }

  extern "C" void init_applications_asm();

  void init_applications() {
    processes = new utility::id_allocator<process *>();
    paused_threads = new utility::queue<thread *>();
    threads_waiting_for_input = new utility::queue<thread *>();
    all_socket_listeners = new utility::list<socket_listener *>();
    init_applications_asm();
  }

  //only called from non-interruptable contexts.
  //cpu argument not on stack.
  extern "C" [[noreturn]] void resume_thread(const cpu_state &cpu);

  extern "C" void *copy_syscall_stack(uint8_t *rsp) {
    uint64_t size = 0xfffffffffffff000 - (uint64_t)rsp;
    uint8_t *buffer = new uint8_t[size];
    for (uint64_t i = 0; i < size; ++i)
      buffer[i] = rsp[i];
    return buffer;
  }

  extern "C" void restore_syscall_stack(const uint8_t *from, uint8_t *rsp) {
    uint64_t size = 0xfffffffffffff000 - (uint64_t)rsp;
    for (uint64_t i = 0; i < size; ++i)
      rsp[i] = from[i];
    delete[] from;
  }

  thread::~thread() {
    for (auto *p = the_process->threads.first; p; p = p->next)
      if (p->value == this) {
        the_process->threads.remove(p);
        break;
      }
    if (the_process->threads.first == 0) {
      the_process->exit_code = exit_code;
      the_process->cleanup();
    }
    if (state != thread_state::running)
      panic(0x9af5e6);
  }

  [[noreturn]] void resume_next() {
    while (paused_threads->count == 0)
      asm volatile ("sti\nhlt\ncli");
    auto *t = paused_threads->take();
    running_thread = t;
    t->state = thread_state::running;
    resume_thread(t->cpu);
  }

  void process::end_process(unsigned exit_code) {
    while (threads.first != 0)
      delete threads.first->value;
    this->exit_code = exit_code;
    cleanup();
  }

  void process::cleanup() {
    //TODO
    panic(0x9af5e6);
  }

  socket_stream::socket_stream(socket *sock, bool are_we_b)
    : sock(sock), are_we_b(are_we_b),
      our_threads_waiting_to_read(are_we_b
        ? sock->process_b_threads_waiting_to_read
        : sock->process_a_threads_waiting_to_read),
      their_threads_waiting_to_read(are_we_b
        ? sock->process_a_threads_waiting_to_read
        : sock->process_b_threads_waiting_to_read),
      them_to_us(are_we_b ? sock->a_to_b : sock->b_to_a),
      us_to_them(are_we_b ? sock->b_to_a : sock->a_to_b),
      them_closed(are_we_b ? sock->a_closed : sock->b_closed),
      us_closed(are_we_b ? sock->b_closed : sock->a_closed) {}

  stream_result socket_stream::seek(seek_origin, int64_t) {
    return stream_result::not_seekable;
  }

  stream_result socket_stream::read(uint64_t count, void *into) {
    uint8_t *buffer = (uint8_t *)into;
    for (uint64_t i = 0; i < count; ++i) {
      while (them_to_us.count == 0) {
        if (them_closed)
          return stream_result::other_end_closed;
        if (!save_thread_state(running_thread->cpu)) {
          running_thread->state = thread_state::waiting;
          our_threads_waiting_to_read.insert(running_thread);
          resume_next();
        }
      }
      buffer[i] = them_to_us.take();
    }
    return stream_result::success;
  }

  stream_result socket_stream::write(uint64_t count, const void *from) {
    if (them_closed)
      return stream_result::other_end_closed;
    const uint8_t *buffer = (const uint8_t *)from;
    for (uint64_t i = 0; i < count; ++i) {
      if (their_threads_waiting_to_read.count > 0) {
        auto *ot = their_threads_waiting_to_read.take();
        ot->state = thread_state::paused;
        paused_threads->insert(ot);
      }
      us_to_them.insert(buffer[i]);
    }
    return stream_result::success;
  }

  stream_result socket_stream::get_length(uint64_t &) {
    return stream_result::not_sized;
  }

  stream_result socket_stream::set_length(uint64_t) {
    return stream_result::not_sized;
  }

  socket_stream::~socket_stream() {
    if (our_threads_waiting_to_read.count > 0)
      panic(0x9af5e6);
    if (them_closed)
      delete sock;
    else {
      us_closed = true;
      while (their_threads_waiting_to_read.count > 0) {
        auto *t = their_threads_waiting_to_read.take();
        t->state = thread_state::paused;
        paused_threads->insert(t);
      }
    }
  }

  vfile_stream::vfile_stream(vfile::vfile &&file)
    : file(utility::move(file)), offset(0) {}

  stream_result vfile_stream::seek(seek_origin origin, int64_t offset) {
    uint64_t start_at = {};
    switch (origin) {
    case seek_origin::beginning:
      start_at = 0;
      break;
    case seek_origin::end:
      start_at = file.dir_entry.length;
      break;
    case seek_origin::current_position:
      start_at = this->offset;
      break;
    }
    if (offset < 0 && (uint64_t)-offset > start_at)
      return stream_result::out_of_bounds;
    if (offset + start_at > file.dir_entry.length)
      return stream_result::out_of_bounds;
    this->offset = start_at + offset;
    return stream_result::success;
  }

  stream_result vfile_stream::read(uint64_t count, void *into) {
    if (offset + count > file.dir_entry.length)
      return stream_result::out_of_bounds;
    if (file.read_file(offset, count, into) != storage::fs_result::success)
      return stream_result::io_error;
    offset += count;
    return stream_result::success;
  }

  stream_result vfile_stream::write(uint64_t count, const void *from) {
    if (offset + count > file.dir_entry.length)
      return stream_result::out_of_bounds;
    (void)from;
    panic(0x9af5e6);
  }

  stream_result vfile_stream::get_length(uint64_t &out) {
    out = file.dir_entry.length;
    return stream_result::success;
  }

  stream_result vfile_stream::set_length(uint64_t to) {
    (void)to;
    panic(0x9af5e6);
  }

}