#include "kernel/prelude/kernel_prelude.hats"

#define ATS_DYNLOADFLAG 0

staload BOOT = "kernel/bootinfo/multiboot.sats"

staload "./frame.sats"
staload "lib/SATS/init.sats"
staload "lib/DATS/init.dats"

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

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

fn containing_area(area : $BOOT.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(allocator : &frame_allocator_t) :<!wrt,!ref> void = let 
  val length = g1ofg0($BOOT.get_memory_mappings_n())
  val next_free_frame = allocator.next_free_frame
  fun loop {n,i : nat | i < n} .<n-i>. ( i : size_t i, n : size_t n) :<!ref> $BOOT.memory_area_t = let 
    val entry = $BOOT.get_memory_mapping(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 (succ(i),n)
    else
      entry //TODO: fix
  end
in 
  if (length > 0) then //TODO: fix
    allocator.current_area := loop(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

  // Private function
  extern fn _allocate_frame(allocator: &frame_allocator_t):<!wrt,!ref> frame_t

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

    implement init(kernel_start, kernel_end, multiboot_start, multiboot_end) = let
      val allocator = @{ 
        next_free_frame = containing_address(the_null_ptr),
        current_area = $BOOT.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

      implement initialize$fwork<frame_allocator_t>(v) = (
        v := allocator; 
        $effmask_ref(choose_next_area(v));
        let prval () = opt_some v in true end
      )
    in
      initialize<frame_allocator_t>(allocator_ref)
    end

    implement allocate_frame(): frame_t = let
      implement exec$fwork<frame_allocator_t><frame_t>(v) = $effmask_ref(_allocate_frame(v))
    in
      exec<frame_allocator_t><frame_t>(allocator_ref)
    end

  end

  implement _allocate_frame(allocator) : frame_t =
    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(allocator);
          _allocate_frame(allocator)
        ) 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};
          _allocate_frame(allocator)
        ) 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};
          _allocate_frame(allocator)
        ) else (
          // frame is unused, increment `next_free_frame` and return it
          allocator.next_free_frame.num := succ(allocator.next_free_frame.num);
          frame
        );
      end
    else
      @{num = i2sz(0)} // No free frames left


end