#include "scheduler.h"
#include "kernel.h"
#include "gdt.h"
#include "log.h"
#include "compositor.h"

#include "smp.h"
#include "mem.h"
#include "fs/elf.h"
#include "asm_x86.h"
#include "asm_task.h"
#include "asm_usermode.h"
#include "kmalloc.h"
#include "e1000.h"

#include "vmem.h"
#include "spinlock.h"
#include "ringbuffer.h"
#include "keyboard.h"
#include "mouse.h"
#include "syscalls.h"
#include "fs/ext2.h"
#include "lib/string/string.h"

#define NO_TASK 0xffffffff

static volatile uint32_t pid=0; // holds next pid, to be assigned to a process.

static uint32_t nextPID()
{
    spinlock_spin(SPINLOCK_PID);
    uint32_t ret=pid;
    pid++;
    spinlock_release(SPINLOCK_PID);
    return ret;
}

// we hold this stuff per cpu
static volatile uint32_t current_task[SMP_MAX_PROC];
static volatile uint32_t last_task[SMP_MAX_PROC]; // last non ring1 task

// we hold this stuff per cpu
static volatile struct task_list_struct
{
    volatile bool     active;  // is this slot used (Y/N)
    volatile uint32_t pid;     // process id 

    volatile uint32_t parent;  // parent process id
    volatile uint32_t esp;     // stack pointer of the task
    volatile uint32_t esp0;    // tss.esp0
    volatile struct   pdirectory *vmem; // number of virtual memory table

    volatile uint32_t brk;     // memory brk pos

    volatile bool     try;     // try to process syscall
    volatile bool     syscall; // syscall not processed yet
    volatile uint32_t eax;     // syscall details
    volatile uint32_t ebx;
    volatile uint32_t ecx;
    volatile uint32_t edx;

    volatile bool thread; // is this a light thread?

    volatile char name[255];
    
}task_list[SMP_MAX_PROC][MAX_TASKS];

// init tasks //
volatile void scheduler_init(uint32_t cpu, void *dir)
{
    for(int i=0;i<MAX_TASKS;i++)
    {
	task_list[cpu][i].active=false;
    }

    current_task[cpu]=0;
//  last_task[cpu]=0;
   
    // need to make space on the esp stacks for pushing vals vias task_pusha

    // this is our main kernel task at slot 0 (per cpu)
    task_list[cpu][0].parent=0;
    task_list[cpu][0].pid=nextPID();
    task_list[cpu][0].active=true;
    task_list[cpu][0].syscall=false;
    task_list[cpu][0].thread=false;
    task_list[cpu][0].vmem=dir;
    task_list[cpu][0].esp = kballoc(4)+4*4096-0x200-8; // 4 pages stack & prealign
    task_list[cpu][0].esp0 = 0; // esp0 not required by kernel space tasks
    strcpy(task_list[cpu][0].name,"kernel worker");
    //fd_init_std_streams(task_list[cpu][0].pid,0);

    // this will go to userspace 
    task_list[cpu][1].parent=0;
    task_list[cpu][1].pid=nextPID();
    task_list[cpu][1].active=true;
    task_list[cpu][1].thread=false;
    task_list[cpu][1].syscall=false;
    task_list[cpu][1].vmem=dir;
    task_list[cpu][1].esp  = kballoc(4)+4*4096-0x200-8; // 4 pages stack & prealign
    task_list[cpu][1].esp  = VMEM_USER_STACK_TOP-0x200-8; 
    task_list[cpu][1].esp0 = kballoc(4)+4*4096;  // esp0 needed by user space tasks
    strcpy(task_list[cpu][1].name,"userspace init");
    fd_init_std_streams(task_list[cpu][1].pid);

    // sleeper
    task_list[cpu][2].parent=0;
    task_list[cpu][2].pid=nextPID();
    task_list[cpu][2].active=true;
    task_list[cpu][2].thread=false;
    task_list[cpu][2].syscall=false;
    task_list[cpu][2].vmem=dir;
    task_list[cpu][2].esp = kballoc(4)+4*4096-0x200-8; // 4 pages stack & prealign
    task_list[cpu][2].esp0 =0;  // esp0 not needed by kernel space tasks
    strcpy(task_list[cpu][2].name,"idle process");
    //fd_init_std_streams(task_list[cpu][2].pid,0);

    // stacks
    task_pusha(task_list[cpu][0].esp);
    task_pusha(task_list[cpu][1].esp);
    task_pusha(task_list[cpu][2].esp);
}

