#include "ccc.h"
#include "codegen.h"
#include "scope.h"
#include "register.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define CGEN_PANIC(format, ...) {\
    fprintf(\
        stderr,\
        "ccc: code gen error: " format "\n" __VA_OPT__(,)\
        __VA_ARGS__);\
    exit(1);\
}

struct lval_def {
    struct storage_location loc;
    unsigned long long sz;
};

static const struct storage_location RV_LOC = {
    .type = STO_REG,
    .reg = &RAX,
};
static const struct storage_location MULDIV_LOC = RV_LOC;
static const struct storage_location MULDIV_OVERFLOW_LOC = {
    .type = STO_REG,
    .reg = &RDX,
};

#define RETURN_LABEL_FMT "%s@coda"
#define FULL_REG_SZ 8

static struct scope* scope;
static const struct fn_decl_node* active_fn;

static struct lval_def allocate_register(unsigned long long sz) {
    return (struct lval_def) {
        .loc = {
            .type = STO_REG,
            .reg = &RAX, /* TODO: no real register coloring happening LOL */
        },
        .sz = sz,
    };
}

static struct lval_def allocate_stack(FILE* outfile, unsigned long long sz) {
    fprintf(outfile, "\tsub rsp, %llu\n", sz);
    scope->bp_offset += sz;
    return (struct lval_def) {
        .loc = {
            .type = STO_STACK,
            .bp_offset = scope->bp_offset,
        },
        .sz = sz,
    };
}

static struct lval_def allocate_temporary(
    FILE* outfile,
    unsigned long long sz
) {
    return allocate_stack(outfile, sz);
}

static void deallocate_temporary(FILE* outfile, const struct lval_def* tmp) {
    if (tmp->loc.type == STO_STACK) {
        fprintf(outfile, "\tadd rsp, %llu\n", tmp->sz);
        scope->bp_offset -= tmp->sz;
    }
}

static void emit_storage_loc(
    FILE* outfile,
    const struct storage_location* loc,
    unsigned long long sz
) {
    switch (loc->type) {
        case STO_LABEL:
            fprintf(outfile, "%s", loc->label);
            break;
        case STO_FN:
            fprintf(outfile, "%s", loc->decl->name);
            break;
        case STO_REG:
            if (sz > 4) fprintf(outfile, "%s", loc->reg->qword);
            else if (sz > 2) fprintf(outfile, "%s", loc->reg->dword);
            else if (sz > 1) fprintf(outfile, "%s", loc->reg->word);
            else fprintf(outfile, "%s", loc->reg->byte);
            break;
        case STO_STACK:
            if (loc->bp_offset < 0)
                fprintf(outfile, "[rbp + %lld]", -loc->bp_offset);
            else if (loc->bp_offset > 0)
                fprintf(outfile, "[rbp - %lld]", loc->bp_offset);
            else
                fprintf(outfile, "[rbp]");
            break;
        case STO_IMM:
            fprintf(outfile, "%llu", loc->value);
            break;
    }
}

static bool locs_equal(
    const struct storage_location* a,
    const struct storage_location* b
) {
    if (a->type != b->type) return false;

    switch (a->type) {
        case STO_IMM:
            return a->value == b->value;
        case STO_REG:
            return strcmp(a->reg->qword, b->reg->qword) == 0;
        case STO_STACK:
            return a->bp_offset == b->bp_offset;
        case STO_LABEL:
            return strcmp(a->label, b->label) == 0;
        case STO_FN:
            return a->decl == b->decl;
    }
    CGEN_PANIC("unhandled storage type case");
}

static void emit_size_const(FILE* outfile, unsigned long long sz) {
    if (sz > 4) fprintf(outfile, "qword ");
    else if (sz > 2) fprintf(outfile, "dword ");
    else if (sz > 1) fprintf(outfile, "word ");
    else fprintf(outfile, "byte ");
}

