// ptsk.h
#ifndef __PTSK_H__
#define __PTSK_H__
#include <setjmp.h>

#ifdef __cplusplus
extern "C" {
#endif

#define PTSK_CFG_FAKE_TIME
#define PTSK_CFG_PANIC_EXIT

typedef long ptsk_time_t;
ptsk_time_t ptsk_time(ptsk_time_t shift);

typedef struct ptsk_tcb_tag {
    void (*task)(void* args); void *args;
    int   stack_size;
    int  *stack_top;
    jmp_buf ctx;
    struct ptsk_tcb_tag *next, *prev;
} ptsk_tcb_t;

void ptsk_init(int main_stack_size);
void ptsk_done();
void ptsk_idle();
void ptsk_addtask(ptsk_tcb_t *task);
void ptsk_deltask(ptsk_tcb_t *task); // if current it'll die
void ptsk_die();                     // kill current task
ptsk_tcb_t* ptsk_current();

void ptsk_sleep(ptsk_time_t ticks);
void ptsk_sleep_till(ptsk_time_t ticks);

#ifdef __cplusplus
} // extern "C"
#endif

#endif // __PTSK_H__

// ptsk.c
//#include "ptsk.h"

#ifdef PTSK_CFG_PANIC_EXIT
#include <stdlib.h>
#endif

#ifdef __cplusplus
extern "C" {
#endif


#ifdef PTSK_CFG_FAKE_TIME
    static ptsk_time_t ptsk_time_value=0;
    ptsk_time_t ptsk_time(ptsk_time_t from) {
        // use unsigned to avoid signed overflow UB
        return (long)((unsigned long)ptsk_time_value-(unsigned long)from);
    }
    static void ptsk_time_advance() { ptsk_time_value+=100; }
    #define PTSK_TIME_ADVANCE() ptsk_time_advance()
#else
    #define PTSK_TIME_ADVANCE()
    extern ptsk_time_t ptsk_time(ptsk_time_t from);
#endif

static void ptsk_panic() {
    #ifdef PTSK_CFG_PANIC_EXIT
    exit(0);
    #else
    for(;;) {} // halt
    #endif
}

//------------------------------------------------

enum { 
    PTSK_1ST,  // first call 
    PTSK_WAKE, // wake task
    PTSK_NEW,  // new task
    PTSK_RET   // task return
};

static char        ptsk_active=0;
static ptsk_tcb_t* ptsk_head=0;
static ptsk_tcb_t* ptsk_tail=0;
static ptsk_tcb_t* ptsk_curr=0;
static jmp_buf     ptsk_last, ptsk_main;
static int         ptsk_stack_last;

void ptsk_init(int main_stack_size) {
    ptsk_active=0;
    ptsk_head=ptsk_tail=ptsk_curr=0;
    ptsk_stack_last=main_stack_size;
}
void ptsk_done() {
    if (ptsk_active) longjmp(ptsk_main,PTSK_RET);
}

static void ptsk_switch() {
    ptsk_curr=ptsk_curr->next; 
    if (!ptsk_curr) {
        ptsk_curr=ptsk_head;
        if (!ptsk_curr) return ptsk_done();
        PTSK_TIME_ADVANCE();
    }
    longjmp(ptsk_curr->ctx,PTSK_WAKE);          
}
void ptsk_die() {
    if (!ptsk_curr) { ptsk_done(); ptsk_panic(); return; } // panic
    if (ptsk_curr->prev) ptsk_curr->prev->next=ptsk_curr->next; else ptsk_head=ptsk_curr->next;
    if (ptsk_curr->next) ptsk_curr->next->prev=ptsk_curr->prev; else ptsk_tail=ptsk_curr->prev;
    ptsk_switch();
}
void ptsk_deltask(ptsk_tcb_t *task) {
    if (ptsk_curr==task) ptsk_die();
    if (task->prev) task->prev->next=task->next; else ptsk_head=task->next;
    if (task->next) task->next->prev=task->prev; else ptsk_tail=task->prev;
}
void ptsk_run(ptsk_tcb_t* task) {
    task->stack_top=(int*)&task;
    task->task(task->args);
    ptsk_die();
}
static int ptsk_stackalloc(int alloc) {
    volatile int dummy=0; ptsk_tcb_t* tsk; int sz;
    static char* stk;
    if (alloc==1) stk=(char*)&dummy; 
    sz=(char*)&dummy-stk; if (sz<0) sz=-sz;
    if (sz<ptsk_stack_last) ptsk_stackalloc(0);
    ptsk_stack_last=ptsk_tail->stack_size;
    tsk=ptsk_tail;
    if (setjmp(ptsk_last)==PTSK_NEW) ptsk_stackalloc(1);
    switch(setjmp(tsk->ctx)) {
        case PTSK_1ST:  longjmp(ptsk_main,PTSK_RET);
        case PTSK_WAKE: ptsk_run(tsk);
    }
    ptsk_die(); return dummy;
}
void ptsk_addtask(ptsk_tcb_t *task) {
    ptsk_tcb_t* last;
    last=ptsk_tail;
    task->prev=ptsk_tail;
    task->next=0;
    if (ptsk_tail) ptsk_tail->next=task; else ptsk_head=task;
    ptsk_tail=task;
    switch(setjmp(ptsk_main)) {
        case PTSK_1ST:  if (last) longjmp(ptsk_last,PTSK_NEW); else ptsk_stackalloc(1);
        //case PTSK_RET: break;
    }
}
void ptsk_idle() {
    if (!ptsk_curr) { 
        ptsk_curr=ptsk_head; if (!ptsk_curr) return;
        ptsk_active=1;
        switch(setjmp(ptsk_main)) {
            case PTSK_1ST: ptsk_run(ptsk_curr);
            //case PTSK_RET: break;
        }
        ptsk_active=0;
        return;
    }
    switch(setjmp(ptsk_curr->ctx)) {
        case PTSK_1ST:  ptsk_switch();
        //case PTSK_WAKE: break;
    }
}
ptsk_tcb_t* ptsk_getcurrent() {
    return ptsk_curr;
}
void ptsk_sleep_till(ptsk_time_t time) {
    do {
        ptsk_idle();
    } while(ptsk_time(time)<0);
}
void ptsk_sleep(ptsk_time_t ticks) {
    ptsk_sleep_till(ptsk_time(-ticks));
}

#ifdef __cplusplus
} // extern "C"
#endif

