#include <kernel/address-space.h>
#include <kernel/iovec.h>
#include <kernel/libc/stdio.h>
#include <kernel/object.h>
#include <kernel/panic.h>
#include <kernel/printk.h>
#include <kernel/util.h>
#include <kernel/vm-object.h>
#include <mango/status.h>

/*** STATIC DATA + MACROS *****************************************************/

#define ADDRESS_SPACE_CAST(p)                                                  \
	OBJECT_C_CAST(struct address_space, s_base, &address_space_type, p)
#define INVALID_OFFSET ((off_t) - 1)

enum get_entry_flags {
	GET_ENTRY_EXACT = 0,
	GET_ENTRY_CLOSEST_LEFT,
	GET_ENTRY_CLOSEST_RIGHT,
};

/* iterates over a range of mapped virtual memory in a region, and provides
 * a moving buffer through which the memory can be accessed */
struct vm_iterator {
	struct address_space *it_region;
	struct vm_area *it_mapping;
	virt_addr_t it_base;
	vm_prot_t it_prot;
	void *it_buf;
	size_t it_max;
};

/* iterates over the areas in an address space */
struct area_iterator {
	struct address_space *it_root;
	struct vm_area *it_area;
	virt_addr_t it_search_base, it_search_limit;
	virt_addr_t it_base, it_limit;
};

enum search_direction {
	SEARCH_LEFT,
	SEARCH_RIGHT,
};

static kern_status_t address_space_object_destroy(struct object *obj);

static struct object_type address_space_type = {
	.ob_name = "address-space",
	.ob_size = sizeof(struct address_space),
	.ob_header_offset = offsetof(struct address_space, s_base),
};

static struct vm_cache vm_area_cache = {
	.c_name = "vm-area",
	.c_obj_size = sizeof(struct vm_area),
};

/*** INTERNAL UTILITY FUNCTION ************************************************/

/* this function must be called with `parent` locked */
static void put_entry(struct btree *tree, struct vm_area *child)
{
	struct btree_node *cur = tree->b_root;
	if (!cur) {
		tree->b_root = &child->vma_node;
		btree_insert_fixup(tree, &child->vma_node);
		return;
	}

	while (cur) {
		struct vm_area *cur_entry
			= BTREE_CONTAINER(struct vm_area, vma_node, cur);

		struct btree_node *next = NULL;

		if (child->vma_limit < cur_entry->vma_base) {
			next = btree_left(cur);
		} else if (child->vma_base > cur_entry->vma_limit) {
			next = btree_right(cur);
		} else {
			panic("tried to add an overlapping entry [%zx-%zx] to "
			      "vm-region (overlaps [%zx-%zx])",
			      child->vma_base,
			      child->vma_limit,
			      cur_entry->vma_base,
			      cur_entry->vma_limit);
		}

		if (next) {
			cur = next;
			continue;
		}

		if (child->vma_limit < cur_entry->vma_base) {
			btree_put_left(cur, &child->vma_node);
		} else {
			btree_put_right(cur, &child->vma_node);
		}

		btree_insert_fixup(tree, &child->vma_node);
		break;
	}
}

static struct vm_area *get_entry(
	struct address_space *region,
	virt_addr_t address,
	enum get_entry_flags flags)
{
	/* `x` must be to the left of `y` */
#define LEFT_DIFF(x, y)  ((y) ? ((y)->vma_base - (x)) : ((size_t)-1))
	/* `x` must be to the right of `y` */
#define RIGHT_DIFF(x, y) ((y) ? ((y)->vma_limit - (x)) : ((size_t)-1))

	struct btree_node *cur = region->s_mappings.b_root;
	if (!cur) {
		return NULL;
	}

	struct vm_area *result = NULL;
	struct vm_area *closest_left = NULL;
	struct vm_area *closest_right = NULL;

	while (cur) {
		struct vm_area *child
			= BTREE_CONTAINER(struct vm_area, vma_node, cur);

		struct btree_node *next = NULL;

		if (address < child->vma_base) {
			next = btree_left(cur);
			if (LEFT_DIFF(address, child)
			    < LEFT_DIFF(address, closest_left)) {
				closest_left = child;
			}
		} else if (address > child->vma_limit) {
			next = btree_right(cur);
			if (RIGHT_DIFF(address, child)
			    < RIGHT_DIFF(address, closest_right)) {
				closest_right = child;
			}
		} else {
			result = child;
			break;
		}

		cur = next;
	}

	if (result) {
		return result;
	}

	if (flags & GET_ENTRY_CLOSEST_LEFT) {
		return closest_left;
	}

	if (flags & GET_ENTRY_CLOSEST_RIGHT) {
		return closest_right;
	}

	return NULL;
#undef LEFT_DIFF
#undef RIGHT_DIFF
}

