#include <vslc.h>

static void node_print ( node_t *root, int nesting );
static void simplify_tree ( node_t **simplified, node_t *root );
static void node_finalize ( node_t *discard );

typedef struct stem_t *stem;
struct stem_t { const char *str; stem next; };
static void
tree_print(node_t* root, stem head);


static void destroy_subtree ( node_t *discard );

static void prune_children(node_t **simplified, node_t *root);
static void resolve_constant_expressions(node_t **simplified, node_t *root);
static void flatten(node_t **simplified, node_t *root);


/* External interface */
void
destroy_syntax_tree ( void )
{
    destroy_subtree ( root );
}


void
simplify_syntax_tree ( void )
{
    simplify_tree ( &root, root );
}


extern bool new_print_style;
void
print_syntax_tree ( void )
{
    if (new_print_style)
        tree_print ( root, 0 );
    // Old tree printing
    else
        node_print ( root, 0 );
}


// Changed so it returns the pointer to the new node, can be used as before, but makes the parser file cleaner
node_t*
node_init (node_t *nd, node_index_t type, void *data, uint64_t n_children, ...)
{
    va_list child_list;
    *nd = (node_t) {
        .type = type,
        .data = data,
        .entry = NULL,
        .n_children = n_children,
        .children = (node_t **) malloc ( n_children * sizeof(node_t *) )
    };
    va_start ( child_list, n_children );
    for ( uint64_t i=0; i<n_children; i++ )
        nd->children[i] = va_arg ( child_list, node_t * );
    va_end ( child_list );

    return nd;
}


static void
tree_print(node_t* root, stem head) 
{
    static const char *sdown = " │", *slast = " └", *snone = "  ";
    struct stem_t col = {0, 0}, *tail;

    // Print stems of branches coming further down
    for (tail = head; tail; tail = tail->next) {
        if (!tail->next) {
            if (!strcmp(sdown, tail->str))
                printf(" ├");
            else 
                printf("%s", tail->str);
            break;
        }
        printf("%s", tail->str);
    }
    
    if (root == NULL) {
        // Secure against null pointers sent as root
        printf("─(nil)\n");
        return;
    }
    printf("─%s", node_string[root->type]);
    if ( root->type == IDENTIFIER_DATA  ||
         root->type == STRING_DATA      ||
         root->type == EXPRESSION       ||
         root->type == RELATION)  
        printf("(%s)", (char *) root->data);
    else if (root->type == NUMBER_DATA)
        printf("(%ld)", *((int64_t *)root->data));
    putchar('\n');
 
    if (!root->n_children) return;
 
    if (tail && tail->str == slast)
        tail->str = snone;
 
    if (!tail)  tail = head = &col;
    else        tail->next = &col;

    for ( int64_t i=0; i < root->n_children; i++ ) {
        col.str = root->n_children - i - 1 ? sdown : slast;
        tree_print(root->children[i], head);
    }
    tail->next = 0;
}

/* Internal choices */
static void
node_print ( node_t *root, int nesting )
{
    if ( root != NULL )
    {
        printf ( "%*c%s", nesting, ' ', node_string[root->type] );
        if ( root->type == IDENTIFIER_DATA ||
             root->type == STRING_DATA ||
             root->type == EXPRESSION ) 
            printf ( "(%s)", (char *) root->data );
        else if ( root->type == NUMBER_DATA )
            printf ( "(%ld)", *((int64_t *)root->data) );
        putchar ( '\n' );
        for ( int64_t i=0; i<root->n_children; i++ )
            node_print ( root->children[i], nesting+1 );
    }
    else
        printf ( "%*c%p\n", nesting, ' ', root );
}


static void
node_finalize ( node_t *discard )
{
    if ( discard != NULL )
    {
        free ( discard->data );
        free ( discard->children );
        free ( discard );
    }
}


static void
destroy_subtree ( node_t *discard )
{
    if ( discard != NULL )
    {
        for ( uint64_t i=0; i<discard->n_children; i++ )
            destroy_subtree ( discard->children[i] );
        node_finalize ( discard );
    }
}


static void
flatten(node_t **simplified, node_t *root)
{
    /* This will flatten left-expanded lists */
    if (!root)
        return;

    /* Do this recursivly */
    for (int i = 0; i < root->n_children; i++)
        flatten(&root->children[i], root->children[i]);

    node_t **new_children, *result = root;
    switch (root->type)
    {
    case GLOBAL_LIST:
    case STATEMENT_LIST:
    case PRINT_LIST:
    case EXPRESSION_LIST:
    case VARIABLE_LIST:
    case DECLARATION_LIST:
        // Check if node have more than two children
        if (root->n_children < 2)
            break;
        
        result = root->children[0];
        result->n_children++;

        // Realloc the array of children to the new size 
        if (!(new_children = realloc(result->children, result->n_children * sizeof(node_t*))))
            break;
        // if successs, insert the new array
        result->children = new_children;

        // Insert child at the end
        result->children[result->n_children - 1] = root->children[1];

        node_finalize(root);
        break;
    }

    *simplified = result;
}


