#include <kernel/equeue.h>
#include <kernel/sched.h>
#include <kernel/thread.h>
#include <kernel/util.h>
#include <kernel/vm-controller.h>
#include <kernel/vm-object.h>
#include <magenta/signal.h>

#define VM_CONTROLLER_CAST(p)                                                  \
	OBJECT_C_CAST(struct vm_controller, vc_base, &vm_controller_type, p)

BTREE_DEFINE_SIMPLE_GET(
	struct vm_request,
	uint64_t,
	req_node,
	req_id,
	get_request)

static struct object_type vm_controller_type = {
	.ob_name = "vm-controller",
	.ob_size = sizeof(struct vm_controller),
	.ob_header_offset = offsetof(struct vm_controller, vc_base),
};

static struct vm_cache vm_request_cache = {
	.c_name = "vm-request",
	.c_obj_size = sizeof(struct vm_request),
};

kern_status_t vm_controller_type_init(void)
{
	vm_cache_init(&vm_request_cache);
	return object_type_register(&vm_controller_type);
}

struct vm_controller *vm_controller_cast(struct object *obj)
{
	return VM_CONTROLLER_CAST(obj);
}

struct vm_controller *vm_controller_create(void)
{
	struct object *ctrl_object = object_create(&vm_controller_type);
	if (!ctrl_object) {
		return NULL;
	}

	struct vm_controller *ctrl = VM_CONTROLLER_CAST(ctrl_object);

	return ctrl;
}

static struct vm_request *get_next_request(struct vm_controller *ctrl)
{
	struct btree_node *cur = btree_first(&ctrl->vc_requests);
	while (cur) {
		struct vm_request *req
			= BTREE_CONTAINER(struct vm_request, req_node, cur);
		spin_lock(&req->req_lock);
		switch (req->req_status) {
		case VM_REQUEST_PENDING:
			req->req_status = VM_REQUEST_IN_PROGRESS;
			ctrl->vc_requests_waiting--;
			return req;
		case VM_REQUEST_ASYNC:
			btree_delete(&ctrl->vc_requests, &req->req_node);
			ctrl->vc_requests_waiting--;
			return req;
		default:
			break;
		}

		spin_unlock(&req->req_lock);
		cur = btree_next(cur);
	}

	return NULL;
}

static kern_status_t try_enqueue(struct btree *tree, struct vm_request *req)
{
	if (!tree->b_root) {
		tree->b_root = &req->req_node;
		btree_insert_fixup(tree, &req->req_node);
		return true;
	}

	struct btree_node *cur = tree->b_root;
	while (1) {
		struct vm_request *cur_node
			= BTREE_CONTAINER(struct vm_request, req_node, cur);
		struct btree_node *next = NULL;

		if (req->req_id > cur_node->req_id) {
			next = btree_right(cur);

			if (!next) {
				btree_put_right(cur, &req->req_node);
				break;
			}
		} else if (req->req_id < cur_node->req_id) {
			next = btree_left(cur);

			if (!next) {
				btree_put_left(cur, &req->req_node);
				break;
			}
		} else {
			return false;
		}

		cur = next;
	}

	btree_insert_fixup(tree, &req->req_node);
	return true;
}

static kern_status_t send_request_async(
	struct vm_controller *ctrl,
	struct vm_request *req)
{
	fill_random(&req->req_id, sizeof req->req_id);
	while (!try_enqueue(&ctrl->vc_requests, req)) {
		req->req_id++;
	}

	ctrl->vc_requests_waiting++;
	object_assert_signal(
		&ctrl->vc_base,
		VM_CONTROLLER_SIGNAL_REQUEST_RECEIVED);

	return KERN_OK;
}

