#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <strings.h>

#include "debug.h"
#include "code.h"
#include "mem.h"
#include "desc.h"
#include "gc.h"

// For compressing the source code a bit

#define instackb(addr) ((char*)addr >= (char*) &stack_b[0] && (char*)addr < (char*) &stack_b[STACK_SIZE_B])

#define follow_thunk(thunk) if (thunk->desc == (Desc*) __FORWARD_PTR__) thunk = thunk->_forward_ptr;

#define forward_thunk(thunk, frame_ptr) \
        Thunk* dst = get_dst(frame_ptr); \
        if(dst != NULL){ \
        dst->desc = (Desc*) __FORWARD_PTR__; \
        dst->_forward_ptr = thunk; \
        }

#define placeholder() \
        push_a(alloc_b());

#define arg_from_code(descarg, arg) \
        if(((FunEntry*) (descarg))->boxing & argmask) \
        { \
            placeholder(); \
            exec(arg, frame_ptr, stack_top_a); \
        } \
        else if(((FunEntry*) (descarg))->strictness & argmask) \
        { \
            push_a(NULL); \
            exec(arg, frame_ptr, stack_top_a); \
        } \
        else \
        { \
            arg->create_thunk(arg, &stack_a[stack_top_a++], frame_ptr); \
        } \
        argmask <<= 1;

void create_thunk_app_static(Code* expr, Thunk** target, int frame_ptr)
{
    Thunk* thunk = (Thunk*) alloc_heap(((AppEntry*) expr)->f->thunk_size);
    *target = thunk;
    thunk->desc = ((AppEntry*) expr)->f;

    assert(thunk->desc->arity == expr->nr_args);

    for (int i = 0; i < expr->nr_args; i++) {
         ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], &thunk->_args[i], frame_ptr);
    }                     
}

void create_thunk_app_dyn(Code* expr, Thunk** target, int frame_ptr)
{
    push_a(local(frame_ptr, ((AppEntry*)expr)->var.index));

    int tmp = gc_enabled;
    gc_enabled = 0;
    peek_a()->desc->eval();
    gc_enabled = tmp;
    
    Thunk* basethunk = pop_a();
    
    Desc* slice =
            get_slice(basethunk->desc->type == FT_SLICE ?
                      ((SliceEntry*) basethunk->desc)->forward_ptr : basethunk->desc, basethunk->desc->arity + expr->nr_args);

    Thunk* thunk = (Thunk*) alloc_heap(slice->thunk_size);
    thunk->desc = slice;
    *target = thunk;
    
    assert(thunk->desc->arity == basethunk->desc->arity + expr->nr_args);            

    memcpy(&thunk->_args, &basethunk->_args, sizeof(Thunk*) * basethunk->desc->arity);

    for (int i = 0; i < expr->nr_args; i++) {
        ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], &thunk->_args[basethunk->desc->arity + i], frame_ptr);
    } 
}

void create_thunk_var(Code* expr, Thunk** target, int frame_ptr)
{
    *target = local(frame_ptr, ((VarEntry*) expr)->index);    
}

void create_thunk_var_unboxed(Code* expr, Thunk** target, int frame_ptr)
{
    Thunk* arg = local(frame_ptr, ((VarEntry*) expr)->index);

    if(instackb(arg))
    {
        // The likely case
        Thunk* newthunk = (Thunk*) alloc_heap(sizeof (Thunk));
        memcpy(newthunk, arg, sizeof(Thunk));
        *target = newthunk;        
    }
    else
    {
        *target = arg;
    }
}

void create_thunk_thunk(Code* expr, Thunk** target, int frame_ptr)
{
    *target = &((ThunkEntry*) expr)->thunk;
}

void set_create_thunk_fun(Code* code)
{
    switch(code->type)
    {
        case CT_APP_PRIM1:
        case CT_APP_PRIM2:
        case CT_APP_PRIM:            
        case CT_APP_FUN:
        case CT_APP_FUN_TR:
        case CT_APP_FUN1:
        case CT_APP_FUN2:
        case CT_APP_THUNK:
            code->create_thunk = create_thunk_app_static;
            break;
        case CT_APP_DYN:
            code->create_thunk = create_thunk_app_dyn;
            break;
        case CT_VAR:
        case CT_VAR_STRICT:            
            code->create_thunk = create_thunk_var;
            break;
        case CT_VAR_UNBOXED:
            code->create_thunk = create_thunk_var_unboxed;
            break;
        case CT_THUNK:
            code->create_thunk = create_thunk_thunk;
            break;
        case CT_SELECT_ADT:
        case CT_SELECT_LIT:    
        case CT_IF:
        case CT_LET:
            code->create_thunk = NULL;
            break;            
    }
}

