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

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

// For compressing the source code a bit

#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); \
        dst->desc = (Desc*) __FORWARD_PTR__; \
        dst->_forward_ptr = thunk; 

#define arg_from_thunk(descarg, arg) \
        if(((FunEntry*) (descarg))->strictness & argmask && !arg->desc->hnf) \
        { \
            push_a(eval(arg)); \
        } \
        else \
        { \
            push_a(arg); \
        } \
        argmask <<= 1;

#define arg_from_code(descarg, arg) \
        if(((FunEntry*) (descarg))->strictness & argmask) \
        { \
            Thunk* phl = alloc_b(); \
            phl->desc = (Desc*) __STACK_PLACEHOLDER__; \
            push_a(phl); \
            exec(arg, frame_ptr, stack_top_a); \
        } \
        else \
        { \
            push_a(arg->create_thunk(arg, frame_ptr)); \
        } \
        argmask <<= 1;

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

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

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

    return thunk;                        
}

struct Thunk* create_thunk_app_dyn(Code* expr, int frame_ptr)
{
    Thunk* basethunk = local(frame_ptr, ((AppEntry*)expr)->var.index);
    if(!basethunk->desc->hnf) basethunk = eval(basethunk);

    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;

    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++) {
        thunk->_args[basethunk->desc->arity + i] 
                = ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], frame_ptr);
    }

    return thunk;    
}

struct Thunk* create_thunk_var(Code* expr, int frame_ptr)
{
    return local(frame_ptr, ((VarEntry*) expr)->index);    
}

struct Thunk* create_thunk_var_strict(Code* expr, int frame_ptr)
{
    Thunk* arg = local(frame_ptr, ((VarEntry*) expr)->index);

    if(arg->desc->unboxable) // unboxable means it is on the B stack
    {
        Thunk* target = (Thunk*) alloc_heap(sizeof (Thunk));
        memcpy(target, arg, sizeof(Thunk));
        return target;
    }
    else
    {
        return arg;
    }
}

struct Thunk* create_thunk_thunk(Code* expr, int frame_ptr)
{
    return &((ThunkEntry*) expr)->thunk;
}

void set_create_thunk_fun(Code* code)
{
    switch(code->type)
    {
        case CT_APP_PRIM1:
        case CT_APP_PRIM_S:
        case CT_APP_PRIM2:
        case CT_APP_PRIM_ST:
        case CT_APP_PRIM_TS:
        case CT_APP_PRIM_SS:
        case CT_APP_PRIM_AT:
        case CT_APP_PRIM_TA:
        case CT_APP_PRIM_AS:
        case CT_APP_PRIM_SA:
        case CT_APP_FUN:
        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:
            code->create_thunk = create_thunk_var;
            break;
        case CT_VAR_STRICT:
            code->create_thunk = create_thunk_var_strict;
            break;
        case CT_THUNK:
            code->create_thunk = create_thunk_thunk;
            break;
        case CT_SELECT:
        case CT_IF:
            code->create_thunk = NULL;
            break;            
    }
}