/* does not consider reserved areas! */
static bool is_area_free(
	const struct address_space *space,
	virt_addr_t base,
	size_t len)
{
	virt_addr_t limit = base + len - 1;

	if (base < space->s_base_address) {
		return false;
	}

	if (base + limit > space->s_limit_address) {
		return false;
	}

	struct btree_node *cur = space->s_mappings.b_root;
	if (!cur) {
		return true;
	}

	while (cur) {
		struct vm_area *cur_area
			= BTREE_CONTAINER(struct vm_area, vma_node, cur);

		if (base >= cur_area->vma_base && base <= cur_area->vma_limit) {
			return false;
		}

		if (limit >= cur_area->vma_base
		    && limit <= cur_area->vma_limit) {
			return false;
		}

		if (base < cur_area->vma_base && limit > cur_area->vma_limit) {
			return false;
		}

		if (base > cur_area->vma_limit) {
			cur = btree_right(cur);
		} else if (limit < cur_area->vma_base) {
			cur = btree_left(cur);
		} else {
			/* what */
			panic("unhandled case in is_area_free. base=%zx, "
			      "len=%zx, "
			      "limit=%zx, cur_area=[%zx-%zx]",
			      base,
			      len,
			      limit,
			      cur_area->vma_base,
			      cur_area->vma_limit);
		}
	}

	return true;
}

/* ONLY considers reserved areas! */
static bool is_area_reserved(
	const struct address_space *space,
	virt_addr_t base,
	size_t len)
{
	virt_addr_t limit = base + len - 1;

	if (base < space->s_base_address) {
		return false;
	}

	if (base + limit > space->s_limit_address) {
		return false;
	}

	struct btree_node *cur = space->s_reserved.b_root;
	if (!cur) {
		return false;
	}

	while (cur) {
		struct vm_area *cur_area
			= BTREE_CONTAINER(struct vm_area, vma_node, cur);

		if (base >= cur_area->vma_base && base <= cur_area->vma_limit) {
			return true;
		}

		if (limit >= cur_area->vma_base
		    && limit <= cur_area->vma_limit) {
			return true;
		}

		if (base < cur_area->vma_base && limit > cur_area->vma_limit) {
			return false;
		}

		if (base > cur_area->vma_limit) {
			cur = btree_right(cur);
		} else if (limit < cur_area->vma_base) {
			cur = btree_left(cur);
		} else {
			/* what */
			panic("unhandled case in is_area_reserved. base=%zx, "
			      "len=%zx, "
			      "limit=%zx, cur_area=[%zx-%zx]",
			      base,
			      len,
			      limit,
			      cur_area->vma_base,
			      cur_area->vma_limit);
		}
	}

	return false;
}

static virt_addr_t generate_address(
	virt_addr_t lower_bound,
	virt_addr_t upper_bound)
{
	virt_addr_t result = 0;
	fill_random(&result, sizeof result);

	virt_addr_t mask = upper_bound;

	result += lower_bound;
	result &= mask;
	result &= ~VM_PAGE_MASK;

	return result;
}

static virt_addr_t find_free_area(
	struct address_space *space,
	size_t target_length)
{
	virt_addr_t search_limit = space->s_base_address << 16;

	int attempt = 0;
	while (1) {
		virt_addr_t base = generate_address(
			space->s_base_address,
			search_limit - 1);
		bool ok = true;

		if (is_area_reserved(space, base, target_length)) {
			ok = false;
		}

		if (ok && !is_area_free(space, base, target_length)) {
			ok = false;
		}

		if (ok) {
			return base;
		}

		attempt++;
		if (attempt >= 3) {
			search_limit <<= 4;
			attempt = 0;
		}
	}

	return 0;
}

/* this function should be called with `region` locked */
static void vm_iterator_begin(
	struct vm_iterator *it,
	struct address_space *region,
	virt_addr_t base,
	vm_prot_t prot)
{
	memset(it, 0x0, sizeof *it);
	it->it_base = base;
	it->it_region = region;
	it->it_prot = prot;

	it->it_mapping = get_entry(region, base, GET_ENTRY_EXACT);
	if (!it->it_mapping) {
		return;
	}

	if ((it->it_mapping->vma_prot & prot) != prot) {
		return;
	}

	off_t object_offset = base - it->it_mapping->vma_base
	                    + it->it_mapping->vma_object_offset;
	struct vm_page *pg = NULL;
	enum vm_object_flags flags = 0;
	if (prot & VM_PROT_WRITE) {
		flags |= VMO_ALLOCATE_MISSING_PAGE;
	}

	pg = vm_object_get_page(
		it->it_mapping->vma_object,
		object_offset,
		flags,
		NULL);

	if (!pg) {
		return;
	}

	void *buffer_base = vm_page_get_vaddr(pg);
	phys_addr_t pg_addr = vm_page_get_paddr(pg);
	size_t buffer_size = vm_page_get_size_bytes(pg);

	while (1) {
		struct btree_node *next_node = btree_next(&pg->p_bnode);
		struct vm_page *next
			= BTREE_CONTAINER(struct vm_page, p_bnode, next_node);
		if (!next) {
			break;
		}

		phys_addr_t next_addr = vm_page_get_paddr(next);
		if (pg_addr + vm_page_get_size_bytes(pg) != next_addr) {
			break;
		}

		pg = next;
		pg_addr = next_addr;
		buffer_size += vm_page_get_size_bytes(next);
	}

	it->it_buf = (char *)buffer_base + (object_offset & VM_PAGE_MASK);
	it->it_max = buffer_size - (object_offset & VM_PAGE_MASK);
}

