#include "kernel/prelude/kernel_prelude.hats"

#define ATS_DYNLOADFLAG 0

staload "./multiboot.sats"
staload UN = "prelude/SATS/unsafe.sats"

local

var boot_info : boot_info_t

in

implement boot_info_init (p : Ptr1) = let
  extern praxi __assert{l:addr} (ptr: ptr (l)): vtakeout0 (boot_info_t@l) // takeout proof  UNSAFE:
  prval (pf, fpf) = __assert (addr@boot_info)
  val total_size = $UN.ptr0_get<uint>(p)

  fun loop {l : agz} (pf : !boot_info_t@l | boot_p : ptr l , p : ptr): void = let
    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_p->memory_map := $UN.ptr0_get<memory_areas_t>(p); 
               boot_p->memory_map.entries := ptr_add<uint32>(p,4)
              )
      | 9u => (
               boot_p->elf_tag := $UN.ptr0_get<elf_tag_t>(p);
               boot_p->elf_tag.headers := ptr_add<uint32>(p,5);
              )
      | _ => ();

    if (type != 0u) then
      loop(pf | boot_p, ptr_add<uint8>(p, next))
  end

in
  boot_info.total_size := total_size;
  loop(pf | addr@boot_info, ptr_add<uint32>(p,2));

  (pf , fpf | addr@boot_info)
end

end

extern castfn ui2sz (n : uint) : [n : nat] size_t n

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

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

// TODO: here use optional datatype (fix runtime alloc)

implement get_memory_mapping(pf | p,n) =  (
  assertloc(n < get_memory_mappings_n(pf | p));
  $UN.ptr0_get<memory_area_t>(ptr_add<memory_area_t>(p->memory_map.entries,n))
)

//TODO: use foreach with fwork to loop through memory maps
implement print_memory_mappings(p) = let
  val length = get_memory_mappings_n(p.0 | p.2)
  fun loop {n,i : nat | i < n} .<n-i>. (p : !bootptr, i : size_t i, n : size_t n) : void = let 
    val entry = get_memory_mapping(p.0 | p.2 ,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 (p,succ(i),n)
  end
in
  if (length > 0) then (
    println!("Memory areas:");
    loop(p,i2sz(0),length))
  else
    println!("No memory mappings")
end

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

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

implement get_elf_headers_n(p) = ui2sz p.2->elf_tag.num

implement get_elf_header (p,n) = (
  assertloc(n < get_elf_headers_n(p));
  $UN.ptr0_get<elf64_shdr_t>(ptr_add<elf64_shdr_t>(p.2->elf_tag.headers,n))
)

implement get_kernel_ranges(p) = let
  val length = get_elf_headers_n(p)
  fun loop {n,i : nat | i < n} .<n-i>. (p : !bootptr, i : size_t i,n : size_t n, min : Ptr, max : Ptr) : (Ptr,Ptr) = let
    val header = get_elf_header(p,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(p,succ(i),n,new_min,new_max)
    else 
      (new_min,new_max)
  end

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

implement get_bootinfo_ranges(p) = (p.2, ptr_add<uint8>(p.2,p.2->total_size))

implement print_elf_headers(p) = let
  val length = get_elf_headers_n(p)
  fun loop {n,i : nat | i < n} .<n-i>. (p : !bootptr, i : size_t i, n : size_t n) : void = let
    val header = get_elf_header(p,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(p,succ(i),n)

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

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