#include "kernel/prelude/kernel_prelude.hats"

#define ATS_DYNLOADFLAG 0

staload "./multiboot.sats"
staload "lib/SATS/init.sats"
staload "lib/DATS/init.dats"
staload UN = "prelude/SATS/unsafe.sats"

local

  // Private functions
  extern castfn ui2sz (n : uint) :<> [n : nat] size_t n
  extern fn _get_memory_mappings_n(boot_info: !boot_info_t):<> [n:nat] size_t n
  extern fn _get_memory_mapping(boot_info: !boot_info_t,n : size_t):<> memory_area_t
  extern fn _print_memory_mappings(boot_info: !boot_info_t):<!wrt> void
  extern fn _get_elf_headers_n(boot_info: !boot_info_t):<> [n:nat] size_t n
  extern fn _get_elf_header(boot_info: !boot_info_t, n: size_t):<> elf64_shdr_t
  extern fn _print_elf_headers(boot_info: !boot_info_t):<!wrt> void
  extern fn _get_kernel_ranges(boot_info: !boot_info_t):<> (Ptr,Ptr)
  extern fn _get_multiboot_ranges(boot_info: !boot_info_t):<> (Ptr0,Ptr0)

in

  local
    // Static variable
    var _v : boot_info_t?
    prval () = opt_none _v
    var v = uninitialized<boot_info_t>(_v)
    // Reference to writer
    val bootinfo_ref = ref_make_viewptr{initializable boot_info_t} (view@v | addr@v)
  in

    implement init (p : Ptr1) = let
      val total_size = g1ofg0($UN.ptr0_get<uint>(p))

      fun loop (boot_info : &boot_info_t? >> boot_info_t, p : ptr):<!wrt,!ntm> void = let
        extern castfn unsafe_init(x : &boot_info_t? >> boot_info_t) :<> void
        val type= $UN.ptr0_get<uint>(p)
        val size = $UN.ptr0_get<uint>(ptr_succ<uint>(p))
        val next = (if (size % 8u != 0) then size + 8u - (size % 8u) else size) : uint
      in 
        case+ type of 
          | 6u => (
                   boot_info.memory_map := $UN.ptr0_get<memory_areas_t>(p); 
                   boot_info.memory_map.entries := ptr_add<uint32>(p,4)
                  )
          | 9u => (
                   boot_info.elf_tag := $UN.ptr0_get<elf_tag_t>(p);
                   boot_info.elf_tag.headers := ptr_add<uint32>(p,5);
                  )
          | _ => ();

        if (type != 0u) then
            loop(boot_info,ptr_add<uint8>(p, next))
        else
          unsafe_init(boot_info) //UNSAFE:
      end

      implement initialize$fwork<boot_info_t>(v) = (
        if (total_size > 0u) then (
          v.address := p;
          v.total_size := total_size;
          $effmask_ntm(loop(v, ptr_add<uint32>(p,2)));
          let prval () = opt_some v in true end)
        else
          let prval () = opt_none v in false end
      )
    in 
      initialize<boot_info_t>(bootinfo_ref)
    end

    implement get_memory_mappings_n() = let
      implement exec$fwork<boot_info_t><size_t>(v) = _get_memory_mappings_n(v)
    in
      $effmask_wrt(exec<boot_info_t><size_t>(bootinfo_ref))
    end

    implement get_memory_mapping(n) = let
      implement exec$fwork<boot_info_t><memory_area_t>(v) = _get_memory_mapping(v,n)
    in
      $effmask_wrt(exec<boot_info_t><memory_area_t>(bootinfo_ref))
    end

    implement print_memory_mappings() = let
      implement exec_void$fwork<boot_info_t>(v) = _print_memory_mappings(v)
    in
      exec_void<boot_info_t>(bootinfo_ref)
    end

    implement get_elf_headers_n() = let
      implement exec$fwork<boot_info_t><size_t>(v) = _get_elf_headers_n(v)
    in
      exec<boot_info_t><size_t>(bootinfo_ref)
    end

    implement get_elf_header(n) = let
      implement exec$fwork<boot_info_t><elf64_shdr_t>(v) = _get_elf_header(v,n)
    in
      exec<boot_info_t><elf64_shdr_t>(bootinfo_ref)
    end

    implement print_elf_headers() = let
      implement exec_void$fwork<boot_info_t>(v) = _print_elf_headers(v)
    in
      exec_void<boot_info_t>(bootinfo_ref)
    end

    implement get_kernel_ranges() = let
      implement exec$fwork<boot_info_t><(Ptr,Ptr)>(v) = _get_kernel_ranges(v)
    in
      exec<boot_info_t><(Ptr,Ptr)>(bootinfo_ref)
    end

    implement get_multiboot_ranges() = let
      implement exec$fwork<boot_info_t><(Ptr0,Ptr0)>(v) = _get_multiboot_ranges(v)
    in
      exec<boot_info_t><(Ptr0,Ptr0)>(bootinfo_ref)
    end

  end

  //------------Memory------------------------

  implement _get_memory_mappings_n(boot_info) = let
    val size = boot_info.memory_map.tag.size 
    val n_entries = (if(size >= 16u) then (size - 16u) / boot_info.memory_map.entry_size else 0u) : uint
  in
    ui2sz n_entries
  end

  implement _get_memory_mapping(boot_info,n) =  (
    assertloc (n < _get_memory_mappings_n boot_info);
    $UN.ptr0_get<memory_area_t>(ptr_add<memory_area_t>(boot_info.memory_map.entries,n))
  )

  //TODO: use foreach with fwork to loop through memory maps
  implement _print_memory_mappings(boot_info) = let
    val length = _get_memory_mappings_n(boot_info)
    fun loop {n,i : nat | i < n} .<n-i>. (i : size_t i, n : size_t n):<!wrt> void = let 
      val entry = _get_memory_mapping(boot_info,i)
    in
      if (entry.type = 1) then (
        print("   Start Address: ");
        print_hex(entry.base_addr);
        print(", Length: ");
        println_hex(entry.length);
      );
      if (i < n-1) then
        loop (succ(i),n)
    end
  in
    if (length > 0) then (
      println!("Memory areas:");
      loop(i2sz(0),length))
    else
      println!("No memory mappings")
  end


  //----------------------------------------------

  //---------------------ELF----------------------

  implement _get_elf_headers_n(boot_info) = ui2sz boot_info.elf_tag.num

  implement _get_elf_header (boot_info,n) = (
    assertloc (n < _get_elf_headers_n boot_info);
    $UN.ptr0_get<elf64_shdr_t>(ptr_add<elf64_shdr_t>(boot_info.elf_tag.headers,n))
  )

  implement _get_kernel_ranges(boot_info) = let
    val length = _get_elf_headers_n(boot_info)
    fun loop {n,i : nat | i < n} .<n-i>. (i : size_t i,n : size_t n, min : Ptr, max : Ptr) :<> (Ptr,Ptr) = let
      val header = _get_elf_header(boot_info,i)
      val new_min = (if ((header.addr < min || min = the_null_ptr) && header.size != 0) then header.addr else min)
      val end_p = ptr_add<uint8>(header.addr,header.size)
      val new_max = (
        if (end_p > max && header.size != 0) 
          then $UN.cast2Ptr1(end_p) 
        else max
      )
    in
      if (i < n - 1) then
        loop(succ(i),n,new_min,new_max)
      else 
        (new_min,new_max)
    end

  in
    if (length > 0) then
      loop(i2sz(0),length,the_null_ptr,the_null_ptr)
    else
      (the_null_ptr,the_null_ptr)
  end

  implement _get_multiboot_ranges(boot_info) = (boot_info.address, $UN.cast2Ptr1(ptr_add<uint8>(boot_info.address,boot_info.total_size)))

  implement _print_elf_headers(boot_info) = let
    val length = _get_elf_headers_n(boot_info)
    fun loop {n,i : nat | i < n} .<n-i>. (i : size_t i, n : size_t n) :<!wrt> void = let
      val header = _get_elf_header(boot_info,i)
    in
      print("   addr: ");
      print_hex(header.addr);
      print(", size: ");
      print_hex(header.size);
      print(", flags: ");
      println_hex(header.flags);
        
      if (i < n - 1) then
        loop(succ(i),n)

    end
  in
    if (length > 0) then (
      println!("Elf section headers: ");
      loop(i2sz(0),length))
    else
      println!("No elf section headers")
  end

end