#include <vslc.h>

#define ASM(opcode, args...)    puts("\t"#opcode"\t"#args)
#define LABEL(label)            printf("_%s:\n", (char*)label)

/**Generate table of strings in a rodata section. */
void generate_stringtable ( void );
/**Declare global variables in a bss section */
void generate_global_variables ( void );
/**Generate function entry code
 * @param function symbol table entry of function */
void generate_function ( symbol_t *function );
/**Generate code for a node in the AST, to be called recursively from
 * generate_function
 * @param node root node of current code block */
static void generate_node ( node_t *node );
/**Initializes program (already implemented) */
void generate_main ( symbol_t *first );

#define MIN(a,b) (((a)<(b)) ? (a):(b))

static const char *record[6] = {
    "%rdi", "%rsi", "%rdx", "%rcx", "%r8", "%r9"
};

// Helper funcs for generating different nodes
static void generate_print(node_t* node);
static void generate_var_ident(node_t *node);
static void solve_expressions(node_t *node);

// Helper func for fetching all symbols in a table
static uint64_t fetch_symbols(tlhash_t* symbol_table, symbol_t*** symbol_list);

void
generate_program ( void )
{
    /* TODO: Emit assembly instructions for functions, function calls,
     * print statements and expressions.
     * The provided function 'generate_main' creates a program entry point
     * for the function symbol it is given as argument.
     */

    // TODO: Implement
    // - Generate code for all functions
    // - Generate main (function already implemented) by assigning either the
    //   function named main or the first function of the source file if no 
    //   main exists.
    generate_stringtable();
    generate_global_variables();

    symbol_t **global_list;
    uint64_t no_globals = fetch_symbols(global_names, &global_list);

    for (uint64_t g = 0; g < no_globals; g++)
    {
        if (global_list[g]->type == SYM_FUNCTION)
            generate_function(global_list[g]);
    }
}

void
generate_stringtable ( void )
{
	/* These can be used to emit numbers, strings and a run-time
	 * error msg. from main
	 */
    puts("# DATA SECTION");
	puts(".data");
	puts(".intout:\t.asciz \"\%ld \"");
	puts(".strout:\t.asciz \"\%s \"");
	puts(".errout:\t.asciz \"Wrong number of arguments\"");

    for (uint64_t s = 0; s < stringc; s++) 
    {
        printf(".STR%03ld:\t.asciz %s\n", s, string_list[s]);
    }
    putchar('\n');
}

void
generate_global_variables ( void )
{
    symbol_t **global_list;
    uint64_t no_globals = fetch_symbols(global_names, &global_list);

    puts("# GLOBAL VARIABLES");
    puts(".bss");
    puts(".align 8");

    for (uint64_t g = 0; g < no_globals; g++) {
        if (global_list[g]->type == SYM_GLOBAL_VAR)
            printf(".%s:\n", global_list[g]->name);
    }
    putchar('\n');
    free(global_list);

}

void
generate_function ( symbol_t *function )
{
    // TODO: Generate code for declaring and entering function, then generate its body

    printf("# func %s(nparams: %ld)\n", function->name, function->nparms);
    puts(".text");
    printf(".global _%s\n", function->name);
    LABEL(function->name);
    ASM(pushq, %rbp);
    ASM(movq, %rsp, %rbp);

    // Push params to stack
    for (int arg = 0; arg < MIN(6,function->nparms); arg++)
        printf("\tpushq\t%s\n", record[arg] );

    // How many local variables are inside function
    uint64_t no_locals = function->locals->size - function->nparms;

    // IF the stack alignment is not 16 bytes, 
    // add one now as all local var also is 0
    if ((MIN(6,function->nparms) + no_locals) % 2)
        ASM(pushq, $0);

    // Make room for the local vars
    while(no_locals--)
        ASM(pushq, $0);

    // Now the stack ptr should be 16 byte aligned.

    generate_node(function->node);

    putchar('\n');
}


void
generate_node ( node_t *node)
{
    // TODO: Generate code corresponding to node

    switch (node->type)
    {
    case ASSIGNMENT_STATEMENT:
        break;
    case ADD_STATEMENT:
        break;
    case SUBTRACT_STATEMENT:
        break;
    case MULTIPLY_STATEMENT:
        break;
    case DIVIDE_STATEMENT:
        break;
    case PRINT_STATEMENT:
        generate_print(node);
        break;
    case RETURN_STATEMENT:
        ASM(movq, %rbp, %rsp);
        ASM(popq, %rbp);
        ASM(ret);
        break;
    case NULL_STATEMENT:
        break;

    case IF_STATEMENT:
    case WHILE_STATEMENT:
        /* DO NOTHING YET */
        break;
    
    default:
        for (int c = 0; c < node->n_children; c++)
            generate_node(node->children[c]);
        break;
    }
}


