#include "task.h"
#include "sched.h"
#include <core/sys.h>
#include <lib/cpp.h>
#include <mem/seg.h>
#include <mem/gdt.h>
#include "timer.h"

#include <linker/elf.h>

#include <vfs/node.h>
#include <lib/earray.h>
#include <ui/vt.h>

#define KSTACKSIZE 0x8000

//Static routines for handling threads exiting and all cleanup
static void thread_exit_stackJmp(uint32_t reason);
static void thread_exit2(uint32_t reason);

//From task_.asm
extern "C" uint32_t read_eip();
extern "C" void task_idle(void*);

earray<process> *proc_by_pid;

process *kernel_process = 0, *current_process = 0;
thread *current_thread = 0, *idle_thread = 0;

int proc_count = 0;

uint32_t tasking_tmpStack[KSTACKSIZE];

/*	Sets up tasking. Called by kmain on startup.
	Creates a kernel process and an IDLE thread in it. */
void tasking_init() {
	proc_by_pid = new earray<process>(12, 32);

	cli();
	kernel_process = new process();
	kernel_process->uid = kernel_process->thread_count = 0;
	kernel_process->pid = proc_by_pid->add(kernel_process);
	kernel_process->privilege = PL_KERNEL;
	kernel_process->parent = kernel_process;
	kernel_process->pagedir = kernel_pagedir;
	kernel_process->threads = kernel_process->threads_waiting = 0;
	kernel_process->finished = false;
	kernel_process->data = kernel_process->stack = 0;
	current_thread = 0;
	idle_thread = new thread(kernel_process, task_idle, 0, 0);
	sti();
}

/*	Called by the paging functions when a page table is allocated in the kernel space (>K_HIGHHALF_ADDR).
	Updates the page directories of all the processes. */
void tasking_updateKernelPagetable(uint32_t idx, page_table *table, uint32_t tablephysical) {
	if (idx < FIRST_KERNEL_PAGETABLE) return;
	for (int i = 0; i < proc_by_pid->count(); i++) {
		process *it = proc_by_pid->at(i);
		if (it == 0) continue;
		it->pagedir->tables[idx] = table;
		it->pagedir->tablesPhysical[idx] = tablephysical;
	}
}

/*	Called when a timer IRQ fires. Does a context switch. */
void schedule() {
	asm volatile("cli");

	uint32_t esp, ebp, eip;

	asm volatile("mov %%esp, %0" : "=r"(esp));
	asm volatile("mov %%ebp, %0" : "=r"(ebp));
	eip = read_eip();

	if (eip == 0x12345) {
		return;
	}

	if (current_thread != 0) {
		current_thread->esp = esp;
		current_thread->ebp = ebp;
		current_thread->eip = eip;
		if (current_thread->state == TS_RUNNING) sched_enqueue(current_thread);
	}

	current_thread = sched_dequeue();
	ASSERT(current_thread != 0);
	current_process = current_thread->process;

	pagedir_switch(current_process->pagedir);

	gdt_setKernelStack(((uint32_t)current_thread->kernelStack_addr) + current_thread->kernelStack_size);

	asm volatile("			\
			mov %0, %%ebp;	\
			mov %1, %%esp;	\
			mov %2, %%ecx;	\
			mov $0x12345, %%eax;	\
			jmp *%%ecx;"
		: : "r"(current_thread->ebp), "r"(current_thread->esp), "r"(current_thread->eip));
}

/*	Called when an exception happens. Provides a stack trace if it was in kernel land.
	Ends the thread for most exceptions, ends the whole process for page faults. */
uint32_t tasking_handleException(registers *regs) {
	if (current_thread == 0) return 0;	//No tasking yet
	NL; WHERE; ke_vt->writeStr("Exception: `");
	char *exception_messages[] = {"Division By Zero","Debug","Non Maskable Interrupt","Breakpoint",
    "Into Detected Overflow","Out of Bounds","Invalid Opcode","No Coprocessor", "Double Fault",
    "Coprocessor Segment Overrun","Bad TSS","Segment Not Present","Stack Fault","General Protection Fault",
    "Page Fault","Unknown Interrupt","Coprocessor Fault","Alignment Check","Machine Check"};
	*ke_vt << exception_messages[regs->int_no];
	*ke_vt << "'\teip:" << regs->eip;
	if (regs->eip >= K_HIGHHALF_ADDR) {
		*ke_vt << "\nException in kernel. Stack trace:\n";
		stack_trace(regs->ebp);
	}
	if (regs->int_no == 14) {
		*ke_vt << "\n>>> Process exiting.\n";
		thread_exit_stackJmp(EX_PR_EXCEPTION);
	} else {
		*ke_vt << "\n>>> Thread exiting.\n";
		thread_exit_stackJmp(EX_TH_EXCEPTION);
	}
	PANIC("This should never have happened. Please report this.");
	return 0;
}