static kern_status_t vm_iterator_seek(struct vm_iterator *it, size_t nr_bytes)
{
	if (nr_bytes < it->it_max) {
		it->it_base += nr_bytes;
		it->it_buf = (char *)it->it_buf + nr_bytes;
		it->it_max -= nr_bytes;
		return KERN_OK;
	}

	it->it_base += nr_bytes;

	struct vm_area *next_mapping
		= get_entry(it->it_region, it->it_base, GET_ENTRY_EXACT);
	if (!next_mapping) {
		it->it_buf = NULL;
		it->it_max = 0;
		return KERN_MEMORY_FAULT;
	}

	if ((next_mapping->vma_prot & it->it_prot) != it->it_prot) {
		it->it_buf = NULL;
		it->it_max = 0;
		return KERN_MEMORY_FAULT;
	}

	off_t object_offset = it->it_base - it->it_mapping->vma_base
	                    + it->it_mapping->vma_object_offset;

	struct vm_page *pg = NULL;
	enum vm_object_flags flags = 0;
	if (it->it_prot & VM_PROT_WRITE) {
		flags |= VMO_ALLOCATE_MISSING_PAGE;
	}

	pg = vm_object_get_page(
		it->it_mapping->vma_object,
		object_offset,
		flags,
		NULL);

	if (!pg) {
		return KERN_NO_MEMORY;
	}

	void *buffer_base = vm_page_get_vaddr(pg);
	phys_addr_t pg_addr = vm_page_get_paddr(pg);
	size_t buffer_size = vm_page_get_size_bytes(pg);

	while (1) {
		struct btree_node *next_node = btree_next(&pg->p_bnode);
		struct vm_page *next
			= BTREE_CONTAINER(struct vm_page, p_bnode, next_node);
		if (!next) {
			break;
		}

		phys_addr_t next_addr = vm_page_get_paddr(next);
		if (pg_addr + vm_page_get_size_bytes(pg) != next_addr) {
			break;
		}

		pg = next;
		pg_addr = next_addr;
		buffer_size += vm_page_get_size_bytes(next);
	}

	it->it_buf = (char *)buffer_base + (object_offset & VM_PAGE_MASK);
	it->it_max = buffer_size;
	return KERN_OK;
}

static void vm_iterator_finish(struct vm_iterator *it)
{
	memset(it, 0x0, sizeof *it);
}

/* this function must be called with `space` locked */
static void area_iterator_begin(
	struct area_iterator *it,
	struct address_space *space,
	virt_addr_t base,
	virt_addr_t limit)
{
	memset(it, 0x0, sizeof *it);

	struct vm_area *area = get_entry(space, base, GET_ENTRY_CLOSEST_RIGHT);
	if (!area) {
		return;
	}

	if (area->vma_base > limit) {
		return;
	}

	it->it_search_base = base;
	it->it_search_limit = limit;
	it->it_root = space;
	it->it_area = area;
	it->it_base = area->vma_base;
	it->it_limit = area->vma_base;

	if (it->it_base < base) {
		it->it_base = base;
	}

	if (it->it_limit > limit) {
		it->it_limit = limit;
	}
}

static void area_iterator_finish(struct area_iterator *it)
{
	memset(it, 0x0, sizeof *it);
}

static kern_status_t area_iterator_move_next(struct area_iterator *it)
{
	if (!it->it_root || !it->it_area) {
		return KERN_NO_ENTRY;
	}

	struct btree_node *next = btree_next(&it->it_area->vma_node);
	if (!next) {
		goto end;
	}

	struct vm_area *area = BTREE_CONTAINER(struct vm_area, vma_node, next);
	if (!area) {
		goto end;
	}

	if (area->vma_base > it->it_search_limit) {
		goto end;
	}

	it->it_area = area;
	it->it_base = area->vma_base;
	it->it_limit = area->vma_base;

	if (it->it_base < it->it_search_base) {
		it->it_base = it->it_search_base;
	}

	if (it->it_limit > it->it_search_limit) {
		it->it_limit = it->it_search_limit;
	}

	return KERN_OK;

end:
	memset(it, 0x0, sizeof *it);
	return KERN_NO_ENTRY;
}

