#ifndef __TASK_H__
#define __TASK_H__

extern "C" {
#include "multiboot.h"
}

#include "types.h"
#include "queue.h"
#include "memmgr.h"
#include "timer.h"
#include "gdt.h"

class TaskScheduler;

struct TaskContext
{
    uint32 ds;
    uint32 es;
    uint32 fs;
    uint32 gs;
    uint32 edi;
    uint32 esi;
    uint32 ebp;
    uint32 esp; // do not use this value
    uint32 ebx;
    uint32 edx;
    uint32 ecx;
    uint32 eax;
    uint32 eip;
    uint32 cs;
    uint32 eflags;
} __attribute__ ((packed));

struct TCB
{
    uint32        StackSegment;
    TaskContext*  Context;
};

class Task : public Object
{
public:
    enum TaskStatus {
        RUNNING = 0,
        READY   = 1,
        WAIT    = 2,
    };

    enum TaskPriority {
        MIN_PRIORITY    = 0,
        NORMAL_PRIORITY = 5,
        MAX_PRIORITY    = 9,
    };

    Task(size_t stack_size);
    Task(uint8* sp, size_t stack_size) : 
        Stack(sp), StackSize(stack_size), 
        Priority(NORMAL_PRIORITY), isStackInMemPool(false) {}

    TCB& getTcb(void) {return Tcb;}
    void setStatus(TaskStatus st) {Status = st;}
    TaskStatus getStatus(void) {return Status;}

    TaskPriority getPriority(void) {return Priority;}
    TaskPriority getSubPriority(void) {return SubPriority;}
    void setPriority(TaskPriority pri) {Priority = pri;}

    virtual void run(void) {}

private:
    TCB          Tcb;
    uint8*       Stack;
    size_t       StackSize;
    TaskStatus   Status;
    TaskPriority Priority;
    TaskPriority SubPriority;
    Queue        EventQueue;

    bool isStackInMemPool;
};

class TaskScheduler : public TimerJob
{
public:
    TaskScheduler(void) : TimerJob(1), ReadyQueue(TaskCmp) {}
    void  init(Timer& switch_timer, Task& current_task);
    void  add(Task& task) {addReadyQueue(task);}
    void  dispatch(bool dispatch_immd = true);
    void  sleep(Task& task, time_t duration = 0);
    void  wakeup(Task& task);

    Task* getCurrentTask(void) {return CurrentTask;}
    TCB*  getCurrentTcb(void);
    static TCB* CurrentTcb;

    void onExpire(void) {
        SwitchTimer->add(*this);
        if (--TimeSlice <= 0) dispatch(false);
    }

private:
    enum {
        HIGH_TIME_SLICE   = 4,
        NORMAL_TIME_SLICE = 2, 
        LOW_TIME_SLICE    = 1, 
    };

    Task*         CurrentTask;
    PriorityQueue ReadyQueue;
    Timer*        SwitchTimer;
    int           TimeSlice;

    class TaskComparator : public Comparator {
        int compare(Object& src1, Object& src2) {
            Task& t1 = (Task&)src1;
            Task& t2 = (Task&)src2;
            return int(t1.getPriority() - t2.getPriority());
        }
    } TaskCmp;

    void addReadyQueue(Task& task) {
        ReadyQueue.add(task);
        task.setStatus(Task::READY);
    }
};

extern TaskScheduler SystemScheduler;

#endif /* __TASK_H__ */