/*	Puts the current thread in an inactive state. */
void thread_goInactive() {
	current_thread->state = TS_WAKEWAIT;
	schedule();
}

/*	Wakes up the given thread. */
void thread::wakeUp() {
	if (state == TS_WAKEWAIT) {
		state = TS_RUNNING;
		sched_enqueue(this);
	}
}

/*	Returns the privilege level of the current process. */
int proc_priv() {
	if (current_thread == 0 || current_process == 0) return PL_UNKNOWN;
	return current_process->privilege;
}

/*	For internal use only. Called by thread_exit_stackJmp on a stack that will not be deleted.
	Exits current thread or process, depending on the reason. */
void thread_exit2(uint32_t reason) {		//See EX_TH_* defines in task.h
	/*
	 * if reason == EX_TH_NORMAL, it is just one thread exiting because it has finished
	 * if reason == EX_TH_EXCEPTION, it is just one thread exiting because of an exception
	 * if reason is none of the two cases above, it is the whole process exiting (with error code = reason)
	 */
	thread *th = current_thread;
	process* pr;
	if (th == 0 || th->process == 0) goto retrn;
	pr = th->process;
	if ((reason == EX_TH_NORMAL || reason == EX_TH_EXCEPTION) && pr->thread_count > 1) {
		delete th;
	} else {
		pr->finish(reason);
	}
	retrn:
	schedule();
}

/*	For internal use only. Called by thread_exit and process_exit.
	Switches to a stack that will not be deleted when current thread is deleted. */
void thread_exit_stackJmp(uint32_t reason) {
	asm volatile("cli");
	uint32_t *stack;
	stack = tasking_tmpStack + (KSTACKSIZE / 4);
	stack--; *stack = reason;
	stack--; *stack = 0;
	asm volatile("			\
			mov %0, %%esp;	\
			mov %1, %%ebp;	\
			mov %2, %%ecx;	\
			mov %3, %%cr3;	\
			jmp *%%ecx;" : :
			"r"(stack), "r"(stack), "r"(thread_exit2), "r"(kernel_pagedir->physicalAddr));
}

/*	System call. Exit the current thread. */
void thread_exit() {
	thread_exit_stackJmp(EX_TH_NORMAL);
}

/*	System call. Exit the current process. */
void process_exit(size_t retval) {
	if (retval == EX_TH_NORMAL || retval == EX_TH_EXCEPTION) retval = EX_PR_EXCEPTION;
	thread_exit_stackJmp(retval);
}

/*	For internal use only. This is called when a newly created thread first runs
	(its address is the value given for EIP).
	It switches to user mode if necessary and calls the entry point. */
static void thread_run(void* u_esp, thread *thread, thread_entry entry_point, void *data) {
	pagedir_switch(thread->process->pagedir);
	if (thread->process->privilege >= PL_USER) {	//User mode !
		uint32_t *stack = (uint32_t*)u_esp;

		stack--; *stack = (uint32_t)data;
		stack--; *stack = 0;
		size_t esp = (size_t)stack, eip = (size_t)entry_point;
		//Setup a false structure for returning from an interrupt :
		//value for esp is in ebx, for eip is in ecx
		//- update data segments to 0x23 = user data segment with RPL=3
		//- push value for ss : 0x23 (user data seg rpl3)
		//- push value for esp
		//- push flags
		//- update flags, set IF = 1 (interrupts flag)
		//- push value for cs : 0x1B = user code segment with RPL=3
		//- push eip
		//- return from fake interrupt
		asm volatile("				\
				mov $0x23, %%ax;	\
				mov %%ax, %%ds;		\
				mov %%ax, %%es;		\
				mov %%ax, %%fs;		\
				mov %%ax, %%gs;		\
									\
				pushl $0x23;		\
				pushl %%ebx;		\
				pushf;				\
				pop %%eax;			\
				or $0x200, %%eax;	\
				push %%eax;			\
				pushl $0x1B;		\
				push %%ecx;			\
				iret;				\
				" : : "b"(esp), "c"(eip));
	} else {
		asm volatile("sti");
		entry_point(data);
	}
	thread_exit();
}