static void area_iterator_erase(struct area_iterator *it)
{
}

/*** PUBLIC API ***************************************************************/

kern_status_t address_space_type_init(void)
{
	vm_cache_init(&vm_area_cache);
	return object_type_register(&address_space_type);
}

struct address_space *address_space_cast(struct object *obj)
{
	return ADDRESS_SPACE_CAST(obj);
}

/* this function should be called with `parent` locked (if parent is
 * non-NULL)
 */
kern_status_t address_space_create(
	virt_addr_t base,
	virt_addr_t limit,
	struct address_space **out)
{
	if (!base || !limit || limit <= base) {
		return KERN_INVALID_ARGUMENT;
	}

	if ((base & VM_PAGE_MASK) || ((limit + 1) & VM_PAGE_MASK)) {
		return KERN_INVALID_ARGUMENT;
	}

	struct object *region_object = object_create(&address_space_type);
	if (!region_object) {
		return KERN_NO_MEMORY;
	}

	struct address_space *space = ADDRESS_SPACE_CAST(region_object);

	space->s_base_address = base;
	space->s_limit_address = limit;
#ifdef TRACE
	tracek("creating address space at [%llx-%llx]", base, limit);
#endif

	*out = space;
	return KERN_OK;
}

kern_status_t address_space_map(
	struct address_space *root,
	virt_addr_t map_address,
	struct vm_object *object,
	off_t object_offset,
	size_t length,
	vm_prot_t prot,
	virt_addr_t *out)
{
	if (object_offset & VM_PAGE_MASK) {
		object_offset &= ~VM_PAGE_MASK;
	}

	if (length & VM_PAGE_MASK) {
		length &= ~VM_PAGE_MASK;
		length += VM_PAGE_SIZE;
	}

	if (map_address != MAP_ADDRESS_ANY && (map_address & VM_PAGE_MASK)) {
		map_address &= ~VM_PAGE_MASK;
	}

	tracek("address_space_map(%zx, %zx)", map_address, length);

	if (!root || !object) {
		tracek("null pointer");
		return KERN_INVALID_ARGUMENT;
	}

	if ((prot & object->vo_prot) != prot) {
		tracek("protection error");
		return KERN_INVALID_ARGUMENT;
	}

	if (!length || object_offset + length > object->vo_size) {
		tracek("length exceeds object bounds");
		return KERN_INVALID_ARGUMENT;
	}

	if (map_address == MAP_ADDRESS_ANY) {
		map_address = find_free_area(root, length);

		if (map_address == MAP_ADDRESS_INVALID) {
			tracek("no virtual memory available");
			return KERN_NO_MEMORY;
		}
	} else if (!is_area_free(root, map_address, length)) {
		tracek("area already in use");
		return KERN_INVALID_ARGUMENT;
	}

	struct vm_area *area = vm_cache_alloc(&vm_area_cache, VM_NORMAL);
	if (!area) {
		return KERN_NO_MEMORY;
	}

	object_ref(&object->vo_base);
	area->vma_object = object;
	area->vma_prot = prot;
	area->vma_object_offset = object_offset;
	area->vma_base = map_address;
	area->vma_limit = map_address + length - 1;

#ifdef TRACE
	tracek("mapping %s at [%llx-%llx]",
	       object->vo_name,
	       area->vma_base,
	       area->vma_base + length);
#endif
	put_entry(&root->s_mappings, area);

	unsigned long lock_flags;
	vm_object_lock_irqsave(object, &lock_flags);
	queue_push_back(&object->vo_mappings, &area->vma_object_entry);
	vm_object_unlock_irqrestore(object, lock_flags);

	if (out) {
		*out = map_address;
	}

	return KERN_OK;
}

/* unmap some pages in the middle of an area, splitting it into two
 * separate mappings */
