#include "mem.h"

#include "communication_interface.h"
#include "task.h"
#include "scheduler.h"

extern struct MTask *current_task;
extern struct TaskTree *events;

static uint16_t memsize
	= MEMSIZE - 1 - (MEMSIZE - 1) % sizeof (struct TaskTree);
static void *mtmem_real[MEMSIZE/sizeof(void *)] = {NULL};
static uint8_t *mtmem = (uint8_t *)&mtmem_real;

//Variable pointers
uint16_t mem_task;
static uint16_t mem_heap;

void *mem_ptr(uint16_t i)
{
	return i == MT_NULL ? NULL : (void *)(mtmem + i);
}

uint16_t mem_rptr(void *p)
{
	return p == NULL ? MT_NULL : ((uint8_t *)p)-mtmem;
}

void mem_reset(void)
{
	mem_task = 0;
	mem_heap = memsize;
	msg_log(SC("sizeof(struct MTMessageTo)  : %02x\n"),
		sizeof(struct MTMessageTo));
	msg_log(SC("sizeof(struct MTMessageFrom): %02x\n"),
		sizeof(struct MTMessageFro));
	msg_log(SC("sizeof(struct MTask)        : %02x\n"),
		sizeof(struct MTask));
	msg_log(SC("sizeof(struct TaskTree)     : %02x\n"),
		sizeof(struct TaskTree));
	msg_log(SC("mem_top: %p\n"), mem_ptr(memsize));
	msg_log(SC("mem_bottom: %p\n"), mem_ptr(0));
	msg_log(SC("free bytes: %llu\n"), (size_t)mem_heap - (size_t)mem_task);
}

uint16_t *mem_stack(void)
{
	//This is aligned because tasks are always PTRSIZE bytes aligned
	return mem_ptr(mem_task);
}

uint16_t *mem_max_sp(void)
{
	return mem_ptr(mem_heap);
}

void mem_set_stack_ptr(void *ptr)
{
	mem_task = mem_rptr(ptr);
}

void *mem_alloc_task(size_t size)
{
	msg_debug(SC("mem_alloc_task(%lu) mem_task: %p (%u)\n"),
		size, mem_ptr(mem_task), mem_task);

	if (size == 0)
		return mem_ptr(mem_task);

#ifdef REQUIRE_ALIGNED_MEMORY_ACCESS
	mem_task += PALIGND(size, mem_ptr(mem_task));
#endif /* REQUIRE_ALIGNED_MEMORY_ACCESS */
	void *r = mem_ptr(mem_task);

	msg_debug(SC("mem_alloc_task(%lu) mem_task: %p (%u)\n"),
		size, r, mem_task);

	if (mem_task + size > mem_heap) {
		send_message((struct MTMessageFro)
			{ .cons=MTFException_c
			, .data={.MTFException=
				{ .cons = MTEOutOfMemory_c
				, .data = {.MTEOutOfMemory=0}
				}
			}});
		return NULL;
	}
	mem_task += size;
	return r;
}

struct MTask *mem_task_head()
{
	return mem_task == 0 ? NULL : mem_cast_task(0);
}

struct MTask *mem_task_next_unsafe(struct MTask *t)
{
	return (struct MTask *)\
		(t->instructions.elements+t->instructions.size);
}
struct MTask *mem_task_next(struct MTask *t)
{
	struct MTask *r = mem_task_next_unsafe(t);
	if (mem_rptr(r) >= mem_task)
		r = NULL;
	return r;
}

struct TaskTree *mem_alloc_tree()
{
	msg_debug(SC("general alloc\n"));
	if (mem_heap - sizeof(struct TaskTree) < mem_task) {
		msg_log(SC("Not enough heap\n"));
		send_message((struct MTMessageFro)
			{ .cons=MTFException_c
			, .data={.MTFException=
				{.cons=MTEHeapUnderflow_c
				,.data={.MTEHeapUnderflow=current_task->taskid}
				}
			}});
		return NULL;
	}
	msg_debug(SC("at: %p\n"), mem_cast_tree(mem_heap));

	return (struct TaskTree *)mem_ptr(mem_heap -= sizeof (struct TaskTree));
}

