#include "quads.h"

#include "ast.h"
#include "ast_print.h"
#include "asmgen.h"
#include "parser.tab.h" // for token binop values
#include "quads_cf.h"
#include "quads_print.h"
#include "symtab.h"
#include "util.h"
#include "types.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

unsigned temp_count = 0;

void gen_ret(astn n);

// allocate temporary
astn qtemp_alloc(unsigned size) {
    astn n = astn_alloc(ASTN_QTEMP);
    n->Qtemp.tempno = ++temp_count;
    n->Qtemp.size = size;
    sym f = st_lookup(cursor.fn, NS_MISC);
    //printf("found fn: "); st_dump_entry(f);
    f->fn_scope->stack_total += 4;
    n->Qtemp.stack_offset = f->fn_scope->stack_total;
    return n;
}

enum quad_op binop_to_quad_op(int binop) {
    switch (binop) {
        case '*': return Q_MUL;
        case '/': return Q_DIV;
        case '%': return Q_MOD;
        case '+': return Q_ADD;
        case '-': return Q_SUB;
        case SHL: return Q_SHL;
        case SHR: return Q_SHR;
        case '&': return Q_BWAND;
        case '^': return Q_BWXOR;
        case '|': return Q_BWOR;
        // condexpr stuff
        default: die("unhandled binop");
    }
    return -1;
}

// definitiely leaky
astn decay_array(astn a) {
    if (a->type == ASTN_SYMPTR) return decay_array(a->Symptr.e->type);
    if (a->type == ASTN_TYPE && a->Type.is_derived && a->Type.derived.type == t_ARRAY) {
        astn n = astn_alloc(ASTN_TYPE);
        n->Type.is_derived = true;
        n->Type.derived.type = t_PTR;
        n->Type.derived.target = a->Type.derived.target;
        return n;
    }
    return a;
}

// 
/*
astn check_ptr_binop(astn n) {
    if (n->type == ASTN_SYMPTR) return check_ptr_binop(n->Symptr.e->type);
    if (n->type == ASTN_TYPE && n->Type.is_derived) {
        if (n->Type.derived.type == t_ARRAY) return decay_array(n);
        else if (n->Type.derived.type == t_PTR) return n;
        else die("undefined derived type");
    } else {
        return n;
    }
    return NULL;
}
*/

void quad_error(const char* msg) {
    eprintf("quad generation error: %s\n", msg);
    exit(-1);
}

// triggers on either ptr or array
bool isptr(astn n) {
    if (n->type == ASTN_SYMPTR) return isptr(n->Symptr.e->type);
    return (n->type == ASTN_TYPE && n->Type.is_derived);
}

bool isarr(astn n) {
    if (n->type == ASTN_SYMPTR) return isarr(n->Symptr.e->type);
    return (n->type == ASTN_TYPE && n->Type.is_derived && n->Type.derived.type == t_ARRAY);
}

bool isfn(astn n) {
    return (n->type == ASTN_SYMPTR && n->Symptr.e->entry_type == STE_FN);
}

static astn const_alloc(int num) {
    astn n=astn_alloc(ASTN_NUM);
    n->Num.number.integer = num;
    n->Num.number.is_signed = true;
    n->Num.number.aux_type = s_INT;
    return n;
}

// works on arr or ptr
static astn ptr_target(astn n) {
    if (n->type == ASTN_SYMPTR) return ptr_target(n->Symptr.e->type);
    return n->Type.derived.target;
}

_Noreturn void todo(const char* msg) {
    eprintf("TODO: %s\n", msg);
    exit(-1);
}

void cond_rvalue(astn n, astn left, astn right, astn target);

// emit args in reverse
void put_args(astn head) {
    if (!head) return;
    put_args(list_next(head));
    emit(Q_ARG, gen_rvalue(list_data(head), NULL), NULL, NULL);
}