static void emit_mov(
    FILE* outfile,
    const struct lval_def* dst,
    const struct storage_location* src
) {
    /* first optimization: if dst == src, emit nothing */
    if (locs_equal(&dst->loc, src)) return;

    switch (dst->loc.type) {
        case STO_REG:
            if (src->type == STO_REG && dst->sz < 4) {
                fprintf(outfile, "\tmovzx ");
                emit_storage_loc(outfile, &dst->loc, FULL_REG_SZ);
            } else {
                fprintf(outfile, "\tmov ");
                emit_storage_loc(outfile, &dst->loc, dst->sz);
            }

            fprintf(outfile, ", ");
            emit_storage_loc(outfile, src, dst->sz);
            break;
        case STO_STACK:
            if (src->type == STO_STACK) {
                /* `mov mem, mem` is illegal in x86_64 */
                struct lval_def tmp = allocate_register(dst->sz);
                emit_mov(outfile, &tmp, src);
                emit_mov(outfile, dst, &tmp.loc);
                return;
            }

            fprintf(outfile, "\tmov ");
            if (src->type == STO_IMM) emit_size_const(outfile, dst->sz);

            emit_storage_loc(outfile, &dst->loc, dst->sz);
            fprintf(outfile, ", ");
            emit_storage_loc(outfile, src, dst->sz);
            break;
        case STO_LABEL:
            CGEN_PANIC("can't move value into label %s", dst->loc.label);
        case STO_FN:
            CGEN_PANIC(
                "can't move value into function %s", dst->loc.decl->name);
        case STO_IMM:
            CGEN_PANIC(
                "can't move value into immediate value %lld", dst->loc.value);
    }
    fprintf(outfile, "\n");
}

/* TODO: move this utility to the type checker and use evaluated_type */
static unsigned long long get_type_size(const struct type_node* type) {
    if (type->ptr_level > 0) return PTR_SIZE;
    return type->def->sz;
}

static inline struct lval_def make_stack_lval(
    FILE* outfile,
    const struct type_node* type
) {
    return allocate_stack(outfile, get_type_size(type));
}

static void emit_expr(
    FILE* outfile,
    const struct expr_node* node,
    const struct lval_def* dst);

static void emit_int_lit(
    FILE* outfile,
    const struct int_lit_node* node,
    const struct lval_def* dst
) {
    if (dst != NULL)
        emit_mov(
            outfile,
            dst,
            &(struct storage_location) {
                .type = STO_IMM,
                .value = node->val,
            });
}

static void emit_float_lit(
    FILE* outfile,
    const struct float_lit_node* node,
    const struct lval_def* dst
) {
    if (dst != NULL) {
        CGEN_PANIC("float literals are not implemented");        
    }
}

static void emit_char_lit(
    FILE* outfile,
    const struct char_lit_node* node,
    const struct lval_def* dst
) {
    if (dst != NULL) {
        emit_mov(
            outfile,
            dst,
            &(struct storage_location) {
                .type = STO_IMM,
                .value = (unsigned long long) node->val
            });
    }
}

static void emit_str_lit(
    FILE* outfile,
    const struct str_lit_node* node,
    const struct lval_def* dst
) {
    if (dst != NULL) {
        CGEN_PANIC("string literals are not implemented");
    }
}

static struct var_def get_var(const char* name) {
    struct var_def var_def;
    if (!scope_get_var(scope, &var_def, name))
        CGEN_PANIC("reference to undefined variable %s", name);
    return var_def;
}

static void emit_var_ref(
    FILE* outfile,
    const struct var_ref_node* node,
    const struct lval_def* dst
) {
    if (dst != NULL) {
        struct var_def var_def = get_var(node->ident);
        emit_mov(outfile, dst, &var_def.loc);
    }
}

static void emit_stmt(FILE* outfile, const struct stmt_node* node);

static struct var_def emit_var_decl(
    FILE* outfile,
    const struct var_decl_node* node
) {
    /* TODO: type checker should define vars, we just set their loc */
    struct lval_def var_dst = make_stack_lval(outfile, &node->type);
    struct var_def var_def = {
        .name = node->ident,
        .loc = var_dst.loc,
        .sz = var_dst.sz,
    };
    scope_define_var(scope, var_def);
    return var_def;
}

