// TODO(allen): add profiling markup

////////////////////////////////
// Coroutine Functions

MR4TH_THREADVAR CO_Ctx* co__ctx = 0;

MR4TH_SYMBOL_STATIC void
co__main(void *udata){
    // read parameters (udata goes bad after the handshake!)
    CO_Ctx *ctx = 0;
    CO_Coroutine *co = 0;
    {  
        void **params = (void**)udata;
        ctx = (CO_Ctx*)params[0];
        co = (CO_Coroutine*)params[1];
    }
    
    // set threadvar
    co__ctx = ctx;
    
    // finish the handshake
    th_mutex_wait_for_lock(&ctx->mutex);
    ctx->did_init = 1;
    th_condvar_signal(&ctx->cv);
    
    for (;;){
        // wait for a wake up
        co__wait_until_active(ctx, co);
        
        // when we wake up, run the func
        Assert(co->next != 0);
        Assert(co->func != 0);
        co->func(co->param);
        
        // remember parent
        CO_Coroutine *parent = co->next;
        
        // mark this coroutine as dead & free
        co->func = 0;
        co->in = 0;
        co->out = 0;
        co->next = 0;
        co->state = CO_State_Dead;
        SLLStackPush(ctx->free, co);
        
        // return control up the chain
        parent->state = CO_State_Active;
        ctx->active = parent;
        co__pass_control_to(parent);
    }
}

MR4TH_SYMBOL_STATIC B32
co_ctx_init(U64 max_co_count){
    Assert(co__ctx == 0);
    B32 success = 0;
    
    if (max_co_count > 0){
        U64 raw_size = sizeof(CO_Ctx) + max_co_count*sizeof(CO_Coroutine);
        U64 size = AlignUpPow2(raw_size, KB(4));
        
        void *memory = os_memory_reserve(size);
        if (os_memory_commit(memory, size)){
            CO_Ctx *ctx = (CO_Ctx*)memory;
            ctx->size = size;
            
            // initialize coroutine memory pool
            ctx->buf = (CO_Coroutine*)(ctx + 1);
            ctx->ptr = ctx->buf;
            ctx->opl = ctx->buf + max_co_count;
            ctx->free = 0;
            ctx->active = 0;
            
            // initialize root
            ctx->root.state = CO_State_Active;
            ctx->active = &ctx->root;
            // (root does not get it's own thread -
            //  it stands for the thread that constructs the ctx)
            th_condvar_init(&ctx->root.cv);
            
            // initialize crit & cv
            th_mutex_init(&ctx->mutex);
            th_condvar_init(&ctx->cv);
            
            // (from now on, always hold the crit on the active thread)
            th_mutex_wait_for_lock(&ctx->mutex);
            
            // save to threadvar
            co__ctx = ctx;
            
            success = 1;
        }
    }
    
    return(success);
}

MR4TH_SYMBOL_STATIC void
co_ctx_release(void){
    Assert(co__ctx != 0);
    CO_Ctx *ctx = co__ctx;
    Assert(ctx->active == &ctx->root);
    
    // free each coroutine
    for (CO_Coroutine *co = ctx->buf;
         co != ctx->ptr;
         co += 1){
        th_thread_force_stop(&co->thread);
        // (cv requires no release)
    }
    // (root has no thread to release)
    
    // free ctx
    th_mutex_release(&ctx->mutex);
    os_memory_release(ctx, ctx->size);
    
    // zero threadvar
    co__ctx = 0;
}

MR4TH_SYMBOL_STATIC CO_Coroutine*
co_create(CO_Func *func, void *param, void *udata){
    CO_Ctx *ctx = co__ctx;
    CO_Coroutine *co = co__alloc(ctx);
    if (co != 0){
        co->state = CO_State_Inactive;
        co->func  = func;
        co->param = param;
        co->udata = udata;
    }
    return(co);
}

MR4TH_SYMBOL_STATIC CO_Coroutine*
co_run(CO_Coroutine *co, void *in, void *out){
    CO_Ctx *ctx = co__ctx;
    CO_Coroutine *parent = ctx->active;
    
    Assert(co->state == CO_State_Inactive);
    Assert(parent->state == CO_State_Active);
    Assert(co != parent);
    
    co->in = in;
    co->out = out;
    co->next = parent;
    
    parent->state = CO_State_Waiting;
    co->state = CO_State_Active;
    ctx->active = co;
    
    co__pass_control_to(co);
    co__wait_until_active(ctx, parent);
    
    CO_Coroutine *result = co;
    if (co->state == CO_State_Dead){
        result = 0;
    }
    return(result);
}

MR4TH_SYMBOL_STATIC void
co_yield(void){
    CO_Ctx *ctx = co__ctx;
    CO_Coroutine *co = ctx->active;
    CO_Coroutine *parent = co->next;
    
    Assert(co->state == CO_State_Active);
    Assert(parent->state == CO_State_Waiting);
    Assert(co != parent);
    
    co->in = 0;
    co->out = 0;
    co->next = 0;
    
    co->state = CO_State_Inactive;
    parent->state = CO_State_Active;
    ctx->active = parent;
    
    co__pass_control_to(parent);
    co__wait_until_active(ctx, co);
}

MR4TH_SYMBOL_STATIC CO_Coroutine*
co_me(void){
    return(co__ctx->active);
}

MR4TH_SYMBOL_STATIC void*
co_param(void){
    return(co__ctx->active->param);
}

MR4TH_SYMBOL_STATIC void*
co_udata(void){
    return(co__ctx->active->udata);
}

MR4TH_SYMBOL_STATIC void*
co_in(void){
    return(co__ctx->active->in);
}

MR4TH_SYMBOL_STATIC void*
co_out(void){
    return(co__ctx->active->out);
}

MR4TH_SYMBOL_STATIC CO_Coroutine*
co__alloc(CO_Ctx *ctx){
    CO_Coroutine *result = ctx->free;
    if (result != 0){
        SLLStackPop(ctx->free);
    }
    else if (ctx->ptr < ctx->opl){
        result = ctx->ptr;
        ctx->ptr += 1;
        
        // init cv & thread
        void *params[] = {ctx, result};
        th_condvar_init(&result->cv);
        ctx->did_init = 0;
        th_thread_create(co__main, params, &result->thread);
        co__wait_for_init(ctx);
    }
    return(result);
}

MR4TH_SYMBOL_STATIC void
co__pass_control_to(CO_Coroutine *co){
    th_condvar_signal(&co->cv);
}

MR4TH_SYMBOL_STATIC void
co__wait_until_active(CO_Ctx *ctx, CO_Coroutine *co){
    for (;co->state != CO_State_Active;){
        th_condvar_wait_for_signal(&co->cv, &ctx->mutex);
    }
}

MR4TH_SYMBOL_STATIC void
co__wait_for_init(CO_Ctx *ctx){
    for (;!ctx->did_init;){
        th_condvar_wait_for_signal(&ctx->cv, &ctx->mutex);
    }
}
