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

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

static const struct storage_location RV_LOC = {
    .type = REGISTER,
    .reg = &RAX,
};
#define FULL_REG_SZ 8

static struct scope* scope;

static void emit_storage_loc(
    FILE* outfile,
    const struct storage_location* loc,
    unsigned long long sz
) {
    switch (loc->type) {
        case REGISTER:
            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 JMP_LABEL:
            fprintf(outfile, "%s", loc->label);
            break;
        case BP_OFFSET:
            if (loc->offset < 0)
                fprintf(outfile, "[rbp + %lld]", -loc->offset);
            else if (loc->offset > 0)
                fprintf(outfile, "[rbp - %lld]", loc->offset);
            else
                fprintf(outfile, "[rbp]");
            break;
        case IMMEDIATE:
            fprintf(outfile, "%llu", loc->value);
            break;
    }
}

static void emit_mov(
    FILE* outfile,
    const struct storage_location* dst,
    const struct storage_location* src,
    unsigned long long sz
) {
    switch (dst->type) {
        case REGISTER:
            if (src->type == REGISTER && sz < 4) {
                fprintf(outfile, "\tmovzx ");
                emit_storage_loc(outfile, dst, FULL_REG_SZ);
            } else {
                fprintf(outfile, "\tmov ");
                emit_storage_loc(outfile, dst, sz);
            }

            fprintf(outfile, ", ");
            emit_storage_loc(outfile, src, sz);
            break;
        case BP_OFFSET:
            if (src->type == BP_OFFSET) {
                /* `mov mem, mem` is illegal in x86_64 */
                emit_mov(outfile, &RV_LOC, src, sz);
                emit_mov(outfile, dst, &RV_LOC, sz);
                return;
            }

            fprintf(outfile, "\tmov ");
            emit_storage_loc(outfile, dst, sz);
            fprintf(outfile, ", ");
            emit_storage_loc(outfile, src, sz);
            break;
        case JMP_LABEL:
            CGEN_PANIC("can't move value into label %s", dst->label);
        case IMMEDIATE:
            CGEN_PANIC(
                "can't move value into immediate value %lld", dst->value);
    }
    fprintf(outfile, "\n");
}

static void emit_expr(
    FILE* outfile,
    const struct expr_node* node,
    const struct storage_location* storage
);

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

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 storage_location* storage
) {
    if (storage != NULL) {
        struct var_def var_def = get_var(node->ident);
        emit_mov(outfile, storage, &var_def.loc, var_def.sz);
    }
}

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

static unsigned long long get_type_size(const struct type_node* type) {
    if (type->ptr_level > 0) return PTR_SIZE;

    struct type_def type_def;
    if (!scope_get_type(scope, &type_def, type->name))
        CGEN_PANIC("size of type %s is not known", type->name);

    return type_def.size;
}

static struct var_def emit_var_decl(
    FILE* outfile,
    const struct var_decl_node* node
) {
    unsigned long long type_sz = get_type_size(&node->type);

    fprintf(outfile, "\tsub rsp, %llu\n", type_sz);
    scope->bp_offset += type_sz;

    struct var_def var_def = {
        .name = node->ident,
        .loc = {
            .type = BP_OFFSET,
            .offset = scope->bp_offset,
        },
        .sz = type_sz,
    };
    scope_define_var(scope, var_def);
    return var_def;
}

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

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->as._var_decl);
            return (struct lval_def) {.loc = var_def.loc, .sz = var_def.sz};
        case LVAL_VAR_REF:
            var_def = get_var(node->as._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 storage_location* storage
) {
    const struct lval_def lval_def = emit_lval(outfile, &node->lval);
    emit_expr(outfile, node->rval, &lval_def.loc);
    if (storage != NULL) emit_mov(outfile, storage, &lval_def.loc, lval_def.sz);
}

static void emit_expr(
    FILE* outfile,
    const struct expr_node* node,
    const struct storage_location* storage
) {
    switch (node->type) {
        case EXPR_INT_LIT:
            emit_int_lit(outfile, &node->as._int_lit, storage);
            break;
        case EXPR_VAR_REF:
            emit_var_ref(outfile, &node->as._var_ref, storage);
            break;
        case EXPR_ASSIGN:
            emit_assignment(outfile, &node->as._assign, storage);
    }
}

static void emit_return(FILE* outfile, const struct return_node* node) {
    if (node->ret_val != NULL) emit_expr(outfile, node->ret_val, &RV_LOC);
    fprintf(outfile, "\tmov rsp, rbp\n");
    fprintf(outfile, "\tpop rbp\n");
    fprintf(outfile, "\tret\n");
}

static void emit_group(FILE* outfile, const struct group_node* node) {
    const struct stmt_node* body_node = node->body_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) {
    scope_push(&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_pop(&scope);
}

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->as._var_decl);
            break;
        case STMT_RETURN:
            emit_return(outfile, &node->as._return);
            break;
        case STMT_EXPR:
            emit_expr(outfile, &node->as._expr, NULL);
            break;
        case STMT_GROUP:
            emit_stmt_group(outfile, &node->as._group);
            break;
    }
}

static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
    fprintf(outfile, "%s:\n", node->name);
    fprintf(outfile, "\tpush rbp\n");
    fprintf(outfile, "\tmov rbp, rsp\n");

    scope_push(&scope);
    scope->bp_offset = 0;

    long long spilled_bp_ofs = -16; // return address + old bp
    unsigned long long arg_regnum = 0;
    struct var_decl_node* args_node = node->args_head;
    while (args_node != NULL) {
        unsigned long long type_sz = get_type_size(&args_node->type);

        scope->bp_offset += type_sz;
        struct var_def var_def = {
            .name = args_node->ident,
            .loc = {
                .type = BP_OFFSET,
                .offset = scope->bp_offset,
            },
            .sz = type_sz,
        };
        scope_define_var(scope, var_def);

        struct storage_location arg_src;
        if (arg_regnum < CC_N_REGS) {
            arg_src = (struct storage_location) {
                .type = REGISTER,
                .reg = CALLING_CONV[arg_regnum++]
            };
        } else {
            arg_src = (struct storage_location) {
                .type = BP_OFFSET,
                .offset = spilled_bp_ofs,
            };
            spilled_bp_ofs -= type_sz;
        }
        emit_mov(outfile, &var_def.loc, &arg_src, type_sz);

        args_node = args_node->next;
    }

    emit_group(outfile, &node->body);

    scope_pop(&scope);

    fprintf(outfile, "\tmov rsp, rbp\n");
    fprintf(outfile, "\tpop rbp\n");
    fprintf(outfile, "\tret\n");
}

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

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

    fprintf(outfile, "section .text\n");
    scope_push(&scope);
    scope_install_default_types(scope);

    /* output all non-static function declarations as globals */
    const struct root_node* node = ast;
    while (node != NULL) {
        if (node->type == ROOT_FN_DECL) {
            const char* fn_name = node->as._fn_decl.name;
            scope_define_var(scope, (struct var_def) {
                .name = fn_name,
                .loc = {
                    .type = JMP_LABEL,
                    .label = fn_name, 
                },
            });
            fprintf(outfile, "global %s\n", fn_name);
        }

        node = node->next;
    }

    fprintf(outfile, "\n");

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

    scope_pop(&scope);
    fclose(outfile);
}