static struct lval_def emit_lval(
    FILE* outfile,
    const struct lval_node* node
) {
    struct var_def var_def;
    switch (node->type) {
        case LVAL_VAR_DECL:
            var_def = emit_var_decl(outfile, &node->inner.var_decl);
            return (struct lval_def) {.loc = var_def.loc, .sz = var_def.sz};
        case LVAL_VAR_REF:
            var_def = get_var(node->inner.var_ref.ident);
            return (struct lval_def) {.loc = var_def.loc, .sz = var_def.sz};
    }
    CGEN_PANIC("unknown lval type: %d", node->type);
}

static void emit_assignment(
    FILE* outfile,
    const struct assign_node* node,
    const struct lval_def* dst
) {
    const struct lval_def lval_def = emit_lval(outfile, &node->lval);
    emit_expr(outfile, node->rval, &lval_def);
    if (dst != NULL) emit_mov(outfile, dst, &lval_def.loc);
}

static void emit_call(
    FILE* outfile,
    const struct call_node* node,
    const struct lval_def* dst
) {
    unsigned long long orig_bp_offset = scope->bp_offset;
    unsigned long long arg_bp_offset = orig_bp_offset;
 
    struct args_decl_node* arg_decl = node->called_fn_ref->args;
    struct args_eval_node* arg_eval = node->args;
    while (arg_decl != NULL && arg_eval != NULL) {
        struct lval_def arg_dst =
            make_stack_lval(outfile, &arg_decl->decl->type);
        emit_expr(outfile, arg_eval->expr, &arg_dst);

        arg_decl = arg_decl->next;
        arg_eval = arg_eval->next;
    }

    if (arg_decl != NULL)
        CGEN_PANIC(
            "too many arguments to function %s",
            node->called_fn_ref->name);
    if (arg_eval != NULL)
        CGEN_PANIC(
            "missing arguments to function %s",
            node->called_fn_ref->name);

    unsigned char arg_regnum = 0;
    arg_decl = node->called_fn_ref->args;
    while (arg_decl != NULL) {
        unsigned long long type_sz = get_type_size(&arg_decl->decl->type);
        arg_bp_offset += type_sz;

        struct lval_def arg_dst;
        if (arg_regnum < CC_N_REGS)
            arg_dst = (struct lval_def) {
                .loc = (struct storage_location) {
                    .type = STO_REG,
                    .reg = CALLING_CONV[arg_regnum++],
                },
                .sz = type_sz,
            };
        else
            arg_dst = make_stack_lval(outfile, &arg_decl->decl->type);

        emit_mov(outfile, &arg_dst, &(struct storage_location) {
            .type = STO_STACK,
            .bp_offset = arg_bp_offset,
        });
        arg_decl = arg_decl->next;
    }

    fprintf(outfile, "\tcall %s\n", node->called_fn_ref->name);
    if (dst != NULL) {
        if (get_type_size(&node->called_fn_ref->return_type) == 0)
            CGEN_PANIC("can't assign the result of a void function");

        emit_mov(outfile, dst, &RV_LOC);
    }

    /* mass-pop our argument temporaries off the stack */
    scope->bp_offset = orig_bp_offset;
    if (orig_bp_offset > 0)
        fprintf(outfile, "\tlea rsp, [rbp - %llu]\n", orig_bp_offset);
    else
        fprintf(outfile, "\tmov rsp, rbp\n");
}

static void emit_unary(
    FILE* outfile,
    const struct unary_node* node,
    const struct lval_def* dst
) {
    emit_expr(outfile, node->expr, dst);
    if (dst == NULL) return;

    switch (node->op) {
        case UNARY_NEG:
            fprintf(outfile, "\tneg ");
            if (dst->loc.type == STO_STACK) emit_size_const(outfile, dst->sz);
            emit_storage_loc(outfile, &dst->loc, dst->sz);
            fprintf(outfile, "\n");
            break;
    }
}