//----------------------------------------------------------------------------

#include <stdio.h>

void task1(void* arg) {
    for(int i=0;i<=9;++i) {
        printf("task1 \ta=%d\n",i);
        ptsk_idle();
    }
}
void task2(void* arg) {
    for(int i=0;i<=7;++i) {
        printf("task2 \t\tb=%d\n",i);
        ptsk_sleep(500);
    }
}
void task3(void* arg) {
    for(int i=0;i<=5;++i) {
        printf("task3 \t\t\tc=%d\n",i);
        ptsk_sleep(700);
    }
}
void task4(void* arg) {
    for(int i=0;i<=4;++i) {
        printf("task4 \t\t\t\td=%d\n",i);
        ptsk_sleep(100);
    }
}
void task5(void* arg) {
    for(int i=0;i<=3;++i) {
        printf("task5 \t\t\t\t\te=%d\n",i);
        ptsk_sleep(300);
    }
}



void ptsk_test() {
    enum { stack_size=4096 };
    ptsk_init(stack_size);
    ptsk_tcb_t t[]={
        { task1,(void*)1,stack_size },
        { task2,(void*)2,stack_size },
        { task3,(void*)3,stack_size },
        { task4,(void*)4,stack_size },
        { task5,(void*)5,stack_size }
    };
    ptsk_addtask(&t[0]);
    ptsk_addtask(&t[1]);
    ptsk_addtask(&t[2]);
    ptsk_addtask(&t[3]);
    ptsk_addtask(&t[4]);
    ptsk_idle();
    ptsk_done();
}

int main(int argc,char** argv) {
    ptsk_test();
    return 0;
}