static void
prune_children(node_t **simplified, node_t *root) 
{
    if (!root)
        return;

    /* Do this recursivly */
    for (int i = 0; i < root->n_children; i++)
        prune_children(&root->children[i], root->children[i]);

    node_t *result = root;
    switch (root->type)
    {
    case PROGRAM:
    case GLOBAL:
    //case ARGUMENT_LIST:     // For this to work, need to change order of operations
    //case PARAMETER_LIST:    // For this to work, need to change order of operations
    //case VARIABLE_LIST:
    //case EXPRESSION_LIST:
    case DECLARATION:
    case STATEMENT:
    case PRINT_ITEM:
    case PRINT_STATEMENT:
        result = root->children[0];

        // The print_statement only contains a print_list, still need a print_statement.
        if (root->type == PRINT_STATEMENT)
            result->type = PRINT_STATEMENT;
        
        node_finalize(root);
        break;
    }

    *simplified = result;
}


static void
resolve_constant_expressions(node_t **simplified, node_t *root)
{
    if (!root)
        return;

    /* Do this recursivly */
    for (int i = 0; i < root->n_children; i++)
        resolve_constant_expressions(&root->children[i], root->children[i]);
    
    if (root->type != EXPRESSION)
        return;
    
    node_t *result = root;
    
    switch (root->n_children)
    {
    case 1:
        result = root->children[0];
        
        if (root->data && 
            result->type == NUMBER_DATA && 
            result->data)
        {
            switch (*((char*)root->data))
            {
            case '-':
                *((int64_t*)result->data) *= -1;
                break;
            case '~':
                *((int64_t*)result->data) = ~*((int64_t*)result->data);
                break;
            }
        }

        node_finalize(root);
        break;
    
    case 2:
        // Both children needs to be numbers to resolve constants
        if (root->children[0]->type == NUMBER_DATA && 
            root->children[1]->type == NUMBER_DATA)
        {
            // Check if children does not contain null pointers
            if (!root->children[0]->data)
                break;
            if (!root->children[1]->data)
                break;

            // Check if data field is not null pointer
            if (!root->data)
                break;
            
            result = root->children[0];
            int64_t
                *lhs = result->data,
                *rhs = root->children[1]->data;

            switch (*(char*)root->data)
            {
            /* Assignments */
            case '|': *lhs |= *rhs; break;
            case '^': *lhs ^= *rhs; break;
            case '&': *lhs &= *rhs; break;
            case '+': *lhs += *rhs; break;
            case '-': *lhs -= *rhs; break;
            case '*': *lhs *= *rhs; break;
            case '/': *lhs /= *rhs; break;
            }

            node_finalize(root->children[1]);
            node_finalize(root);
        }            
        
        
        break;
    }
    *simplified = result;
}


static void
resolve_constant_relations( node_t** simplified, node_t* root)
{
    if (!root)
        return;

    /* Do this recursivly */
    for (int i = 0; i < root->n_children; i++)
        resolve_constant_relations(&root->children[i], root->children[i]);
    
    if (root->type != RELATION)//|| root->type != RELATION)
        return;

    node_t *result = root;

    if (root->n_children != 2)
        return;

    // Both children must be constant numbers
    if (root->children[0]->type != NUMBER_DATA ||
        root->children[1]->type != NUMBER_DATA)
        return;

    // Check if children does not contain null pointers
    if (!root->children[0]->data)
        return;
    if (!root->children[1]->data)
        return;

    // Check if data field is not null pointer
    if (!root->data)
        return;

    result = root->children[0];
    int64_t
        *lhs = result->data,
        *rhs = root->children[1]->data;

    switch (*(char*)root->data)
    {
    /* Relations */
    case '=': *lhs = (*lhs == *rhs); break;
    case '<': *lhs = (*lhs <  *rhs); break;
    case '>': *lhs = (*lhs >  *rhs); break;
    }

    node_finalize(root->children[1]);
    node_finalize(root);

    *simplified = result;
}


static void
simplify_tree ( node_t **simplified, node_t *root )
{
    if (!root) 
        return;

    /*  
        Each of  the functions do their operations recursivly.
        This opens up for a lot more flexibility, like removing 
        variable list after it is flatten
    */
    flatten(&root, root);
    prune_children(&root, root);
    resolve_constant_expressions(&root, root);

    // The following is experimental, will resolve the constant relations
    resolve_constant_relations(&root, root);

    *simplified = root;
}