static void mem_print_heap()
{
#if LOGLEVEL == 2
	uint16_t walker;
	walker = memsize - sizeof(struct TaskTree);
	while (walker >= mem_heap) {
		struct TaskTree *t = mem_cast_tree(walker);
		msg_debug(SC("%d (rptr: %p, trash: %d, refr: [%u, %u]): "),
			walker, t->ptr, t->trash,
			t->refresh_min * (t->seconds ? 1000u : 1u),
			t->refresh_max * (t->seconds ? 1000u : 1u));
		tasktree_print_node(t);
		walker -= sizeof(struct TaskTree);
		msg_debug(SC("\n"));
	}
#endif /* LOGLEVEL == 2 */
}

static void mem_print_task_heap()
{
#if LOGLEVEL == 2
	msg_debug(SC("TASKS\n"));
	struct MTask *t = mem_task_head();
	while (t != NULL) {
		task_print(t, true);
		t = mem_task_next(t);
	}
	msg_debug(SC("ENDTASKS\n"));
#endif /* LOGLEVEL == 2 */
}

//* Destructively update the pointer
static void mem_update_ptr(
	struct TaskTree *t, struct TaskTree *old, struct TaskTree *new)
{
	switch (t->task_type) {
	case BCStableNode_c:
		t->data.stablenode.next = new;
		break;
	case BCUnstableNode_c:
		t->data.unstablenode.next = new;
		break;
	case BCStep_c:
		t->data.step.lhs = new;
		break;
	case BCStepStable_c:
		t->data.steps.lhs = new;
		break;
	case BCStepUnstable_c:
		t->data.stepu.lhs = new;
		break;
	case BCSeqStable_c:
		t->data.seqs.lhs =
			t->data.seqs.lhs == old ? new : t->data.seqs.lhs;
		t->data.seqs.rhs =
			t->data.seqs.rhs == old ? new : t->data.seqs.rhs;
		break;
	case BCSeqUnstable_c:
		t->data.sequ.lhs =
			t->data.sequ.lhs == old ? new : t->data.sequ.lhs;
		t->data.sequ.rhs =
			t->data.sequ.rhs == old ? new : t->data.sequ.rhs;
		break;
	case BCRepeat_c:
		t->data.repeat.tree =
			t->data.repeat.tree == old ? new : t->data.repeat.tree;
		t->data.repeat.oldtree =
			t->data.repeat.oldtree == old
			? new : t->data.repeat.oldtree;
		break;
	case BCTOr_c:
		t->data.tor.lhs =
			t->data.tor.lhs == old ? new : t->data.tor.lhs;
		t->data.tor.rhs =
			t->data.tor.rhs == old ? new : t->data.tor.rhs;
		break;
	case BCTAnd_c:
		t->data.tand.lhs =
			t->data.tand.lhs == old ? new : t->data.tand.lhs;
		t->data.tand.rhs =
			t->data.tand.rhs == old ? new : t->data.tand.rhs;
		break;
	case BCSdsSet_c:
		t->data.sdsset.data = new;
		break;
	case BCRateLimit_c:
		t->data.ratelimit.task =
			t->data.ratelimit.task == old ? new : t->data.ratelimit.task;
		t->data.ratelimit.storage =
			t->data.ratelimit.storage == old ? new : t->data.ratelimit.storage;
		break;
	case BCSdsUpd_c:
		t->data.sdsupd.ctx = new;
		break;
	case BCEvent_c:
		if (t == NULL) {
			events = new;
		} else {
			t->data.event.next = new;
		}
		break;
	default:
		break;
	}
}