static uint32_t scheduler_schedule(uint32_t idx)
{

	uint32_t cpu=smp_get(SMP_APIC_ID);

	if(task_list[cpu][idx].active && !task_list[cpu][idx].syscall)
	{
            //if(current_task[cpu]!=0)last_task[cpu]=current_task[cpu];
            //if(current_task[cpu]==idx)return task_list[cpu][idx].esp;

	    current_task[cpu]=idx;
            
             //klog("idx %d",idx);
             //klog("name: %s",task_list[cpu][idx].name);

             //klog("cpu %d / idx %d / pid %d / name: %5s",cpu,idx,task_list[cpu][idx].pid,task_list[cpu][idx].name);
             //klog("%d idx %d",last_task[cpu],current_task[cpu]);

	    install_tss(cpu,task_list[cpu][idx].esp0);
	    x86_set_page_directory(task_list[cpu][idx].vmem);
	    return task_list[cpu][idx].esp;
	}
	return 0;
}

//
// REMEMBER WE ARE INSIDE AN INTERRUPT HERE - DON'T WASTE TIME!
//
// oldesp - is the adress of the stack pointer when pit_interrupt_handler was entered.
// registers have been pushed with pusha to this old stack.
//
// stack pointer was moved to the 16kb stack we have from multiboot.s
//
// we need to return a NEW stack pointer where popa will get the registers the new task requires
//

volatile uint32_t scheduler_run(uint32_t oldesp,uint32_t preference)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t init=smp_get(SMP_SCHEDULER_INIT);

    if(init)
    {
	scheduler_init(cpu,x86_get_page_directory());
        smp_set(SMP_SCHEDULER_INIT,0);
        klog("Scheduler initialized for cpu %d",cpu);
        return task_list[cpu][current_task[cpu]].esp;
    }
    else task_list[cpu][current_task[cpu]].esp=oldesp;

    for(int i=0;i<MAX_TASKS;i++)
    {
	int idx=(current_task[cpu]+1+i)%MAX_TASKS; // schedule round robin style
	if(idx==2)continue;// skip sleeper here
	uint32_t esp=scheduler_schedule(idx);
	if(esp)return esp;
    }

    // force the sleeper task here ...
    return scheduler_schedule(2);
}


void scheduler_func()
{
    // we need enable here again (since the pushed eflags have it disabled)? TODO: why they disabled it!???
    x86_sti(); 

    uint32_t cpu=smp_get(SMP_APIC_ID);

    fixme("this will dadlock on context switch during log if never switched back before finish");

    if(current_task[cpu]==0)
    {
        /*
        while(1)
        {
            uint64_t t0=x86_rdtscp();
            asm("hlt"); // sleeper task
            uint64_t t1=x86_rdtscp();
            klog("task 0 / slept cycles: l:%d h:%d",(t1-t0));
        }
        */

        while(1)
        {
        if(cpu==0)
        {
         task_syscall_worker();
        }
        else
        {
            asm("hlt"); // sleeper task
        }
        }
        //task_list[cpu][0].syscall=true; // sleep
	//__asm__("int $0x81"); // wake scheduler! with IPI
        
    }

    if(current_task[cpu]==1)
    {
        if(cpu==0)
        {
            uint32_t alloc;
            uint32_t entry_global=load_elf(BIN_INIT,&alloc);
//            uint32_t entry_global=load_elf("/bin/vesa",&alloc);
            task_set_brk(task_get_current_pid(),alloc);
            asm_usermode(entry_global); 
            kpanic("init died on cpu %d",cpu);
        }

        while(1)
        {
            uint64_t t0=x86_rdtscp();
            asm("hlt"); // sleeper task
            uint64_t t1=x86_rdtscp();
//            klog("task 1 / slept cycles: l:%d h:%d",(t1-t0));
        }
    }

    if(current_task[cpu]==2)
    {
        while(1)
        {
            uint64_t t0=x86_rdtscp();
            asm("hlt"); // sleeper task
            uint64_t t1=x86_rdtscp();
 //           klog("task 2 (sleeper) / slept cycles: l:%d h:%d",(t1-t0));
        }
    }

    kpanic("unknwon task");

}

