#include <stdint.h>
#include <string.h>
#include <sys.h>
#include <panic.h>

#include "memory.h"
#include "boot/tag.h"
#include "print.h"

struct MemoryArea {
    uint32_t len;
    struct MemoryArea *prev;
    struct MemoryArea *next;
};

typedef unsigned char page[4096];

extern unsigned char kernel_start, kernel_end;
static uintptr_t kernel_start_addr, kernel_end_addr;
static uint32_t *bitmap;
static uint32_t total_memory;
static uint32_t free_memory;
static uint32_t page_count;
static uint32_t page_free_start;
static struct MemoryArea *page_start;

int memory_lock(void) {
    int_disable();
    return 0;
}

int memory_unlock(void) {
    int_enable();
    return 0;
}

static int n_pages(const struct MemoryArea *m) {
	return (m->len - sizeof(*m)) / sizeof(page);
}

static void *page_at(int i) {
	int cur_page = 0;
	for (struct MemoryArea *m = page_start; m != NULL; m = m->next) {
		int pages = n_pages(m);
		if (i - cur_page < pages) {
			page *page_array = (page *) (m + 1);
			return page_array[i - cur_page];
		}
		cur_page += pages;
	}
	return NULL;
}

static int page_idx(page p) {
	uintptr_t addr = (uintptr_t) p;
	int cur_page = 0;
	for (struct MemoryArea *m = page_start; m != NULL; m = m->next) {
		if ((uintptr_t) m + m->len > addr) {
			return cur_page + (addr - (uintptr_t) m) / sizeof(page);
		}
		cur_page += n_pages(m);
	}
	return -1;
}

static inline bool bitmap_get(int i) {
	return (bitmap[i / 32] >> i % 32) & 1;
}

static inline void bitmap_set(int i, bool v) {
	int idx = i / 32;
	bitmap[idx] &= ~(1 << i % 32);
	bitmap[idx] |=  (v << i % 32);
}

void *memory_alloc_page(int pages) {
	if (pages < 1) return NULL;

	int n_contiguous = 0;
	int free_region_start = 0;
	bool first = true;
	for (uint32_t i = page_free_start; i < page_count; i++) {
		bool free = !bitmap_get(i);

		if (first) {
			first = false;
			page_free_start = i;
		}

		if (free) {
			if (n_contiguous == 0) free_region_start = i;
			n_contiguous++;
			if (n_contiguous == pages) {
				for (int j = 0; j < pages; j++)
					bitmap_set(free_region_start + j, true);
				return page_at(free_region_start);
			}
		} else n_contiguous = 0;
	}

	return NULL;
}

int memory_free_page(void *ptr, int pages) {
	int idx = page_idx(ptr);
	if (idx == -1) return 1;

	if ((unsigned) idx < page_free_start) page_free_start = idx;

	for (int i = 0; i < pages; i++)
		bitmap_set(idx + pages, false);
	return 0;
}

void memory_init(void) {

    debugk("Loading memory pages");

    memory_lock();

    bitmap = NULL;
    total_memory = 0;
    free_memory = 0;
    page_count = 0;
    page_free_start = 0;
    page_start = NULL;

    kernel_start_addr = (uintptr_t) &kernel_start;
    kernel_end_addr = (uintptr_t) &kernel_end;

    struct BootTag *tag;
    if (!get_boot_tag(iD_MEMORYMAP, &tag)) {
        panic("No multiboot memory map found");
    }

    uintptr_t end = (uintptr_t) tag->data.memory_map;
    end += tag->size;

    struct MemoryArea *prev = NULL;
    struct MemorySegment *segment = &tag->data.memory_map->entries[0];
    for(; (uintptr_t) segment < end; segment++) {

        if (segment->type != 1) continue;
        if (segment->addr >= UINT32_MAX) continue;
        if (segment->addr < kernel_start_addr) continue;

        uint32_t length;
        if (segment->addr + segment->len > UINT32_MAX) {
            length = UINT32_MAX - segment->addr;
        } else {
            length = segment->len;
        }

        uintptr_t addr;
        if (segment->addr < kernel_end_addr) {
            addr = kernel_end_addr;
            length -= addr - segment->addr;
        } else {
            addr = segment->addr;
        }

        struct MemoryArea *current = (struct MemoryArea *) addr;
        current->prev = prev;
        current->next = NULL;
        current->len = length;

        if (prev != NULL) {
            prev->next = current;
        } else {
            page_start = current;
        }

        page_count += n_pages(current);
        total_memory += length;

        prev = current;

    }

	int bitmap_pages = page_count / 32 / sizeof(page) + 1;
	bitmap = (uint32_t *) page_at(page_count - bitmap_pages);
    page_count -= bitmap_pages;
	memset(bitmap, 0, bitmap_pages * sizeof(page));
	free_memory = page_count * sizeof(page);

    memory_unlock();

    succek("Memory loaded. %k total %k free", total_memory, free_memory);

}

uint32_t memory_total(void) {
    return total_memory;
}

uint32_t memory_free(void) {
    return free_memory;
}

uint32_t memory_used(void) {
    return total_memory - free_memory;
}