void mem_mark_trash_and_destroy(struct TaskTree *tt, struct MTask *t,
				uint16_t *stack)
{
	*stack++ = MT_NULL;
	*stack++ = mem_rptr(tt);
	while (*--stack != MT_NULL) {
		struct TaskTree *ctree = mem_cast_tree(*stack);
		msg_debug(SC("mem_mark_trash_and_destroy: %u "), *stack);
		tasktree_print(ctree, 0);
		msg_debug(SC("\n"));
		switch (ctree->task_type) {
			case BCStableNode_c:
				*stack++ =
					mem_rptr(ctree->data.stablenode.next);
				break;
			case BCUnstableNode_c:
				*stack++ =
					mem_rptr(ctree->data.unstablenode.next);
				break;
			case BCStep_c:
				*stack++ = mem_rptr(ctree->data.step.lhs);
				break;
			case BCStepStable_c:
				*stack++ = mem_rptr(ctree->data.steps.lhs);
				break;
			case BCStepUnstable_c:
				*stack++ = mem_rptr(ctree->data.stepu.lhs);
				break;
			case BCSeqStable_c:
				*stack++ = mem_rptr(ctree->data.seqs.lhs);
				*stack++ = mem_rptr(ctree->data.seqs.rhs);
				break;
			case BCSeqUnstable_c:
				*stack++ = mem_rptr(ctree->data.sequ.lhs);
				*stack++ = mem_rptr(ctree->data.sequ.rhs);
				break;
			case BCRepeat_c:
				*stack++ =
					mem_rptr(ctree->data.repeat.oldtree);
				if (ctree->data.repeat.tree != NULL)
					*stack++ = mem_rptr(
						ctree->data.repeat.tree);
				break;
			case BCTOr_c:
				*stack++ = mem_rptr(ctree->data.tor.lhs);
				*stack++ = mem_rptr(ctree->data.tor.rhs);
				break;
			case BCTAnd_c:
				*stack++ = mem_rptr(ctree->data.tand.lhs);
				*stack++ = mem_rptr(ctree->data.tand.rhs);
				break;
			case BCSdsSet_c:
				*stack++ = mem_rptr(ctree->data.sdsset.data);
				break;
			case BCRateLimit_c:
				*stack++ = mem_rptr(ctree->data.ratelimit.task);
				if (ctree->data.ratelimit.storage != NULL)
					*stack++ = mem_rptr(
						ctree->data.ratelimit.storage);
				break;
			case BCSdsUpd_c:
				*stack++ = mem_rptr(ctree->data.sdsupd.ctx);
				break;
			default:
				break;
		}
		ctree->trash = true;
		destroy_task_node(ctree, t);
	}
}

void mem_node_move(struct TaskTree *t, struct TaskTree *f)
{
	msg_debug(SC("move %lu - %lu\n"), mem_rptr(t), mem_rptr(f));
	mem_print_heap();
	msg_debug(SC("----\n"));
	//Copy the data
	t->task_type = f->task_type;
	t->trash = 0;
	t->seconds = f->seconds;
	t->refresh_min = f->refresh_min;
	t->refresh_max = f->refresh_max;
	t->data = f->data;
	//We leave the ptr in tact
	//t->ptr = f->ptr;

	//Set the rptrs
	switch (t->task_type) {
	case BCStableNode_c:
		t->data.stablenode.next->ptr = t;
		break;
	case BCUnstableNode_c:
		t->data.unstablenode.next->ptr = t;
		break;
	case BCStep_c:
		t->data.step.lhs->ptr = t;
		break;
	case BCStepStable_c:
		t->data.steps.lhs->ptr = t;
		break;
	case BCStepUnstable_c:
		t->data.stepu.lhs->ptr = t;
		break;
	case BCSeqStable_c:
		t->data.seqs.lhs->ptr = t;
		t->data.seqs.rhs->ptr = t;
		break;
	case BCSeqUnstable_c:
		t->data.sequ.lhs->ptr = t;
		t->data.sequ.rhs->ptr = t;
		break;
	case BCRepeat_c:
		if (t->data.repeat.tree != NULL)
			t->data.repeat.tree->ptr = t;
		t->data.repeat.oldtree->ptr = t;
		break;
	case BCTOr_c:
		t->data.tor.lhs->ptr = t;
		t->data.tor.rhs->ptr = t;
		break;
	case BCTAnd_c:
		t->data.tand.lhs->ptr = t;
		t->data.tand.rhs->ptr = t;
		break;
	case BCSdsSet_c:
		t->data.sdsset.data->ptr = t;
		break;
	case BCRateLimit_c:
		t->data.ratelimit.task->ptr = t;
		if (t->data.ratelimit.storage != NULL)
			t->data.ratelimit.storage->ptr = t;
		break;
	case BCSdsUpd_c:
		t->data.sdsupd.ctx->ptr = t;
		break;
	case BCEvent_c:
		if (t->data.event.next != NULL) {
			t->data.event.next->ptr = t;
		}
		events = events == f ? t : events;
		break;
	default:
		break;
	}

	//Mark trash (non recursively)
	if (mem_rptr(f) >= mem_heap) {
		f->trash = true;
		f->ptr = NULL;
	}
	msg_debug(SC("done moving\n"));
	mem_print_heap();
}