static kern_status_t split_area(
	struct vm_area *mapping,
	struct address_space *root,
	virt_addr_t unmap_base,
	virt_addr_t unmap_limit)
{
	struct vm_area *left = mapping;
	struct vm_area *right = vm_cache_alloc(&vm_area_cache, VM_NORMAL);
	if (!right) {
		return KERN_NO_MEMORY;
	}

	virt_addr_t left_base = mapping->vma_base;
	virt_addr_t right_base = unmap_limit;
	off_t left_object_offset = mapping->vma_object_offset;
	size_t left_length = unmap_base - mapping->vma_base;
	size_t right_length = mapping->vma_limit - unmap_limit;
	off_t right_object_offset = mapping->vma_limit - right_length;

	tracek("mapping=[%zx-%zx]->[%zx-%zx]",
	       mapping->vma_base,
	       mapping->vma_limit,
	       mapping->vma_object_offset,
	       mapping->vma_object_offset
	               + (mapping->vma_limit - mapping->vma_base));
	tracek("left=[%zx-%zx]->[%zx-%zx], right=[%zx-%zx]->[%zx-%zx]",
	       left_base,
	       left_base + left_length,
	       left_object_offset,
	       left_object_offset + left_length,
	       right_base,
	       right_base + right_length,
	       right_object_offset,
	       right_object_offset + right_length);

	left->vma_object_offset = left_object_offset;
	left->vma_base = left_base;
	left->vma_limit = left_base + left_length - 1;

	right->vma_object = left->vma_object;
	right->vma_prot = left->vma_prot;
	right->vma_object_offset = right_object_offset;
	right->vma_base = right_base;
	right->vma_limit = right_base + right_length - 1;

	if (!mapping->vma_object) {
		put_entry(&root->s_reserved, right);
		/* just a reservation, no page tables to update */
		return KERN_OK;
	}

	put_entry(&root->s_mappings, right);
	for (size_t i = unmap_base; i < unmap_limit; i += VM_PAGE_SIZE) {
		tracek("unmapping %zx", i);
		pmap_remove(root->s_pmap, i);
	}

	return KERN_OK;
}

/* unmap some pages from the left-side of a mapping to somewhere in the
 * middle. */
static kern_status_t left_reduce_area(
	struct vm_area *mapping,
	struct address_space *root,
	virt_addr_t unmap_base,
	virt_addr_t unmap_limit)
{
	tracek("left reduce mapping [%zx-%zx] subtract [%zx-%zx]",
	       mapping->vma_base,
	       mapping->vma_limit,
	       unmap_base,
	       unmap_limit);

	virt_addr_t base = mapping->vma_base;
	virt_addr_t limit = unmap_limit;
	size_t length = limit - base + 1;

	mapping->vma_base += length;
	mapping->vma_object_offset += length;

	if (!mapping->vma_object) {
		/* just a reservation, no page tables to update */
		tracek(" unreserving %zx-%zx (%zx bytes)",
		       base,
		       base + length,
		       length);
		return KERN_OK;
	}

	tracek(" unmapping %zx-%zx (%zx bytes)", base, base + length, length);
	for (size_t i = base; i < limit; i += VM_PAGE_SIZE) {
		pmap_remove(root->s_pmap, i);
	}

	return KERN_OK;
}

/* unmap some pages from the middle of a mapping to the right-side. */
static kern_status_t right_reduce_area(
	struct vm_area *mapping,
	struct address_space *root,
	virt_addr_t unmap_base,
	virt_addr_t unmap_limit)
{
	/* unmap_base falls somwwhere between mapping_offset and
	 * mapping_offset+length */
	tracek("right reduce mapping [%zx-%zx] subtract [%zx-%zx]",
	       mapping->vma_base,
	       mapping->vma_limit,
	       unmap_base,
	       unmap_limit);

	virt_addr_t base = unmap_base;
	virt_addr_t limit = mapping->vma_limit;
	size_t length = limit - base + 1;
	mapping->vma_limit -= length;

	if (!mapping->vma_object) {
		/* just a reservation, no page tables to update */
		tracek(" unreserving %zx-%zx (%zx bytes)",
		       base,
		       base + length,
		       length);
		return KERN_OK;
	}

	tracek(" unmapping %zx-%zx (%zx bytes)", base, base + length, length);
	for (size_t i = base; i < limit; i += VM_PAGE_SIZE) {
		pmap_remove(root->s_pmap, i);
	}

	return KERN_OK;
}

/* completely unmap and delete an entire mapping */
static kern_status_t delete_area(
	struct vm_area *mapping,
	struct address_space *root)
{
	if (!mapping->vma_object) {
		/* just a reservation, no page tables to update */
		return KERN_OK;
	}

	tracek("delete mapping [%zx-%zx]",
	       mapping->vma_base,
	       mapping->vma_limit);

	for (size_t i = mapping->vma_base; i < mapping->vma_limit;
	     i += VM_PAGE_SIZE) {
		pmap_remove(root->s_pmap, i);
	}

	struct vm_object *object = mapping->vma_object;
	unsigned long flags;
	vm_object_lock_irqsave(mapping->vma_object, &flags);
	queue_delete(
		&mapping->vma_object->vo_mappings,
		&mapping->vma_object_entry);
	mapping->vma_object = NULL;
	vm_object_unlock_irqrestore(mapping->vma_object, flags);
	object_unref(&object->vo_base);

	/* don't actually delete the mapping yet. that will be done by
	 * address_space_unmap */

	return KERN_OK;
}