/* TODO: tighten up/enforce with regard to up-casting smaller operands */
static void emit_binary(
    FILE* outfile,
    const struct binary_node* node,
    const struct lval_def* dst
) {
    if (dst == NULL) {
        emit_expr(outfile, node->lhs, NULL);
        emit_expr(outfile, node->rhs, NULL);
        return;
    }

    /* LHS goes in RAX explicitly because imul and idiv are weird */
    struct lval_def rhs_dst = allocate_temporary(outfile, dst->sz);
    emit_expr(outfile, node->rhs, &rhs_dst);
    struct lval_def lhs_dst = (struct lval_def) {
        .loc = MULDIV_LOC,
        .sz = dst->sz,
    };
    emit_expr(outfile, node->lhs, &lhs_dst);

    switch (node->op) {
        case BINARY_ADD:
            fprintf(outfile, "\tadd ");
            emit_storage_loc(outfile, &lhs_dst.loc, dst->sz);
            fprintf(outfile, ", ");
            break;
        case BINARY_SUB:
            fprintf(outfile, "\tsub ");
            emit_storage_loc(outfile, &lhs_dst.loc, dst->sz);
            fprintf(outfile, ", ");
            break;
        case BINARY_MUL:
            fprintf(outfile, "\timul ");
            if (rhs_dst.loc.type == STO_STACK)
                emit_size_const(outfile, dst->sz);
            break;
        case BINARY_DIV:
            /* nothing in the top half reg */
            fprintf(outfile, "\txor ");
            emit_storage_loc(outfile, &MULDIV_OVERFLOW_LOC, FULL_REG_SZ);
            fprintf(outfile, ", ");
            emit_storage_loc(outfile, &MULDIV_OVERFLOW_LOC, FULL_REG_SZ);
            fprintf(outfile, "\n");

            fprintf(outfile, "\tidiv ");
            if (rhs_dst.loc.type == STO_STACK)
                emit_size_const(outfile, dst->sz);
            break;
    }

    emit_storage_loc(outfile, &rhs_dst.loc, dst->sz);
    fprintf(outfile, "\n");

    /* TODO: deal with RDX overflow shit for imul and idiv */
    emit_mov(outfile, dst, &lhs_dst.loc);

    deallocate_temporary(outfile, &lhs_dst);
    deallocate_temporary(outfile, &rhs_dst);
}

static void emit_expr(
    FILE* outfile,
    const struct expr_node* node,
    const struct lval_def* dst
) {
    switch (node->type) {
        case EXPR_INT_LIT:
            emit_int_lit(outfile, &node->inner.int_lit, dst);
            break;
        case EXPR_FLOAT_LIT:
            emit_float_lit(outfile, &node->inner.float_lit, dst);
            break;
        case EXPR_CHAR_LIT:
            emit_char_lit(outfile, &node->inner.char_lit, dst);
            break;
        case EXPR_STR_LIT:
            emit_str_lit(outfile, &node->inner.str_lit, dst);
            break;
        case EXPR_VAR_REF:
            emit_var_ref(outfile, &node->inner.var_ref, dst);
            break;
        case EXPR_ASSIGN:
            emit_assignment(outfile, &node->inner.assign, dst);
            break;
        case EXPR_CALL:
            emit_call(outfile, &node->inner.call, dst);
            break;
        case EXPR_UNARY:
            emit_unary(outfile, &node->inner.unary, dst);
            break;
        case EXPR_BINARY:
            emit_binary(outfile, &node->inner.binary, dst);
            break;
    }
}

static void emit_return(FILE* outfile, const struct return_node* node) {
    if (active_fn == NULL) CGEN_PANIC("must be inside a function to return");
    unsigned long long return_type_sz = get_type_size(&active_fn->return_type);

    if (node->ret_val != NULL) {
        if (return_type_sz == 0)
            CGEN_PANIC(
                "returning a value from void function %s", active_fn->name);

        emit_expr(
            outfile,
            node->ret_val,
            &(struct lval_def) {
                .loc = RV_LOC,
                .sz = return_type_sz,
            });
    } else if (return_type_sz > 0) {
        CGEN_PANIC(
            "non-void function %s should return a value", active_fn->name);
    }

    fprintf(outfile, "\tjmp " RETURN_LABEL_FMT "\n", active_fn->name);
}