void
generate_print(node_t* node)
{
    for (uint64_t p = 0; p < node->n_children; p++)
    {
        node_t *curr_print = node->children[p];

        // Push rdi and rsi to stack incase there are data in them
        //ASM(pushq, %rdi);
        //ASM(pushq, %rsi);
        switch (curr_print->type)
        {
        case EXPRESSION:
            ASM(movq, $.intout, %rdi);
            // Since this is a stack machine, rax allways 
            // contain the result of expression
            ASM(movq, %rax, %rsi);
            break;

        case STRING_DATA:
            ASM(movq, $.strout, %rdi);
            printf("\tmovq\t$.STR%03ld, %%rsi\n", *(uint64_t*)curr_print->data);
            break;

        case IDENTIFIER_DATA:
            ASM(movq, $.intout, %rdi);
            // Need to find the sequence number of var
            printf("\tmovq\t");
            generate_var_ident(curr_print);
            printf(", %%rsi\n");
            break;
        default:
            break;
        }
        ASM(call, printf);


        // Adds a newline
        ASM(movq, $'\n', %rdi);
        ASM(call, putchar);

        //ASM(popq, %rsi);
        //ASM(popq, %rdi);


    }
}

void
generate_var_ident(node_t *node)
{
    symbol_t *ident_sym = node->entry;
    switch (ident_sym->type)
    {
    case SYM_GLOBAL_VAR:
        printf("$.%s", ident_sym->name);
        break;

    case SYM_PARAMETER:
        // If it is a paramter is one of the first 6, seacrch below bp
        if (ident_sym->seq < 6)
            printf("%ld(%%rbp)", -8 * (ident_sym->seq + 1));
        else
            // This requires that the parameters on 
            // stack is in reversed order
            printf("%ld(%%rbp)", 8 * (ident_sym->seq - 6 + 1 ));
        break;

    case SYM_LOCAL_VAR:
        printf("%ld(%%rbp)", -8 * (ident_sym->seq + 1));
        break;
    }
}


void
solve_expressions(node_t *node)
{

}

/**Generates the main function with argument parsing and calling of our
 * main function (first, if no function is named main)
 * @param first Symbol table entry of our main function */
void
generate_main ( symbol_t *first )
{
    puts ( ".globl main" );
    puts ( ".section .text" );
    puts ( "main:" );
    puts ( "\tpushq   %rbp" );
    puts ( "\tmovq    %rsp, %rbp" );

    printf ( "\tsubq\t$1,%%rdi\n" );
    printf ( "\tcmpq\t$%zu,%%rdi\n", first->nparms );
    printf ( "\tjne\tABORT\n" );
    printf ( "\tcmpq\t$0,%%rdi\n" );
    printf ( "\tjz\tSKIP_ARGS\n" );

    printf ( "\tmovq\t%%rdi,%%rcx\n" );
    printf ( "\taddq $%zu, %%rsi\n", 8*first->nparms );
    printf ( "PARSE_ARGV:\n" );
    printf ( "\tpushq %%rcx\n" );
    printf ( "\tpushq %%rsi\n" );

    printf ( "\tmovq\t(%%rsi),%%rdi\n" );
    printf ( "\tmovq\t$0,%%rsi\n" );
    printf ( "\tmovq\t$10,%%rdx\n" );
    printf ( "\tcall\tstrtol\n" );

/*  Now a new argument is an integer in rax */

    printf ( "\tpopq %%rsi\n" );
    printf ( "\tpopq %%rcx\n" );
    printf ( "\tpushq %%rax\n" );
    printf ( "\tsubq $8, %%rsi\n" );
    printf ( "\tloop PARSE_ARGV\n" );

    /* Now the arguments are in order on stack */
    for (int arg = 0; arg < MIN(6,first->nparms); arg++)
        printf ( "\tpopq\t%s\n", record[arg] );

    printf ( "SKIP_ARGS:\n" );
    printf ( "\tcall\t_%s\n", first->name );
    printf ( "\tjmp\tEND\n" );
    printf ( "ABORT:\n" );
    printf ( "\tmovq\t$.errout, %%rdi\n" );
    printf ( "\tcall puts\n" );

    printf ( "END:\n" );
    puts ( "\tmovq    %rax, %rdi" );
    puts ( "\tcall    exit" );
}

static uint64_t 
fetch_symbols(tlhash_t* symbol_table, symbol_t*** symbol_list)
{
    uint64_t no_symbols = tlhash_size(symbol_table);
    *symbol_list = malloc(no_symbols * sizeof(symbol_t));
    tlhash_values(symbol_table, (void **)*symbol_list );

    return no_symbols;
}