#include "kernel/prelude/kernel_prelude.hats"

#define ATS_DYNLOADFLAG 0

staload "./frame.sats"
staload "kernel/bootinfo/multiboot.sats"
staload "lib/SATS/init.sats"

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

assume frame_allocator = frame_allocator_t


implement containing_address(address): frame_t = 
  @{ 
    counter = $UN.cast{size_t}(address) / i2sz(PAGE_SIZE)
  }

fn containing_area(area : memory_area_t) :<> frame_t = containing_address(ptr_add<uint8>(area.base_addr,area.length - i2sz(1)))

//FIX: function doesn't handle no frames left
fn choose_next_area{k : agz}(bf : !boot_info_t @ k | b : ptr k, allocator : &frame_allocator_t) :<!wrt> void = let 
  val length = get_memory_mappings_n(bf | b)
  val next_free_frame = allocator.next_free_frame
  fun loop {n,i : nat | i < n} {l : agz} .<n-i>. (bf : !boot_info_t @ l | b : ptr l, i : size_t i, n : size_t n) :<> memory_area_t = let 
    val entry = get_memory_mapping(bf | b ,i)
    val fr = containing_area(entry)
  in
    if (entry.type = 1 && fr.num >= next_free_frame.num) then
      entry
    else if (i < n-1) then
      loop (bf | b,succ(i),n)
    else
      entry //TODO: fix
  end
in 
  if (length > 0) then //TODO: fix
    allocator.current_area := loop(bf | b, i2sz(0), length);

  if (allocator.current_area.length > 1) then
    let
      val start_frame = containing_address(allocator.current_area.base_addr)
    in
      if (allocator.next_free_frame.num < start_frame.num) then
        allocator.next_free_frame := start_frame
    end
end 

local

var allocator_static : frame_allocator_t?
prval () = opt_none allocator_static
var allocator = empty<frame_allocator_t>(allocator_static)
val allocator_ref = ref_make_viewptr{enablable frame_allocator_t} (view@allocator | addr@allocator)

in

  implement frame_allocator_init(kernel_start, kernel_end, multiboot_start, multiboot_end, b) : void = let
    val allocator = @{ 
      next_free_frame = containing_address(the_null_ptr),
      current_area = invalid_area,
      kernel_start = containing_address(kernel_start),
      kernel_end = containing_address(kernel_end),
      multiboot_start = containing_address(multiboot_start),
      multiboot_end = containing_address(multiboot_end)
    } : frame_allocator_t
    val (vbox pf | p) = ref_get_viewptr(allocator_ref)
  in
    if not p->enabled then let 
        prval () = opt_unnone p->obj
      in 
        p->obj := allocator;
        p->enabled := true;
        (* choose_next_area(b.0 | b.2, !p); *)
        let prval () = opt_some p->obj in end;
      end
  end

implement allocate_frame(bf | b) : frame_t = let
    val (vbox pf | p) = ref_get_viewptr(allocator_ref)
  in
    if p->enabled then let
      (* var allocator = p->obj *)
        prval () = opt_unsome p->obj
      in
        (* if (allocator.current_area.type != 0u) then  let *)
        (*     val area = allocator.current_area *)
        (*     val frame = @{ num = allocator.next_free_frame.num } *)
        (**)
        (*     // last frame of current area *)
        (*     val last_frame_area = containing_area(area) *)
        (*   in *)
        (*     if (frame.num > last_frame_area.num) then ( *)
        (*       // all frames of current area are used, switch to next area *)
        (*       choose_next_area(bf | b, allocator); *)
        (*       let prval () = opt_some allocator in end; *)
        (*       allocate_frame(bf | b) *)
        (*     ) else if (frame.num >= allocator.kernel_start.num && frame.num <= allocator.kernel_end.num) then ( *)
        (*       // frame is used by kernel *)
        (*       allocator.next_free_frame := @{num = allocator.kernel_end.num + 1}; *)
        (*       let prval () = opt_some allocator in end; *)
        (*       allocate_frame(bf | b) *)
        (*     ) else if (frame.num >= allocator.multiboot_start.num && frame.num <= allocator.multiboot_end.num) then ( *)
        (*       // frame is used by multiboot info structure *)
        (*       allocator.next_free_frame := @{num = allocator.multiboot_end.num + 1}; *)
        (*       let prval () = opt_some allocator in end; *)
        (*       allocate_frame(bf | b) *)
        (*     ) else ( *)
        (*       // frame is unused, increment `next_free_frame` and return it *)
        (*       allocator.next_free_frame.num := succ(allocator.next_free_frame.num); *)
        (*       let prval () = opt_some allocator in end; *)
        (*       frame *)
        (*     ); *)
        (*   end *)
        (* else *)
        (*   let prval () = opt_some allocator in end; *)
        (*   @{num = i2sz(0)} // No free frames left *)
        let prval () = opt_some p->obj in end;
        @{num = i2sz(0)}
      end
    else
    @{num = i2sz(0)} // No free frames left
    
end

end

implement deallocate_frame(p) : void = let
in
 
end