kern_status_t address_space_unmap(
	struct address_space *region,
	virt_addr_t unmap_base,
	size_t unmap_length)
{
	if (unmap_length == 0) {
		return KERN_OK;
	}

	if (unmap_base & VM_PAGE_MASK) {
		unmap_base &= ~VM_PAGE_MASK;
	}

	if (unmap_length & VM_PAGE_MASK) {
		unmap_length &= VM_PAGE_MASK;
		unmap_length += VM_PAGE_SIZE;
	}

	kern_status_t status = KERN_OK;
	struct area_iterator it;
	virt_addr_t unmap_limit = unmap_base + unmap_length - 1;
	tracek("unmapping %zx-%zx", unmap_base, unmap_limit);

	area_iterator_begin(&it, region, unmap_base, unmap_limit);
	while (it.it_area) {
		struct vm_area *area = it.it_area;
		virt_addr_t area_base = area->vma_base;
		virt_addr_t area_limit = area->vma_limit;

		bool split
			= (area_base > unmap_base && area_limit < unmap_limit);
		bool delete
			= (area_base <= unmap_base
		           && area_limit >= unmap_limit);
		bool left_reduce
			= (unmap_base <= area_base && unmap_limit < area_limit);
		bool right_reduce
			= (unmap_base > area_base && unmap_limit >= area_limit);

		if (split) {
			status = split_area(
				area,
				region,
				unmap_base,
				unmap_limit);
			delete = true;
		} else if (delete) {
			status = delete_area(area, region);
		} else if (left_reduce) {
			status = left_reduce_area(
				area,
				region,
				unmap_base,
				unmap_limit);
		} else if (right_reduce) {
			status = right_reduce_area(
				area,
				region,
				unmap_base,
				unmap_limit);
		} else {
			panic("don't know what to do with this "
			      "mapping");
		}

		if (delete) {
			area_iterator_erase(&it);
		} else {
			area_iterator_move_next(&it);
		}

		if (status != KERN_OK) {
			break;
		}
	}

	area_iterator_finish(&it);

	return status;
}

kern_status_t address_space_reserve(
	struct address_space *space,
	virt_addr_t base,
	size_t length,
	virt_addr_t *out)
{
	if (length & VM_PAGE_MASK) {
		length &= ~VM_PAGE_MASK;
		length += VM_PAGE_SIZE;
	}

	if (base != MAP_ADDRESS_ANY && (base & VM_PAGE_MASK)) {
		base &= ~VM_PAGE_MASK;
	}

	if (!space) {
		return KERN_INVALID_ARGUMENT;
	}

	if (!length || base + length > space->s_limit_address) {
		return KERN_INVALID_ARGUMENT;
	}

	if (base == MAP_ADDRESS_ANY) {
		base = find_free_area(space, length);

		if (base == MAP_ADDRESS_INVALID) {
			return KERN_NO_MEMORY;
		}
	} else {
		if (is_area_reserved(space, base, length)) {
			/* for now, don't figure out overlapping reservations */
			return KERN_INVALID_ARGUMENT;
		}

		if (!is_area_free(space, base, length)) {
			return KERN_INVALID_ARGUMENT;
		}
	}

	struct vm_area *area = vm_cache_alloc(&vm_area_cache, VM_NORMAL);
	if (!area) {
		return KERN_NO_MEMORY;
	}

	area->vma_base = base;
	area->vma_limit = base + length - 1;

#ifdef TRACE
	tracek("reservation at [%llx-%llx]", area->vma_base, area->vma_limit);
#endif
	put_entry(&space->s_reserved, area);

	if (out) {
		*out = base;
	}

	return KERN_OK;
}

kern_status_t address_space_release(
	struct address_space *space,
	virt_addr_t release_base,
	size_t release_length)
{
	if ((release_base & VM_PAGE_MASK) || (release_length & VM_PAGE_MASK)) {
		return KERN_INVALID_ARGUMENT;
	}

	kern_status_t status = KERN_OK;
	struct area_iterator it;
	virt_addr_t release_limit = release_base + release_length - 1;
	tracek("unreserving %zx-%zx", release_base, release_limit);

	area_iterator_begin(&it, space, release_base, release_limit);
	while (it.it_area) {
		struct vm_area *area = it.it_area;
		virt_addr_t area_base = area->vma_base;
		virt_addr_t area_limit = area->vma_limit;

		bool split
			= (area_base > release_base
		           && area_limit < release_limit);
		bool delete
			= (area_base <= release_base
		           && area_limit >= release_limit);
		bool left_reduce
			= (release_base <= area_base
		           && release_limit < area_limit);
		bool right_reduce
			= (release_base > area_base
		           && release_limit >= area_limit);

		if (split) {
			status = split_area(
				area,
				space,
				release_base,
				release_limit);
			delete = true;
		} else if (delete) {
			status = delete_area(area, space);
		} else if (left_reduce) {
			status = left_reduce_area(
				area,
				space,
				release_base,
				release_limit);
		} else if (right_reduce) {
			status = right_reduce_area(
				area,
				space,
				release_base,
				release_limit);
		} else {
			panic("don't know what to do with this "
			      "mapping");
		}

		if (delete) {
			area_iterator_erase(&it);
		} else {
			area_iterator_move_next(&it);
		}

		if (status != KERN_OK) {
			break;
		}
	}

	area_iterator_finish(&it);

	return status;
}