astn gen_rvalue(astn node, astn target) {
    astn temp;
    if (node->type == ASTN_FNCALL) {
        astn fn = gen_rvalue(node->Fncall.fn, NULL);

        emit(Q_ARGBEGIN, NULL, NULL, NULL);
        astn arg = node->Fncall.args;

        put_args(arg);
        /*
        while (arg) {
            emit(Q_ARG, gen_rvalue(list_data(arg), NULL), NULL, NULL);
            arg = list_next(arg);
        }
        */

        if (!target) target = qtemp_alloc(4);
        emit(Q_CALL, fn, NULL, target);
        return target;
    }
    if (node->type == ASTN_SYMPTR) {
        //struct astn_type *type = &node->Symptr.e->type->Type;
        //if (type->is_derived && type->derived.type == t_ARRAY) {
        if (isarr(node)) {
            astn temp = qtemp_alloc(4); // address type
            emit(Q_LEA, node, NULL, temp);
            return temp;
        }
        else if (isfn(node)) {
            astn temp = astn_alloc(ASTN_IDENT);
            temp->Ident.ident = node->Symptr.e->ident;
            return temp;
        }
        else {
            target = qtemp_alloc(4); // hardcoded
            emit(Q_MOV, node, NULL, target);
            return target;
        }
    }
    if (node->type == ASTN_NUM) return node;
    if (node->type == ASTN_STRLIT) return node;
    if (node->type == ASTN_BINOP) {
        astn left = node->Binop.left;
        astn right = node->Binop.right;

        bool l_isptr = isptr(left);
        bool r_isptr = isptr(right);       

        //decay_array(left);
        //decay_array(right);

        if (l_isptr && r_isptr)
            quad_error("can't add two pointers!\n"); // check for subtraction!!

        if (r_isptr) { // swap them if needed
            temp = right;
            right = left;
            left = temp;
            l_isptr = true; r_isptr = false;
        }

        astn rval_left = gen_rvalue(left, NULL);
        astn rval_right = gen_rvalue(right, NULL);

        if (node->Binop.op == '+') {
            if (l_isptr) { // ptr + int
                if (rval_right->type == ASTN_NUM) { // otherwise we'd emit a constant as an lvalue
                    temp = qtemp_alloc(4);
                    emit(Q_MUL, rval_right, const_alloc(get_sizeof(ptr_target(left))), temp);
                    rval_right = temp;
                }
                else
                    emit(Q_MUL, rval_right, const_alloc(get_sizeof(ptr_target(left))), rval_right);
            }
        }

        if (!target) target=qtemp_alloc(4); // hardcoded int!

        if (node->Binop.op == '<' || node->Binop.op == '>' 
            || node->Binop.op == LTEQ || node->Binop.op == GTEQ ||
            node->Binop.op == EQEQ || node->Binop.op == NOTEQ) {
            cond_rvalue(node, rval_left, rval_right, target);
        } else
            emit(binop_to_quad_op(node->Binop.op), rval_left, rval_right, target);
        return target;
    }
    if (node->type == ASTN_UNOP) {
        astn utarget = node->Unop.target;

        switch (node->Unop.op) {
            case '*':
                //printf("dumping utarget\n"); print_ast(utarget);

                // multidim arrays: needs fixing
                if (isarr(utarget))
                {   eprintf("yup\n");
                    return gen_rvalue(ptr_target(utarget), target);
                }
                ;
                astn addr = gen_rvalue(utarget, NULL);

                if (!target) target = qtemp_alloc(4);
                emit(Q_LOAD, addr, NULL, target);
                return target;

            astn n;
            case MINUSMINUS:
                temp = gen_rvalue(utarget, target);
                n=astn_alloc(ASTN_NUM);
                n->Num.number.integer=1;
                n->Num.number.aux_type = s_INT;
                gen_assign(cassign_alloc('-', utarget, n));
                return temp;
            case PLUSPLUS:
                temp = gen_rvalue(utarget, target);
                n=astn_alloc(ASTN_NUM);
                n->Num.number.integer=1;
                n->Num.number.aux_type = s_INT;
                gen_assign(cassign_alloc('+', utarget, n));
                return temp;
            case '&':
                temp = qtemp_alloc(4); // address type
                emit(Q_LEA, utarget, NULL, temp);
                return temp;
            case '+': // no-op for now
                return utarget;
            case '-':
                temp = qtemp_alloc(4);
                emit(Q_NEG, gen_rvalue(utarget, NULL), NULL, temp);
                return temp;
            case '!':           todo("NOT");
            case '~':
                temp = qtemp_alloc(4);
                emit(Q_BWNOT, gen_rvalue(utarget, NULL), NULL, temp);
                return temp;
            default:
                die("unhandled unop");
        }
    }
    if (node->type == ASTN_ASSIGN) {
        gen_assign(node);
        return gen_rvalue(node->Assign.left, target);
    }
    eprintf("unhandled node type for gen_rvalue: %d\n", node->type);
    die("FUCK!");
    return NULL;
}

astn gen_lvalue(astn node, enum addr_modes *mode) {
    if (node->type == ASTN_SYMPTR) { *mode = MODE_DIRECT; return node; }
    if (node->type == ASTN_NUM) return NULL;
    if (node->type == ASTN_UNOP && node->Unop.op == '*') {
        *mode = MODE_INDIRECT;
        return gen_rvalue(node->Unop.target, NULL);
    }
    return NULL;
}


void gen_assign(astn node) {
    enum addr_modes destmode;
    astn dest = gen_lvalue(node->Assign.left, &destmode);
    if (!dest) {
        eprintf("Error: not an lvalue: "); die("test");
        print_ast(node->Assign.left);
        exit(-1);
    }
    if (destmode == MODE_DIRECT) {
        astn temp = gen_rvalue(node->Assign.right, NULL);
        emit(Q_MOV, temp, NULL, dest);
    } else {
        astn temp = gen_rvalue(node->Assign.right, NULL);
        emit(Q_STORE, temp, NULL, dest);
    }
}