#define task_size(t) ((uintptr_t)mem_task_next_unsafe(t)-(uintptr_t)t)
void mem_gc_tasks(void)
{
	mem_print_task_heap();
	msg_debug(SC("MEMGCTASKS\n"));
	//Task size, location, bytes to move
	uint16_t i, tomove;
	struct MTask *t = mem_task_head();
	while (t != NULL) {
		msg_debug(SC("Task %u, at %p (%u)\n"),
			t->taskid, t, mem_rptr(t));
		if (t->stability == MTRemoved_c) {
			msg_debug(SC("Task is removed, mem_task: %u, %p\n"),
				mem_task, t);
			msg_debug(SC("next task: %p\n"),
				mem_task_next_unsafe(t));
			//Size of the task

			uintptr_t tsize = task_size(t);
			msg_debug(SC("Size: %llu\n"), tsize);

			//Location of the task
			i = mem_rptr(t);
			msg_debug(SC("task loc: %p %u\n"), t, i);

			//Where to stop
			tomove = i;
			for (struct MTask *tt = mem_task_next(t); tt != NULL;
					tt = mem_task_next(tt)) {
				struct TaskTree * task_tree = (struct TaskTree *)tt->tree;
				task_tree->ptr = (struct TaskTree *)((uintptr_t) task_tree->ptr - tsize);
				queue_update_pointer(tt, (struct MTask *)((uintptr_t) tt - tsize));
				tomove += task_size(tt);
			}
			msg_debug(SC("move up to %lu\n"), tomove);

			//Move the bytes
			for (; i<tomove; i++) {
//				msg_log(SC("move mtmem[%u] = mtmem[%u+%u]\n"), i, i, tsize);
				mtmem[i] = mtmem[i+tsize];
			}

			//Decrement the next task pointer
			mem_task -= tsize;
			msg_debug(SC("task removed, %u bytes cleared\n"),
				tsize);
			msg_debug(SC("t: %lu\n"), mem_rptr(t));
			msg_debug(SC("new mem_task: %lu\n"), mem_task);
			if (tomove == 0)
				t = NULL;
		} else {
			t = mem_task_next(t);
		}
	}
	mem_print_task_heap();
}

void mem_gc(void)
{
	mem_print_heap();

	uint16_t walker, hole;
	walker = hole = memsize - sizeof(struct TaskTree);
	while (walker >= mem_heap) {
		//Is trash, pass over but keep hole pointer valid
		if (mem_cast_tree(walker)->trash) {
			walker -= sizeof(struct TaskTree);
		//Not trash and no hole
		} else if (walker == hole) {
			hole = walker -= sizeof(struct TaskTree);
		//Not trash and a hole
		} else {
			msg_debug(SC("no trash and hole, move\n"));
			tasktree_print_node(mem_cast_tree(walker));
			msg_debug(SC("\nwalker: %u - hole: %u\n"),
				walker, hole);

			struct TaskTree *t = mem_cast_tree(hole);
			//Move and correct rptrs
			t->ptr = mem_cast_tree(walker)->ptr;
			mem_node_move(t, mem_cast_tree(walker));

			//Update parent
			if (t->task_type == BCEvent_c && t->ptr == NULL) {
				events = mem_ptr(hole);
			} else if (mem_rptr(t->ptr) < mem_heap) {
				msg_debug(SC("Parent is a task\n"));
				((struct MTask *)t->ptr)->tree = mem_ptr(hole);
			} else {
				msg_debug(SC("Parent is a tasktree\n"));
				mem_update_ptr(mem_cast_tree(hole)->ptr,
					mem_ptr(walker), mem_ptr(hole));
			}
			hole -= sizeof(struct TaskTree);
			walker -= sizeof(struct TaskTree);
		}
	}
	if (walker != hole) {
		msg_debug(SC("mem_heap: %u, walker: %u - hole: %u\n"),
			mem_heap, walker, hole);
		msg_debug(SC("Compacted %u bytes\n"), hole - walker);
		mem_heap = hole+sizeof(struct TaskTree);
		msg_debug(SC("mem_heap: %u, walker: %u - hole: %u\n"),
			mem_heap, walker, hole);
	}
	mem_print_heap();
}