bool address_space_validate_access(
	struct address_space *region,
	virt_addr_t ptr,
	size_t len,
	vm_prot_t prot)
{
	if (len == 0) {
		return true;
	}

	if (ptr < region->s_base_address) {
		return false;
	}

	if (ptr + len > region->s_limit_address) {
		return false;
	}

	virt_addr_t base = ptr & ~VM_PAGE_MASK;
	virt_addr_t limit = (ptr + len) - 1;
	if ((limit + 1) & VM_PAGE_MASK) {
		limit &= ~VM_PAGE_MASK;
		limit += VM_PAGE_SIZE;
		limit -= 1;
	}

	/* TODO improve this to not require a per-page loop */
	for (virt_addr_t i = base; i < limit;) {
		struct vm_area *area = get_entry(region, i, GET_ENTRY_EXACT);
		if (!area) {
			return false;
		}

		if ((area->vma_prot & prot) != prot) {
			return false;
		}

		i = area->vma_limit;
	}

	return true;
}

static kern_status_t request_missing_page(
	struct address_space *region,
	virt_addr_t addr,
	off_t object_offset,
	struct vm_object *object,
	vm_prot_t prot,
	enum pmap_fault_flags flags,
	unsigned long *irq_flags)
{
	/* here:
	 * `region` is locked.
	 * `object` is unlocked.
	 * `irq_flags` must be restored when `region` is unlocked.
	 * the relevant page in `object` may or may not be committed.
	 * if it isn't, it needs to be requested.
	 */
	vm_object_lock(object);
	address_space_unlock(region);

	struct vm_page *pg = vm_object_get_page(
		object,
		object_offset,
		VMO_ALLOCATE_MISSING_PAGE | VMO_REQUEST_MISSING_PAGE,
		irq_flags);
	if (!pg) {
		vm_object_unlock_irqrestore(object, *irq_flags);
		return KERN_FATAL_ERROR;
	}

	/* now: `region` is unlocked, and `object` is locked */

	kern_status_t status = pmap_add(
		region->s_pmap,
		addr,
		vm_page_get_pfn(pg),
		prot,
		PMAP_NORMAL);

	vm_object_unlock_irqrestore(object, *irq_flags);
	return status;
}

/* this function must be called with `region` locked */
kern_status_t address_space_demand_map(
	struct address_space *region,
	virt_addr_t addr,
	enum pmap_fault_flags flags)
{
	addr &= ~VM_PAGE_MASK;
	if (addr < region->s_base_address || addr > region->s_limit_address) {
		return KERN_NO_ENTRY;
	}

	unsigned long irq_flags;
	address_space_lock_irqsave(region, &irq_flags);

	struct vm_area *area = get_entry(region, addr, GET_ENTRY_EXACT);
	if (!area || !area->vma_object) {
		address_space_unlock_irqrestore(region, irq_flags);
		return KERN_NO_ENTRY;
	}

	off_t object_offset = addr - area->vma_base + area->vma_object_offset;
	if (area->vma_object->vo_ctrl) {
		return request_missing_page(
			region,
			addr,
			object_offset,
			area->vma_object,
			area->vma_prot,
			flags,
			&irq_flags);
	}

#if 0
	tracek("vm: tried to access vm-object %s at offset=%05llx",
	       area->vma_object->vo_name,
	       object_offset);
#endif

	/* simple case: this vm-object is not attached to a controller */
	vm_object_lock(area->vma_object);
	struct vm_page *pg = vm_object_get_page(
		area->vma_object,
		object_offset,
		VMO_ALLOCATE_MISSING_PAGE,
		NULL);
	// tracek("vm: mapping %07llx -> %10llx", vm_page_get_paddr(pg),
	// addr);
	kern_status_t status = pmap_add(
		region->s_pmap,
		addr,
		vm_page_get_pfn(pg),
		area->vma_prot,
		PMAP_NORMAL);

	vm_object_unlock(area->vma_object);
	address_space_unlock_irqrestore(region, irq_flags);
	return status;
}

virt_addr_t address_space_get_base_address(const struct address_space *region)
{
	return region->s_base_address;
}