/*	Creates a new thread for given process.
	Allocates a kernel stack and a user stack if necessary.
	Sets up the kernel stack for values to be passed to thread_run. */
thread::thread(class process *proc, thread_entry entry_point, void *data, void *u_esp) {
	process = proc;
	proc->thread_count++;

	if (u_esp == 0) u_esp = (void*)proc->stack;

	kernelStack_addr = kmalloc(KSTACKSIZE);
	kernelStack_size = KSTACKSIZE;

	uint32_t *stack = (uint32_t*)((size_t)kernelStack_addr + kernelStack_size);

	//Pass parameters
	stack--; *stack = (uint32_t)data;
	stack--; *stack = (uint32_t)entry_point;	
	stack--; *stack = (uint32_t)this;
	stack--; *stack = (uint32_t)u_esp;
	stack--; *stack = 0;
	esp = (uint32_t)stack;
	ebp = esp + 16;
	eip = (uint32_t)thread_run;

	next = proc->threads;
	proc->threads = this;

	state = TS_RUNNING;
	sched_enqueue(this);
}

/*	Creates a new process. Creates a struct process and fills it up. */
process::process(process* _parent, uint32_t _uid, uint32_t _privilege) : fd(12, 64) {
	uid = _uid;
	thread_count = 0;
	threads = 0;
	threads_waiting = 0;
	privilege = _privilege;
	parent = _parent;
	pagedir = pagedir_new();
	data = 0;
	dataseg = 0;
	finished = false;
	retval = E_NOT_FINISHED;

	stack = 0;
	if (privilege >= PL_USER) {	//We are running in user mode
		size_t stacksBottom = K_HIGHHALF_ADDR - 0x01000000;
		(new simpleseg(stacksBottom, USER_STACK_SIZE, 1))->map(pagedir, 0);
		stack = stacksBottom + USER_STACK_SIZE - 4;
	}

	pid = proc_by_pid->add(this);

	proc_count++;
}

/*	Deletes given thread, freeing the stack(s). */
thread::~thread() {
	sched_remove(this);
	if (process->threads == this) {
		process->threads = next;
	} else {
		thread *it = process->threads;
		ASSERT(it != 0);
		while (it->next != 0 && it->next != this) it = it->next;
		ASSERT(it->next = this);
		it->next = next;
	}
	if (current_thread == this) current_thread = 0;
	process->thread_count--;
	kfree(kernelStack_addr);
}

/*	Deletes a process. First, deletes all its threads. Also deletes the corresponding page directory. */
void process::finish(int ret) {
	if (finished) return;
	finished = true;

	retval = ret;

	while (threads != 0) {
		delete threads;
	}
	ASSERT(thread_count == 0);

	pagedir_delete(pagedir);

	thread *it = threads_waiting;
	while (it != 0) {
		it->wakeUp();
		it = it->queue_next;
	}

	proc_count--;
	if (proc_count == 0) PANIC("No more processes are running! This is bad! Or is it?");
}

size_t process::sbrk(ptrdiff_t size) {
	if (data == 0) return -1;
	ASSERT(data < K_HIGHHALF_ADDR);
	if (data + size >= K_HIGHHALF_ADDR) return -1;

	size_t ret;
	if (dataseg == 0) {
		size_t start = data;
		if (start & 0x0FFF) start = (start & 0xFFFFF000) + 0x1000;
		size_t end = start + size;
		if (end & 0x0FFF) end = (end & 0xFFFFF000) + 0x1000;
		
		segment *s = new simpleseg(start, end - start, 1);
		if (s == 0) return -5;
		dataseg = s->map(pagedir, 0);
		if (dataseg == 0) return -1;

		ret = start;
		data = start + size;
	} else {
		size_t start = dataseg->start;
		size_t end = data + size;
		if (end <= start) return -1;

		if (end & 0x0FFF) end = (end & 0xFFFFF000) + 0x1000;
		simpleseg *s = (simpleseg*)(dataseg->seg);
		s->resize(dataseg, end - start);

		ret = data;
		data += size;
	}
	/* (DBG) ke_vt->writeStr("(sbrk ");
	ke_vt->writeHex(size);
	ke_vt->writeStr(" ");
	ke_vt->writeHex(ret);
	ke_vt->writeStr(")"); */
	return ret;
}