kern_status_t vm_controller_recv(
	struct vm_controller *ctrl,
	equeue_packet_vm_request_t *out)
{
	struct vm_request *req = NULL;

	req = get_next_request(ctrl);
	if (!req) {
		return KERN_NO_ENTRY;
	}

	if (ctrl->vc_requests_waiting == 0) {
		object_clear_signal(
			&ctrl->vc_base,
			VM_CONTROLLER_SIGNAL_REQUEST_RECEIVED);
	}

	vm_object_lock(req->req_object);
	out->req_id = req->req_id;
	out->req_vmo = req->req_object->vo_key;
	out->req_type = req->req_type;

	switch (req->req_type) {
	case VM_REQUEST_READ:
	case VM_REQUEST_DIRTY:
		out->req_offset = req->req_offset;
		out->req_length = req->req_length;
		break;
	case VM_REQUEST_ATTACH:
		out->req_src_vmo = req->req_object->vo_key;
		break;
	default:
		break;
	}

	vm_object_unlock(req->req_object);
	spin_unlock(&req->req_lock);

	if (req->req_status == VM_REQUEST_ASYNC) {
		put_current_thread(req->req_sender);
		vm_cache_free(&vm_request_cache, req);
	}

	return KERN_OK;
}

kern_status_t vm_controller_recv_async(
	struct vm_controller *ctrl,
	struct equeue *eq,
	equeue_key_t key)
{
	if (ctrl->vc_eq) {
		object_unref(&ctrl->vc_eq->eq_base);
	}

	object_ref(&eq->eq_base);
	ctrl->vc_eq = eq;
	ctrl->vc_eq_key = key;

	return KERN_OK;
}

kern_status_t vm_controller_create_object(
	struct vm_controller *ctrl,
	const char *name,
	size_t name_len,
	equeue_key_t key,
	size_t data_len,
	vm_prot_t prot,
	struct vm_object **out)
{
	struct vm_object *vmo
		= vm_object_create(name, name_len, data_len, prot);
	if (!vmo) {
		return KERN_NO_MEMORY;
	}

	object_ref(&ctrl->vc_base);

	/* TODO expose the VMO_AUTO_DETACH flag to userspace */
	vmo->vo_flags |= VMO_CONTROLLER | VMO_AUTO_DETACH;
	vmo->vo_ctrl = ctrl;
	vmo->vo_key = key;

	*out = vmo;
	return KERN_OK;
}

kern_status_t vm_controller_prepare_attach(
	struct vm_controller *ctrl,
	uint64_t req_id,
	struct vm_object **out_vmo)
{
	struct vm_request *req = get_request(&ctrl->vc_requests, req_id);
	if (!req) {
		return KERN_INVALID_ARGUMENT;
	}

	spin_lock(&req->req_lock);
	req->req_status = VM_REQUEST_IN_PROGRESS;
	*out_vmo = req->req_object;
	spin_unlock(&req->req_lock);

	return KERN_OK;
}

kern_status_t vm_controller_finish_attach(
	struct vm_controller *ctrl,
	uint64_t req_id,
	equeue_key_t new_key)
{
	struct vm_request *req = get_request(&ctrl->vc_requests, req_id);
	if (!req) {
		return KERN_INVALID_ARGUMENT;
	}

	spin_lock(&req->req_lock);
	struct vm_object *vmo = req->req_object;
	spin_unlock(&req->req_lock);

	vm_object_lock(vmo);
	vmo->vo_key = new_key;
	vmo->vo_flags &= ~VMO_LAZY_ATTACH;
	vm_object_unlock(vmo);

	spin_lock(&req->req_lock);
	req->req_status = VM_REQUEST_COMPLETE;
	req->req_result = KERN_OK;
	thread_awaken(req->req_sender);
	spin_unlock(&req->req_lock);

	return KERN_OK;
}

kern_status_t vm_controller_detach_object(
	struct vm_controller *ctrl,
	struct vm_object *vmo)
{
	if (vmo->vo_ctrl != ctrl) {
		return KERN_INVALID_ARGUMENT;
	}