void exec(Code* expr, int frame_ptr, int root_frame_ptr)
{
    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:
        {
            push_a(alloc_b());                        
            exec(((AppEntry*) expr)->args[0], 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_S:
        {
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));                        
            ((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_ST:
        {            
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));                        
            push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[1])->thunk);                        
            ((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_TS:
        {            
            push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[0])->thunk);            
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));                        
            ((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_SS:
        {            
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));
            ((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_TA:
        {            
            push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[0])->thunk);            
            push_a(alloc_b());                        
            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_AT:
        {            
            push_a(alloc_b());                        
            exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
            push_a(&((ThunkEntry*) ((AppEntry*) expr)->args[1])->thunk);            
            ((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_AS:
        {            
            push_a(alloc_b());                        
            exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[1])->index));
            ((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_SA:
        {            
            push_a(local(frame_ptr, ((VarEntry*) ((AppEntry*) expr)->args[0])->index));
            push_a(alloc_b());                        
            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_PRIM2:
        {            
            push_a(alloc_b());                        
            exec(((AppEntry*) expr)->args[0], frame_ptr, stack_top_a);
            push_a(alloc_b());                        
            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_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_THUNK:
        {
            Desc* slice = ((AppEntry*) expr)->f;
            Thunk* thunk = get_dst(root_frame_ptr);
            int newsize = slice->thunk_size;
            
            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++) {
                thunk->_args[i]  
                        = ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[i], frame_ptr);
            }

            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;                                
        }
        case CT_APP_DYN:
        {
            Thunk* basethunk = local(frame_ptr, ((AppEntry*)expr)->var.index);
            if(!basethunk->desc->hnf) basethunk = eval(basethunk);

            Desc* slice =
                    get_slice(basethunk->desc->type == FT_SLICE ?
                              ((SliceEntry*) basethunk->desc)->forward_ptr : basethunk->desc, basethunk->desc->arity + expr->nr_args);

            switch(slice->type) {
            case FT_PRIM1:
            case FT_PRIM2:
            {
                for (int i = 0; i < basethunk->desc->arity; i++) {
                    if(basethunk->_args[i]->desc->hnf)
                    {
                        push_a(basethunk->_args[i]);
                    }
                    else
                    {
                        push_a(eval(basethunk->_args[i]));   
                    }
                }

                for (int i = 0; i < expr->nr_args; i++) {
                    push_a(alloc_b());
                    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 < basethunk->desc->arity; i++) {
                    arg_from_thunk(slice, basethunk->_args[i])
                }

                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);
                int newsize = slice->thunk_size;

                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 == 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++) {
                    thunk->_args[basethunk->desc->arity + i] 
                            = ((AppEntry*) expr)->args[i]->create_thunk(((AppEntry*) expr)->args[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:
        {
            Thunk* arg = local(frame_ptr, ((VarEntry*) expr)->index);
            
            assert(is_hnf(arg));

            // TODO: check how often happens
            if(arg->desc->unboxable)
            {
                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);

            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++) {
                    arg_from_thunk(thunk->desc, thunk->_args[i]);
                }

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

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

                ((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;            
            forward_thunk(thunk, root_frame_ptr);            
            set_return(root_frame_ptr, thunk);
            destroy_stack_frame(root_frame_ptr);
            destroy_stack_frame_b(root_frame_ptr_b);
            return;               
        }
        case CT_SELECT:
        {                        
            bool handled = false;
            
            Thunk* pattern = alloc_b();
            pattern->desc = (Desc*) __STACK_PLACEHOLDER__;
            push_a(pattern);
            exec(((SelectEntry*) expr)->expr, frame_ptr, stack_top_a);
            pattern = pop_a();            
            
            int i = 0;

            do
            {                
                assert(is_hnf(pattern));
                                
                for (;i < expr->nr_cases; i++) {
                    SelectCaseEntry* caseEntry = &((SelectEntry*) expr)->cases[i];
                    
                    if (caseEntry->type == SC_CONS) {

                        // Pattern match
                        if ((Desc*) caseEntry->cons != pattern->desc) continue;

                        // Put the constructor arguments to the stack if matches
                        for (int i = 0; i < pattern->desc->arity; i++) {
                            push_a(pattern->_args[i]);
                        }                  
                    }
                    else if(caseEntry->type == SC_DEFAULT)
                    {
                        // accept it
                    }
                    else if (caseEntry->type == SC_LIT) {
                        assert(caseEntry->lit->thunk.desc != (Desc*) __INT__);
                                        
                        if(caseEntry->lit->thunk._int != pattern->_int) continue;
                    }

                    // must be SC_DEFAULT now
                    expr = caseEntry->body;
                    
                    if(expr->type == CT_SELECT)
                    {
                        i=-1;
                        
                        ((SelectEntry*) expr)->saved_pattern = pattern;
                        
                        pattern = alloc_b();
                        pattern->desc = (Desc*) __STACK_PLACEHOLDER__;
                        push_a(pattern);
                        exec(((SelectEntry*) expr)->expr, frame_ptr, stack_top_a);
                        pattern = pop_a();                        
                        
                        continue;
                    }
                    
                    handled = true;                    
                    break;
                }                    

                if(!handled)
                {
                    if(((SelectEntry*) expr)->fallback != NULL)
                    {
                        expr = (Code*) ((SelectEntry*) expr)->fallback;
                        i = ((SelectEntry*) expr)->fallbackidx;
                        pattern = ((SelectEntry*) expr)->saved_pattern;
                    }
                    else
                    {
                        printf("Exec: no select cases matches");
                        print(pattern, false);
                        exit(-1);                                                
                    }
                }
            }            
            while(!handled);
            
            continue;
        }
        case CT_IF:
        {        
            push_a(alloc_b());
            exec(((IfEntry*) expr)->cond, frame_ptr, stack_top_a);
            Thunk* cond = pop_a();
            
            if (readB(cond)) {
                expr = ((IfEntry*) expr)->texpr;
                continue;                
            }
            else {
                expr = ((IfEntry*) expr)->fexpr;
                continue;     
            }
        }
        }
    }
}

struct Thunk* eval(Thunk* thunk) {
    assert(thunk != NULL);    
        
    follow_thunk(thunk);
    
    switch(thunk->desc->type) {
    case FT_FUN:
    {
        push_a(thunk);
        int frame_ptr = stack_top_a;
        int argmask = 1;
        
        for (int i = 0; i < thunk->desc->arity; i++) {
            arg_from_thunk(thunk->desc, thunk->_args[i]);
        }
        
        exec(((FunEntry*) thunk->desc)->body, frame_ptr, frame_ptr);
        thunk = pop_a();
        return thunk;
    }
    case FT_PRIM1:
    case FT_PRIM2:
    {
        push_a(thunk);
        int frame_ptr = stack_top_a;

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

        ((PrimEntry*) thunk->desc)->exec(frame_ptr);        
        
        stack_top_a = frame_ptr;        
        thunk = pop_a();
        return thunk;
    }
    case FT_CAF:
    case FT_CAF_REDUCED:
        not_implemented("CAF");            
    case FT_SLICE:
    case FT_ADT:
    case FT_RECORD:
    case FT_BOXED_LIT:
        return thunk;
    }
}