void* process::set_args(char** args) {
	if (args == 0) return 0;

	cli();		// we're writing to this process's user space
	page_directory *r = current_pagedir;
	pagedir_switch(pagedir);

	int len = sizeof(char*), argc = 0;
	for (char** a = args; *a != 0; a++) {
		len += strlen(*a) + 1 + sizeof(char*);
		argc++;
	}

	void* p = (void*)sbrk(len);
	char** si = (char**)p;
	char* s = (char*)p + ((argc + 1) * sizeof(char*));
	for (int i = 0; i < argc; i++) {
		si[i] = s;
		strcpy(s, args[i]);
		s += strlen(args[i]) + 1;
	}
	si[argc] = 0;

	pagedir_switch(r);
	sti();

	return p;
}

// =====================  SYSTEM CALLS


size_t process_sbrk(ptrdiff_t size) {
	return current_process->sbrk(size);
}

void process_brk(size_t ptr) {
	ASSERT(ptr < K_HIGHHALF_ADDR);
	current_process->sbrk(ptr - current_process->data);
}

int process_run(char* file, char** args, FILE zero_fd) {
	node* bin = vfs_find(root, file);
	if (bin == 0) return E_NOT_FOUND;
	if ((bin->type & FT_FILE) == 0) return E_INVALID;

	int len = bin->get_size();
	void *v = kmalloc(len);
	int r = bin->read(0, len, (char*)v);
	if (r != len) {
		kfree(v);
		return E_NOT_IMPLEMENTED;
	}
	if (elf_check((uint8_t*)v) != 0) {
		kfree(v);
		return E_NOT_IMPLEMENTED;
	}
	
	// copy arguments to kernel
	char **arg_data;
	int argc = 0;
	if (args != 0) {
		for (char** a = args; *a != 0; a++) argc++;
		arg_data = (char**)kmalloc(sizeof(char*) * (argc + 2));
		for (int i = 0; i < argc; i++) arg_data[i+1] = strdup(args[i]);
	} else {
		arg_data = (char**)kmalloc(sizeof(char*)*2);
	}
	arg_data[0] = strdup(file);
	arg_data[argc + 1] = 0;

	// load
	process *p = elf_exec((uint8_t*)v, PL_USER, arg_data);
	
	kfree(v);
	for (int i = 0; arg_data[i] != 0; i++) kfree(arg_data[i]);
	kfree(arg_data);
	if (p == 0) {
		return E_INVALID;
	} else {
		p->fd.set(0, current_process->fd.at((zero_fd < 0 ? 0 : zero_fd)));
		return p->pid;
	}
}

int process_waitpid(int pid, int block) {
	if (pid <= 0) return E_INVALID_RANGE;

	process *p = proc_by_pid->at(pid);
	if (p == 0) return E_NOT_FOUND;

	asm volatile("cli");
	int ret = p->retval;

	if (!p->finished) {
		if (block == 0) return E_NOT_FINISHED;

		current_thread->queue_next = p->threads_waiting;
		p->threads_waiting = current_thread;
		thread_goInactive();

		asm volatile("cli");
		ASSERT(p->finished);
		ret = p->retval;

		if (p->threads_waiting == current_thread) {
			p->threads_waiting = current_thread->queue_next;
		} else {
			thread *it = p->threads_waiting;
			ASSERT(it != 0);
			while (it->queue_next != 0 && it->queue_next != current_thread) it = it->queue_next;
			ASSERT(it->queue_next == current_thread);
			it->queue_next = current_thread->queue_next;
		}
		current_thread->queue_next = 0;
	}

	if (p->threads_waiting == 0 && p->parent == current_process) {
		proc_by_pid->set(pid, 0);
		delete p;
	}

	return ret;
}