// eval: frame_ptr, frame_ptr
// start: stack_top_a, stack_top_a
// otherwise: frame_ptr, stack_top_a

// frame_ptr: first arguments
// root_frame_ptr: place of the result

void exec(Code* expr, int frame_ptr, int root_frame_ptr)
{
    if(heap_curr > gc_trigger) gc();
    
    int root_frame_ptr_b = stack_top_b;    
    
    while(1)
    {        
        assert(expr != NULL);
        assert(stack_top_a < STACK_SIZE_A);
        assert(stack_top_b < STACK_SIZE_B);

        // TODO: check over application
        // TODO: enforce strictness in ADT/Record
        
        switch (expr->type) {
        case CT_APP_PRIM1:
        {
            switch(expr->arg_pattern)
            {
            case 1:
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));  
                break;
            case 2:
                push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[0])->thunk);            
                break;
            default:
                placeholder();                        
                exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
                break;
            }
            
            ((PrimEntry*) ((AppEntry*) expr)->f)->exec(root_frame_ptr);

            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;                        
        }
        case CT_APP_PRIM2:
        {           
            
            // careful, "exec" may trigger garbage collection
            // read local variables only after the last exec            
            switch(expr->arg_pattern)
            {
            case 1:
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));  
                push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[1])->thunk);
                break;
            case 2:
                push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[0])->thunk);            
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));                        
                break;
            case 3:
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));
                break;                    
            case 4:
                push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[0])->thunk);            
                placeholder();                       
                exec(((AppEntry*) expr)->args[1], frame_ptr, stack_top_a);                    
                break;                    
            case 5:
                placeholder();                     
                exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
                push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[1])->thunk);                                
                break;                    
            case 6:
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));
                placeholder();                       
                exec(((AppEntry*) expr)->args[1], frame_ptr, stack_top_a);
                break;                    
            case 7:
                placeholder();                        
                exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
                push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));
                break;                    
            default:
                placeholder();                       
                exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
                placeholder();                        
                exec(((AppEntry*) expr)->args[1], frame_ptr, stack_top_a);                    
            }
            
            ((PrimEntry*) ((AppEntry*) expr)->f)->exec(root_frame_ptr);

            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;            
        }
        case CT_APP_PRIM:
        {           
            PrimEntry* desc = (PrimEntry*) ((AppEntry*) expr)->f;
            int argmask = 1;            
            
            for (int i = 0; i < desc->base.arity; i++) {

                if(desc->boxingMap & argmask)
                {
                    placeholder();
                    exec(((AppEntry*) expr)->args[i], frame_ptr, stack_top_a);
                }
                else
                {
                    push_a(NULL);
                    exec(((AppEntry*) expr)->args[i], frame_ptr, stack_top_a);
                }

                argmask <<= 1;            
            }

            desc->exec(root_frame_ptr);
            
            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;            
        }        
        case CT_APP_FUN1:
        {
            Desc* slice = ((AppEntry*) expr)->f;

            int argmask = 1;            
            arg_from_code(slice, ((AppEntry*) expr)->args[0]);
              
            expr = ((FunEntry*) slice)->body;
            frame_ptr = stack_top_a - 1;
            continue;                        
        }        
        case CT_APP_FUN2:
        {
            Desc* slice = ((AppEntry*) expr)->f;
                        
            int argmask = 1;
            
            arg_from_code(slice, ((AppEntry*) expr)->args[0]);
            arg_from_code(slice, ((AppEntry*) expr)->args[1]);
              
            expr = ((FunEntry*) slice)->body;
            frame_ptr = stack_top_a - 2;
            continue;                        
        }
        case CT_APP_FUN:
        {            
            Desc* slice = ((AppEntry*) expr)->f;

            int new_frame_ptr = stack_top_a;
            int argmask = 1;
            
            for (int i = 0; i < expr->nr_args; i++) {
                arg_from_code(slice, ((AppEntry*) expr)->args[i]);
            }

            expr = ((FunEntry*) slice)->body;
            frame_ptr = new_frame_ptr;
            continue;            
        }
        case CT_APP_FUN_TR:
        {            
            Desc* slice = ((AppEntry*) expr)->f;
            
            // TODO: B stack?
            int new_frame_ptr = stack_top_a;
            int argmask = 1;
                        
            for (int i = 0; i < expr->nr_args; i++) {
                arg_from_code(slice, ((AppEntry*) expr)->args[i]);
            }

            memcpy(&stack_a[frame_ptr], &stack_a[new_frame_ptr], sizeof(void*) * expr->nr_args);
            stack_top_a = frame_ptr + expr->nr_args;
            
            expr = ((FunEntry*) slice)->body;
            continue;            
        }
        case CT_APP_THUNK:
        {
            Desc* slice = ((AppEntry*) expr)->f;
            Thunk* thunk = get_dst(root_frame_ptr);
            
            // no need to check for array length, thunks in HNF are never overwritten
            int newsize = slice->thunk_size;
            
            if (thunk == NULL)
            {
                thunk = (Thunk*) alloc_heap(newsize);
                set_return(root_frame_ptr, thunk);
            }
            else if (thunk->desc->thunk_size < newsize) {
                Thunk* target = thunk;
                thunk = (Thunk*) alloc_heap(newsize);
                target->desc = (Desc*) __FORWARD_PTR__;
                target->_forward_ptr = thunk;
                set_return(root_frame_ptr, thunk);
            }

            thunk->desc = slice;
            
            assert(thunk->desc->arity == expr->nr_args);
            
            for (int i = 0; i < expr->nr_args; i++) {  
                ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], &thunk->_args[i], frame_ptr);
            }

            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;                                
        }
        case CT_APP_DYN:
        {
            push_a(local(frame_ptr, ((AppEntry*)expr)->var.index));

            Thunk** bt = &peek_a();
            (*bt)->desc->eval();
            
            
            Desc* baseDesc = (*bt)->desc->type == FT_SLICE ?
                              ((SliceEntry*) (*bt)->desc)->forward_ptr : (*bt)->desc;
            
            int newArity = (*bt)->desc->arity + expr->nr_args;
            
            if(newArity > baseDesc->arity)
            {
                
            }
            
            Desc* slice =
                    get_slice(baseDesc, newArity);
            
            switch(slice->type) {
            case FT_PRIM:
            {
                for (int i = 0; i < (*bt)->desc->arity; i++) {
                    push_a((*bt)->_args[i]);
                    /*if(!basethunk->_args[i]->desc->hnf)*/ (*bt)->_args[i]->desc->eval();
                    
                }

                for (int i = 0; i < expr->nr_args; i++) {
                    placeholder();
                    exec(((AppEntry*) expr)->args[i], frame_ptr, stack_top_a);
                } 

                ((PrimEntry*) slice)->exec(root_frame_ptr);
                destroy_stack_frame(root_frame_ptr);
                destroy_stack_frame_b(root_frame_ptr_b);                
                return;                    
            }
            case FT_FUN:
            {
                int new_frame_ptr = stack_top_a;                    
                int argmask = 1;

                for (int i = 0; i < (*bt)->desc->arity; i++) {                    
                    push_a((*bt)->_args[i]);
                    if(((FunEntry*) (slice))->strictness & argmask /*&& !arg->desc->hnf*/) 
                    { 
                        (*bt)->_args[i]->desc->eval(); 
                    } 
                    argmask <<= 1;
                }

                for (int i = 0; i < expr->nr_args; i++) {
                    arg_from_code(slice, ((AppEntry*) expr)->args[i]);    
                } 

                expr = ((FunEntry*) slice)->body;
                frame_ptr = new_frame_ptr;
                continue;                    
            }
            case FT_SLICE:
            case FT_ADT:
            case FT_RECORD:
            {
                Thunk* thunk = get_dst(root_frame_ptr);
                // no need to check for array length, thunks in HNF are never overwritten
                int newsize = slice->thunk_size;

                if (thunk == NULL)
                {
                    thunk = (Thunk*) alloc_heap(newsize);
                    set_return(root_frame_ptr, thunk);
                }
                else if (thunk->desc->thunk_size < newsize) {
                    Thunk* target = thunk;
                    thunk = (Thunk*) alloc_heap(newsize);
                    target->desc = (Desc*) __FORWARD_PTR__;
                    target->_forward_ptr = thunk;
                    set_return(root_frame_ptr, thunk);
                }
                
                thunk->desc = slice;
                
                assert(thunk->desc->arity == (*bt)->desc->arity + expr->nr_args);            

                memcpy(&thunk->_args, &(*bt)->_args, sizeof(Thunk*) * (*bt)->desc->arity);

                for (int i = 0; i < expr->nr_args; i++) {
                    ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], &thunk->_args[(*bt)->desc->arity + i], frame_ptr);
                }

                destroy_stack_frame(root_frame_ptr);
                destroy_stack_frame_b(root_frame_ptr_b);
                return;                    
            }
            case FT_BOXED_LIT:
                abort("Literal unexpected here");                
            case FT_CAF:
            case FT_CAF_REDUCED:
                not_implemented("CAF");
            }
        }
        case CT_VAR_STRICT:
        case CT_VAR_UNBOXED:            
        {
            Thunk* arg = local(frame_ptr, ((VarEntry*) expr)->index);
            
            assert(is_hnf(arg));

            // no need to check for array length, thunks in HNF are never overwritten
            if(get_dst(root_frame_ptr) != NULL && arg->desc->thunk_size <= sizeof(Thunk))
            {
                memcpy(get_dst(root_frame_ptr), arg, sizeof(Thunk));
            }
            else
            {
                forward_thunk(arg, root_frame_ptr);
                set_return(root_frame_ptr, arg);
            }
            
            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);            
            return;                    
        }
        case CT_VAR:
        {
            Thunk* thunk = local(frame_ptr, ((VarEntry*) expr)->index);

            assert(!instackb(thunk));
            
            follow_thunk(thunk);
            forward_thunk(thunk, root_frame_ptr);            
            set_return(root_frame_ptr, thunk);

            switch(thunk->desc->type) {
            case FT_FUN:
            {
                // Destroy stack frame before eval, it is not needed any more
                // Greatly reduces stack usage
                destroy_stack_frame(root_frame_ptr);
                destroy_stack_frame_b(root_frame_ptr_b);  
                                
                frame_ptr = stack_top_a;
                // Here frame_ptr == root_frame_ptr

                int argmask = 1;

                for (int i = 0; i < thunk->desc->arity; i++) {
                    push_a(thunk->_args[i]);
                    if(((FunEntry*) (thunk->desc))->strictness & argmask /*&& !arg->desc->hnf*/) 
                    { 
                        thunk->_args[i]->desc->eval(); 
                        thunk = stack_a[root_frame_ptr-1];
                    }
                    argmask <<= 1;
                }

                expr = ((FunEntry*) thunk->desc)->body;
                continue;
            }
            case FT_PRIM:
            {

                for (int i = 0; i < thunk->desc->arity; i++) {
                    push_a(thunk->_args[i]);
                    thunk->_args[i]->desc->eval();
                    thunk = stack_a[root_frame_ptr-1];
                }

                ((PrimEntry*) thunk->desc)->exec(root_frame_ptr);
                
                destroy_stack_frame(root_frame_ptr);
                destroy_stack_frame_b(root_frame_ptr_b);                
                return;
            }
            case FT_CAF:
            case FT_CAF_REDUCED:
                not_implemented("CAF");            
            case FT_SLICE:
            case FT_ADT:
            case FT_RECORD:
            case FT_BOXED_LIT:
                destroy_stack_frame(root_frame_ptr);
                destroy_stack_frame_b(root_frame_ptr_b);
                return;
            }            
        }
        case CT_THUNK:
        {
            Thunk* thunk = &((ThunkEntry*) expr)->thunk;            
            Thunk* dst = get_dst(root_frame_ptr);

            if(dst != NULL)
            {
                memcpy(dst, thunk, sizeof(Thunk));
                
                if(!instackb(dst))
                {
                    set_return(root_frame_ptr, thunk);
                }                
            }
            else
            {
                set_return(root_frame_ptr, thunk);
            }

            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;               
        }
        case CT_SELECT_LIT:
        {
            placeholder();
            exec(((SelectEntry*) expr)->expr, frame_ptr, stack_top_a);
            Thunk* lit = pop_a();
            
            bool handled = false;
            
            for (int i = 0; i < expr->nr_cases; i++) {
                SelectLitCaseEntry* caseEntry = &((SelectEntry*) expr)->cases[i];

                // NULL means "default", we accept it anyway
                if(caseEntry->lit != NULL)
                {
                    assert(caseEntry->lit->thunk.desc == (Desc*) __INT__ ||
                           caseEntry->lit->thunk.desc == (Desc*) __BOOL__ );
 
                    if(caseEntry->lit->thunk._int != lit->_int) continue;
                }
                
                // must be SC_DEFAULT now
                handled = true;
                expr = caseEntry->body;
                break;
            }
            
            if(handled) continue;
            
            if(((SelectEntry*) expr)->fallback != NULL)
            {
                stack_top_a -= ((SelectEntry*) expr)->fallback_nrargs;
                expr = ((SelectEntry*) expr)->fallback;
                
                continue;
            }
            
            abort("no match");
        }
        case CT_SELECT_ADT:
        {                        
            SelectEntry* select = (SelectEntry*) expr;
            
            push_a(NULL);
            exec(select->expr, frame_ptr, stack_top_a);
            Thunk* cons = pop_a();            
            
            expr = select->bodies[((ADTEntry*)cons->desc)->idx];
            
            if(expr != NULL)
            {
                for (int i = 0; i < cons->desc->arity; i++) {
                    push_a(cons->_args[i]);
                }                  
                
                continue;
            }
            
            if(select->fallback != NULL)
            {
                stack_top_a -= select->fallback_nrargs;
                expr = select->fallback;
                
                continue;
            }
            
            abort("no match");
        }
        case CT_IF:
        {        
            placeholder();
            exec(((IfEntry*) expr)->cond, frame_ptr, stack_top_a);
            Thunk* cond = pop_a();
            
            // safe to do it before read as nothing can overwrite it in between
            stack_top_b--;
            
            if (readB(cond)) {
                expr = ((IfEntry*) expr)->texpr;
                continue;                
            }
            else {
                expr = ((IfEntry*) expr)->fexpr;
                continue;     
            }
        }
        case CT_LET:
        {
            for(int i=0; i<expr->nr_bindings; i++)
            {
                LetBindingEntry* binding = ((LetEntry*) expr)->bindings[i];
                
                // Normal
                if(binding->type == 0)
                {
                    binding->body->create_thunk(binding->body, &stack_a[stack_top_a++], frame_ptr);
                }
                else // strict (including unboxed))
                {
                    push_a(NULL);
                    exec(binding->body, frame_ptr, stack_top_a);
                }
            }
            
            expr = ((LetEntry*) expr)->body;
            continue;
        }
        }
    }
}