kern_status_t address_space_read(
	struct address_space *src_region,
	virt_addr_t src_ptr,
	size_t count,
	void *destp,
	size_t *nr_read)
{
	struct vm_iterator src;
	char *dest = destp;

	vm_iterator_begin(
		&src,
		src_region,
		src_ptr,
		VM_PROT_READ | VM_PROT_USER);

	kern_status_t status = KERN_OK;
	size_t r = 0;

	while (r < count && src.it_max) {
		size_t remaining = count - r;
		size_t to_move = MIN(src.it_max, remaining);
		memmove(dest, src.it_buf, to_move);

		status = vm_iterator_seek(&src, to_move);
		if (status != KERN_OK) {
			break;
		}

		r += to_move;
		dest += to_move;
	}

	vm_iterator_finish(&src);

	if (nr_read) {
		*nr_read = r;
	}

	return status;
}

kern_status_t address_space_write(
	struct address_space *dst_region,
	virt_addr_t dst_ptr,
	size_t count,
	const void *srcp,
	size_t *nr_written)
{
	struct vm_iterator dst;
	const char *src = srcp;

	vm_iterator_begin(
		&dst,
		dst_region,
		dst_ptr,
		VM_PROT_WRITE | VM_PROT_USER);

	kern_status_t status = KERN_OK;
	size_t r = 0;

	while (r < count && dst.it_max) {
		size_t remaining = count - r;
		size_t to_move = MIN(dst.it_max, remaining);
		memmove(dst.it_buf, src, to_move);

		status = vm_iterator_seek(&dst, to_move);
		if (status != KERN_OK) {
			break;
		}

		r += to_move;
		src += to_move;
	}

	vm_iterator_finish(&dst);

	if (nr_written) {
		*nr_written = r;
	}

	return status;
}

kern_status_t address_space_memmove(
	struct address_space *dest_region,
	virt_addr_t dest_ptr,
	struct address_space *src_region,
	virt_addr_t src_ptr,
	size_t count,
	size_t *nr_moved)
{
	struct vm_iterator src, dest;
	vm_iterator_begin(
		&src,
		src_region,
		src_ptr,
		VM_PROT_READ | VM_PROT_USER);
	vm_iterator_begin(
		&dest,
		dest_region,
		dest_ptr,
		VM_PROT_WRITE | VM_PROT_USER);

	kern_status_t status = KERN_OK;
	size_t r = 0;

	while (count && src.it_max && dest.it_max) {
		size_t to_move = MIN(MIN(src.it_max, dest.it_max), count);
		memmove(dest.it_buf, src.it_buf, to_move);

		status = vm_iterator_seek(&src, to_move);
		if (status != KERN_OK) {
			break;
		}

		status = vm_iterator_seek(&dest, to_move);
		if (status != KERN_OK) {
			break;
		}

		count -= to_move;
		r += to_move;
	}

	vm_iterator_finish(&src);
	vm_iterator_finish(&dest);

	if (nr_moved) {
		*nr_moved = r;
	}

	return status;
}

extern kern_status_t address_space_memmove_v(
	struct address_space *dest_region,
	size_t dest_offset,
	const kern_iovec_t *dest_vecs,
	size_t nr_dest_vecs,
	struct address_space *src_region,
	size_t src_offset,
	const kern_iovec_t *src_vecs,
	size_t nr_src_vecs,
	size_t bytes_to_move,
	size_t *nr_bytes_moved)
{
	struct iovec_iterator src, dest;
	iovec_iterator_begin_user(&src, src_region, src_vecs, nr_src_vecs);
	iovec_iterator_begin_user(&dest, dest_region, dest_vecs, nr_dest_vecs);

	iovec_iterator_seek(&src, src_offset);
	iovec_iterator_seek(&dest, dest_offset);

	size_t moved = 0;
	while (bytes_to_move && src.it_len && dest.it_len) {
		size_t to_move
			= MIN(MIN(src.it_len, dest.it_len), bytes_to_move);

		kern_status_t status = address_space_memmove(
			dest_region,
			dest.it_base,
			src_region,
			src.it_base,
			to_move,
			NULL);
		if (status != KERN_OK) {
			return status;
		}

		iovec_iterator_seek(&src, to_move);
		iovec_iterator_seek(&dest, to_move);
		bytes_to_move -= to_move;
		moved += to_move;
	}

	if (nr_bytes_moved) {
		*nr_bytes_moved = moved;
	}

	return KERN_OK;
}

#ifdef TRACE
void address_space_dump(struct address_space *region)
{
	struct btree_node *cur = btree_first(&region->s_mappings);
	while (cur) {
		struct vm_area *area
			= BTREE_CONTAINER(struct vm_area, vma_node, cur);

		tracek("+mapping [%zx-%zx] %s",
		       area->vma_base,
		       area->vma_limit,
		       area->vma_object->vo_name);

		cur = btree_next(cur);
	}
}
#endif