static void emit_group(FILE* outfile, const struct group_node* node) {
    const struct stmt_node* body_node = node->head;
    while (body_node != NULL) {
        emit_stmt(outfile, body_node);
        body_node = body_node->next;
    }
}

static void emit_stmt_group(FILE* outfile, const struct group_node* node) {
    if (node->scope->next_out != scope) CGEN_PANIC("scopes are misaligned");

    scope = node->scope;
    scope->bp_offset = scope->next_out->bp_offset; /* don't reset bp */

    emit_group(outfile, node);

    /* don't reset sp because alloca needs to work */
    scope->next_out->bp_offset = scope->bp_offset;
    scope = scope->next_out;
}

static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
    switch (node->type) {
        case STMT_EMPTY:
            break;
        case STMT_VAR_DECL:
            emit_var_decl(outfile, &node->inner.var_decl);
            break;
        case STMT_RETURN:
            emit_return(outfile, &node->inner.return_);
            break;
        case STMT_EXPR:
            emit_expr(outfile, &node->inner.expr, NULL);
            break;
        case STMT_GROUP:
            emit_stmt_group(outfile, &node->inner.group);
            break;
    }
}

static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
    if (active_fn != NULL)
        CGEN_PANIC(
            "can't define function %s inside function %s",
            node->name,
            active_fn->name);
    active_fn = node;

    fprintf(outfile, "%s:\n", node->name);
    fprintf(outfile, "\tpush rbp\n");
    fprintf(outfile, "\tmov rbp, rsp\n");

    if (node->scope->next_out != scope) CGEN_PANIC("scopes are misaligned");
    scope = node->scope;
    scope->bp_offset = 0;

    long long spilled_bp_ofs = -16; // return address + old bp
    unsigned char arg_regnum = 0;
    struct args_decl_node* arg_decl = node->args;
    while (arg_decl != NULL) {
        struct lval_def arg_dst =
            make_stack_lval(outfile, &arg_decl->decl->type);
        scope_define_var(
            scope,
            (struct var_def) {
                .name = arg_decl->decl->ident,
                .loc = arg_dst.loc,
                .sz = arg_dst.sz,
            });

        struct storage_location arg_src;
        if (arg_regnum < CC_N_REGS) {
            arg_src = (struct storage_location) {
                .type = STO_REG,
                .reg = CALLING_CONV[arg_regnum++]
            };
        } else {
            arg_src = (struct storage_location) {
                .type = STO_STACK,
                .bp_offset = spilled_bp_ofs,
            };
            spilled_bp_ofs -= arg_dst.sz;
        }
        emit_mov(outfile, &arg_dst, &arg_src);

        arg_decl = arg_decl->next;
    }

    emit_group(outfile, &node->body);

    scope = scope->next_out;

    fprintf(outfile, RETURN_LABEL_FMT ":\n", node->name);
    fprintf(outfile, "\tmov rsp, rbp\n");
    fprintf(outfile, "\tpop rbp\n");
    fprintf(outfile, "\tret\n");
    active_fn = NULL;
}

static void emit_root_node(FILE* outfile, const struct root_node* node) {
    switch (node->type) {
        case ROOT_FN_DECL:
            emit_fn_decl(outfile, &node->inner.fn_decl);
            break;
    }
}

void emit_code(struct ast* ast, const char* path) {
    FILE* outfile = fopen(path, "w");
    if (outfile == NULL) CCC_PANIC;

    scope = ast->root_scope;
    fprintf(outfile, "section .text\n");

    /* output all function declarations in the root scope as globals */
    const struct root_node* node = ast->root_node;
    for (; node != NULL; node = node->next) {
        if (node->type != ROOT_FN_DECL) continue;
        fprintf(outfile, "global %s\n", node->inner.fn_decl.name);
    }
    fprintf(outfile, "\n");

    /* actual code body */
    node = ast->root_node;
    while (node != NULL) {
        emit_root_node(outfile, node);
        if (node->next != NULL) fprintf(outfile, "\n");
        node = node->next;
    }

    fclose(outfile);
}