volatile int add_task(uint32_t parent_pid,uint32_t vmem, bool thread, char *name)
{
    uint32_t parent=task_runs(parent_pid);
    uint32_t cpu=smp_get(SMP_APIC_ID);

    for(int i=0;i<MAX_TASKS;i++)
    {
	if(task_list[cpu][i].active!=true) // find a free slot.
	{
	    task_list[cpu][i].pid=nextPID();
	    task_list[cpu][i].parent=parent_pid;
	    task_list[cpu][i].thread=thread;

	    task_list[cpu][i].vmem=vmem; 

	    task_list[cpu][i].esp = kballoc(4)+2*4096; // center
            task_list[cpu][i].esp0 = kballoc(4)+4*4096;

	    task_list[cpu][i].active=true; //TODO: LOCK! (also other similar)
	    task_list[cpu][i].syscall=false;
	    task_list[cpu][i].try=false;

	    task_list[cpu][i].brk=task_list[cpu][parent].brk;

            uint32_t *source=(uint32_t *)task_list[cpu][parent].esp;
	    uint32_t *dst=(uint32_t *)task_list[cpu][i].esp;

	    for(int x=0;x<100;x++) //TODO:  maybe better copy this page too instead of stack
	    {
		*dst=*source;
		dst++;
		source++;
	    }

	    uint32_t *stack=task_list[cpu][i].esp;
	    stack[12]=0x1; 
	    stack[13]=0; // this task returns pid=0 to the caller
    
            strcpy(task_list[cpu][i].name,name);

	    return task_list[cpu][i].pid;
	}
    }
    kpanic("out of task slots!");
}

uint32_t scheduler_wake_worker(uint32_t oldesp)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    task_list[cpu][0].syscall=false; // wake (syscall misused)
    return 0;
//    return scheduler_run(oldesp,0);
}

void scheduler_wake_all()
{
	    uint32_t cpu=smp_get(SMP_APIC_ID);
		// simple approach, any syscall might unblock any other syscall // TODO: better! TODO: all cpus!
		for(int i=0;i<MAX_TASKS;i++)
		{
			    task_list[cpu][i].try=true;
		}
}

/**
 * kernel space worker thread
 *
 * we can get interrupted by an interrupt ANYTIME!
 *
 */
void task_syscall_worker()
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    //task_list[cpu][0].syscall=true; // sleep (syscall misused)
    //return;
    /// TODO: cross check all cpus!

    while(1)
    {
	bool wake=false;
	bool wake_mouse=false;

        //TODO: would be enough only to lock during ringbuffer acces!?
        //TODO: no locking should be required at all!!!
    	//x86_cli(); // disable temporarily mouse/kb/timer interrupts. 
	wake|=keyboard_worker();
        wake_mouse|=mouse_worker();
	//x86_sti();

        if(wake_mouse)compositor_wake();
            compositor_swap_buffers();
	if(wake)scheduler_wake_all();

        //if(cpu==0)compositor_swap_buffers();

	for(int i=0;i<MAX_TASKS;i++)
	{
	    if(task_list[cpu][i].active,task_list[cpu][i].try&&task_list[cpu][i].syscall)
	    {

		uint32_t syscall=task_list[cpu][i].eax;
                //klog("task pid=%d waiting on syscall %d/%s on cpu %d slot %d.",task_list[cpu][i].pid,syscall,syscall_get_name(syscall),cpu,i);

		task_list[cpu][0].vmem=task_list[cpu][i].vmem; // switch syscall worker to pagedir of calling userprog
		x86_set_page_directory(task_list[cpu][0].vmem);

                x86_cli(); // we do not want to be scheduled to other page ... sorry TODO: suuuuckS!
		uint32_t ok = syscall_generic_test(task_list[cpu][i].eax,
			    task_list[cpu][i].edx,
			    task_list[cpu][i].ecx,
			    task_list[cpu][i].ebx,
			    task_list[cpu][i].pid);
                x86_sti();

		if(!ok)
		{
		    task_list[cpu][i].try=false;
		    continue;
		}

                x86_cli(); // we do not want to be scheduled to other page ... sorry TODO: suuuuckS!
		uint32_t ret = syscall_generic(task_list[cpu][i].eax,
			    task_list[cpu][i].edx,
			    task_list[cpu][i].ecx,
			    task_list[cpu][i].ebx,
			    task_list[cpu][i].pid);
                x86_sti();
		
//		klog("... returned : %d",ret);

		scheduler_wake_all();

		uint32_t *stack=task_list[cpu][i].esp;
		stack[12]=0x1; 
		stack[13]=ret; 
	    
		task_list[cpu][i].syscall=false;
	    }
	}

	bool nowork=true;

	x86_cli(); // disable all incoming new work (mouse,kb,timer). software syscalls can not appear since we are here ;)
	for(int i=0;i<MAX_TASKS;i++) /// ... and check if aaany work is pending.
	{
	    if(task_list[cpu][i].active,task_list[cpu][i].try&&task_list[cpu][i].syscall)
	    {
		nowork=false;
		break;
	    }
	}
	if(nowork)
	{
	    task_list[cpu][0].syscall=true; // sleep (syscall misused)
	}
	x86_sti();

	if(nowork)
	{
	    __asm__("int $0x81"); // wake scheduler! with IPI
	}
    }
}