	if (vmo->vo_flags & VMO_LAZY_ATTACH) {
		/* this vmo isn't actually attached to this controller yet.
		 * this can happen if a controller-attached vmo was duplicated
		 * via copy-on-write, and the duplicate vmo has not yet been
		 * accessed. */
		vmo->vo_ctrl = NULL;
		return KERN_OK;
	}

	struct vm_request *req = vm_cache_alloc(&vm_request_cache, VM_NORMAL);
	req->req_type = VM_REQUEST_DETACH;
	req->req_status = VM_REQUEST_ASYNC;
	req->req_object = vmo;
	req->req_sender = get_current_thread();
	send_request_async(ctrl, req);

	vmo->vo_ctrl = NULL;
	object_unref(&ctrl->vc_base);

	return KERN_OK;
}

static void wait_for_reply(
	struct vm_controller *ctrl,
	struct vm_request *req,
	unsigned long *lock_flags)
{
	struct wait_item waiter;
	struct thread *self = get_current_thread();

	wait_item_init(&waiter, self);
	for (;;) {
		self->tr_state = THREAD_SLEEPING;
		if (req->req_status == VM_REQUEST_COMPLETE) {
			break;
		}

		spin_unlock_irqrestore(&req->req_lock, *lock_flags);
		schedule(SCHED_NORMAL);
		spin_lock_irqsave(&req->req_lock, lock_flags);
	}

	self->tr_state = THREAD_READY;
	put_current_thread(self);
}

void vm_controller_fulfill_requests(
	struct vm_controller *ctrl,
	equeue_key_t object,
	off_t offset,
	size_t length,
	kern_status_t result)
{
	off_t limit = offset + length - 1;
	struct btree_node *cur = btree_first(&ctrl->vc_requests);
	while (cur) {
		struct vm_request *req
			= BTREE_CONTAINER(struct vm_request, req_node, cur);
		spin_lock(&req->req_lock);
		bool match = false;
		off_t req_base = req->req_offset;
		off_t req_limit = req->req_offset + req->req_length - 1;

		if (req_base >= offset && req_base <= limit) {
			match = true;
		} else if (req_limit >= offset && req_limit <= limit) {
			match = true;
		}

		vm_object_lock(req->req_object);
		if (req->req_object->vo_key != object) {
			match = false;
		}
		vm_object_unlock(req->req_object);

		if (match) {
			req->req_status = VM_REQUEST_COMPLETE;
			req->req_result = result;
			thread_awaken(req->req_sender);
		}

		spin_unlock(&req->req_lock);
		cur = btree_next(cur);
	}
}

kern_status_t vm_controller_supply_pages(
	struct vm_controller *ctrl,
	struct vm_object *dst,
	off_t dst_offset,
	struct vm_object *src,
	off_t src_offset,
	size_t count)
{
	if (src->vo_flags & VMO_CONTROLLER) {
		return KERN_INVALID_ARGUMENT;
	}

	if (dst->vo_ctrl != ctrl) {
		return KERN_INVALID_ARGUMENT;
	}

	kern_status_t status = vm_object_transfer(
		dst,
		dst_offset,
		src,
		src_offset,
		count,
		NULL);

	return status;
}

kern_status_t vm_controller_send_request(
	struct vm_controller *ctrl,
	struct vm_request *req,
	unsigned long *irq_flags)
{
	fill_random(&req->req_id, sizeof req->req_id);
	while (!try_enqueue(&ctrl->vc_requests, req)) {
		req->req_id++;
	}

	ctrl->vc_requests_waiting++;
	object_assert_signal(
		&ctrl->vc_base,
		VM_CONTROLLER_SIGNAL_REQUEST_RECEIVED);

	vm_controller_unlock(ctrl);
	wait_for_reply(ctrl, req, irq_flags);

	spin_unlock_irqrestore(&req->req_lock, *irq_flags);
	vm_controller_lock_irqsave(ctrl, irq_flags);
	spin_lock(&req->req_lock);

	btree_delete(&ctrl->vc_requests, &req->req_node);

	return KERN_OK;
}