void gen_mov_bb(astn target, int val) {
    astn aval = const_alloc(val);
    emit(Q_MOV, aval, NULL, target);
}

void cond_rvalue(astn n, astn left, astn right, astn target) {
    BB* one = bb_alloc();
    BB* zero = bb_alloc();

    emit(Q_CMP, left, right, NULL);
    emit_branch(n->Binop.op, one, zero);
    
    BB *future = bb_alloc();

    current_bb = one;
    gen_mov_bb(target, 1);
    uncond_branch(future);
    current_bb = zero;
    gen_mov_bb(target, 0);
    uncond_branch(future);

    //astn cond = astn_alloc(ASTN_BINOP);
    current_bb = future;
}

void emit(enum quad_op op, astn src1, astn src2, astn target) {
    quad *q = safe_calloc(1, sizeof(quad));
    q->op = op;
    q->src1 = src1;
    q->src2 = src2;
    q->target = target;
    if (!current_bb->start) {
        current_bb->start = q;
        current_bb->cur = q;
    } else {
        current_bb->cur->next = q;
        q->prev = current_bb->cur;
        current_bb->cur = q;
    }
    //print_quad(q);
}

quad* last_in_bb(BB* b) {
    quad *q = b->start;
    if (!q) return NULL;
    while (q->next) q=q->next;
    return q;
}

void gen_fn(sym e) {
    bb_count = 0;
    cursor.fn = e->ident;

    BB *f = bb_alloc();
    current_bb = f;
    emit(Q_FNSTART, NULL, NULL, NULL);

    astn s = e->body;
    gen_quads(s);
    if (!last_in_bb(current_bb) || last_in_bb(current_bb)->op != Q_RET) {
        if (e->type->Type.is_derived || e->type->Type.scalar.type != t_VOID) {
            if (!strcmp(e->ident, "main")) {
                astn r = astn_alloc(ASTN_RETURN);
                r->Return.ret = const_alloc(0);
                gen_ret(r);
            } else {
                emit(Q_RET, NULL, NULL, NULL);
            }
        }
    }

    f->stack_offset_ez = e->fn_scope->stack_total; // sorry so ugly
    //print_bbs();
    //asmgen(f);
}

void gen_ret(astn n) {
    sym e = st_lookup_fq(cursor.fn, &root_symtab, NS_MISC);
    bool non_void = (e->type->Type.is_derived || e->type->Type.scalar.type != t_VOID);
    if (n->Return.ret) {
        if (!non_void)
            eprintf("Warning: 'return' with a value in a void function\n");
        emit(Q_RET, gen_rvalue(n->Return.ret, NULL), NULL, NULL);
    }
    else {
        if (non_void)
            eprintf("Warning: 'return' without a value in a nonvoid function\n");
        emit(Q_RET, NULL, NULL, NULL);
    }
}

void gen_quads(astn n) {
    switch (n->type) {
        case ASTN_LIST:
            if (!list_data(n))
                break;
            while (n) {
                gen_quads(list_data(n));
                n = list_next(n);
            }
            break;
        case ASTN_ASSIGN:       gen_assign(n);        break;
        case ASTN_IFELSE:       gen_if(n);            break;
        case ASTN_WHILELOOP:
            if (n->Whileloop.is_dowhile)
                gen_dowhile(n);
            else
                gen_while(n);
            break;
        case ASTN_RETURN:       gen_ret(n);           break;
        case ASTN_CONTINUE:     uncond_branch(cursor.cont); break; // may leave behind unreachable code, optimizer would get rid of it
        case ASTN_BREAK:        uncond_branch(cursor.brk);  break;// may leave behind unreachable code, optimizer would get rid of it
        case ASTN_SWITCH:       todo("switch statements");  break;
        case ASTN_FORLOOP:      gen_for(n);           break;
        case ASTN_DECLREC:
            if (n->Declrec.init) {
                if (n->Declrec.e->scope == &root_symtab) {
                    todo("global variable initialization");
                    // ! this doesn't work; we don't even get a declrec right

                    // ! now (it seems) for globals
                }
                // only inside functions one sec
                astn i = astn_alloc(ASTN_ASSIGN);
                astn s = astn_alloc(ASTN_SYMPTR);
                s->Symptr.e = n->Declrec.e;   
                i->Assign.left = s;
                i->Assign.right = n->Declrec.init;
                gen_assign(i);
            }
            break;

        case ASTN_FNCALL:
        case ASTN_UNOP:
        case ASTN_BINOP:
            gen_rvalue(n, NULL);
            break;

        case ASTN_NUM:
        case ASTN_STRLIT:
            eprintf("Warning: useless constant statement: ");
            print_ast(n);
            break;
        case ASTN_NOOP:
            break;
        
        default:
            eprintf("skipping unknown astn for quads %d\n", n->type);
    }
}