// !!! REMEMBER THIS IS INSIDE AN INTERRUPT !!!
volatile uint32_t task_syscall(uint32_t eax,uint32_t ebx, uint32_t ecx, uint32_t edx)
{
    if(eax==2)while(1);
    uint32_t cpu=smp_get(SMP_APIC_ID);
    task_list[cpu][current_task[cpu]].syscall=true;
    task_list[cpu][current_task[cpu]].try=true;
    task_list[cpu][current_task[cpu]].eax=eax;
    task_list[cpu][current_task[cpu]].ebx=ebx;
    task_list[cpu][current_task[cpu]].ecx=ecx;
    task_list[cpu][current_task[cpu]].edx=edx;

    // chance for some preparations //
    int i=current_task[cpu];
    syscall_generic_prep(task_list[cpu][i].eax,
                        task_list[cpu][i].edx,
                        task_list[cpu][i].ecx,
                        task_list[cpu][i].ebx,
                        task_list[cpu][i].pid);
    return 1;
}

volatile uint32_t task_fork(uint32_t pid)
{ 
   uint32_t idx=task_runs(pid);
   uint32_t cpu=smp_get(SMP_APIC_ID);
   char *name=task_list[cpu][idx].name;
   int ret=add_task(pid,vmem_new_space_dir(task_list[cpu][idx].vmem,false),false,name);
   klog("[%d] forked -> [%d]", pid, ret);
   return ret;
}

void task_set_name(uint32_t pid, char *name)
{
   uint32_t idx=task_runs(pid);
   uint32_t cpu=smp_get(SMP_APIC_ID);
   strcpy(task_list[cpu][idx].name,name);
}

volatile uint32_t task_clone(uint32_t pid)
{ 
   uint32_t idx=task_runs(pid);
   uint32_t cpu=smp_get(SMP_APIC_ID);
   char *name=task_list[cpu][idx].name;
   int ret=add_task(pid,vmem_new_space_dir(task_list[cpu][idx].vmem,true),true,name);
   klog("[%d] cloned -> [%d]", pid, ret);
   return ret;
}

volatile uint32_t task_get_brk(uint32_t pid)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t idx=task_idx(pid);
    return task_list[cpu][idx].brk;
}

volatile void task_set_brk(uint32_t pid, uint32_t brk)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t idx=task_idx(pid);
    task_list[cpu][idx].brk=brk;
}

volatile uint32_t task_get_current_pid()
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    return task_list[cpu][current_task[cpu]].pid;
}

volatile uint32_t task_get_parent(uint32_t pid)
{
   uint32_t idx=task_idx(pid);
   uint32_t cpu=smp_get(SMP_APIC_ID);
   return task_list[cpu][idx].parent;
}

volatile int task_add_win(uint32_t pid,ringbuffer *r)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t idx=task_idx(pid);
    
    struct pdirectory *vmem=x86_get_page_directory();//task_list[cpu][idx].vmem;
    vmem_add_framebuffer(vmem);
    klog("add win to compositor");
    compositor_add_window(vmem,pid,r);
    return 0;
}

volatile int task_reset(uint32_t pid, uint32_t entry, uint32_t stack,uint32_t brk)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t idx=task_idx(pid);
    uint32_t *stk=task_list[cpu][idx].esp;
    task_list[cpu][idx].brk=brk;

    stk[14]=entry;
    stk[17]=stack;
    return 1;
}

uint32_t task_idx(uint32_t pid)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);

    for(int i=0;i<MAX_TASKS;i++)
    {
	if(task_list[cpu][i].pid==pid)return i;
    }

    return 0;
}

uint32_t task_runs(uint32_t pid)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);

    for(int i=0;i<MAX_TASKS;i++)
    {
	if(task_list[cpu][i].active==true&& task_list[cpu][i].pid==pid)return i;
    }

    return 0;
}

void task_exit(uint32_t pid)
{
    uint32_t cpu=smp_get(SMP_APIC_ID);
    uint32_t idx=task_runs(pid);

    for(int i=0;i<MAX_TASKS;i++)
    {
	if(task_list[cpu][i].active==true&& task_list[cpu][i].pid==pid)
           task_list[cpu][i].active=false;
    }

    compositor_del_window(task_list[cpu][idx].vmem);
    vmem_free_space_dir(task_list[cpu][idx].vmem,task_list[cpu][idx].thread);
}