void eval_hnf()
{
    return;
}

void eval_fun()
{
    Thunk* thunk = peek_a();    
    int frame_ptr = stack_top_a;
    int argmask = 1;

    for (int i = 0; i < thunk->desc->arity; i++) {
        push_a(thunk->_args[i]);

        if(((FunEntry*) (thunk->desc))->strictness & argmask)
        {
            thunk->_args[i]->desc->eval();
            thunk = stack_a[frame_ptr-1]; // refresh thunk ptr after eval 
        }
        
        argmask <<= 1;
    }

    exec(((FunEntry*) thunk->desc)->body, frame_ptr, frame_ptr);
}

void eval_prim()
{
    Thunk* thunk = peek_a();
    int frame_ptr = stack_top_a;

    for (int i = 0; i < thunk->desc->arity; i++) {
        push_a(thunk->_args[i]);        
        thunk->_args[i]->desc->eval();
        thunk = stack_a[frame_ptr-1]; // refresh thunk ptr after eval 
    }

    ((PrimEntry*) thunk->desc)->exec(frame_ptr);        

    stack_top_a = frame_ptr;        
}

void eval_fwd_ptr()
{
    Thunk* thunk = pop_a();
    follow_thunk(thunk);
    push_a(thunk);
    thunk->desc->eval();
}   
   
void set_eval_fun(Desc* desc)
{
    if(desc == (Desc*) __FORWARD_PTR__)
    {
        desc->eval = eval_fwd_ptr;
        return;
    }
    
    switch(desc->type)
    {
        case FT_BOXED_LIT:
        case FT_RECORD:
        case FT_ADT:
        case FT_CAF:
        case FT_CAF_REDUCED:
        case FT_SLICE:
            desc->eval = eval_hnf;
            break;
        case FT_FUN:
            desc->eval = eval_fun;
            break;
        case FT_PRIM:
            desc->eval = eval_prim;
            break;
    }
}