#include "CodeGenVisitor.h"

#include "generated/ifccParser.h"

antlrcpp::Any CodeGenVisitor::visitStmt(ifccParser::StmtContext *ctx) {
    std::cout << "    ;" << ctx->getText() << "\n";
    return this->visitChildren(ctx);
}

// Declare a new variable: reserve a slot on the stack, mark as uninitialized.
antlrcpp::Any CodeGenVisitor::visitDecl_stmt(ifccParser::Decl_stmtContext *ctx) {
    std::string name = ctx->VAR_NAME()->getText();
    if (symbolTable.find(name) != symbolTable.end()) {
        std::cerr << "error: variable '" << name << "' already declared\n";
        return 1;
    }
    symbolTable[name] = {nextOffset, false};
    nextOffset += 4;
    return 0;
}

// Assign a value to a variable: evaluate val, store on stack, mark as initialized.
antlrcpp::Any CodeGenVisitor::visitSet_stmt(ifccParser::Set_stmtContext *ctx) {
    std::string name = ctx->VAR_NAME()->getText();
    if (symbolTable.find(name) == symbolTable.end()) {
        std::cerr << "error: variable '" << name << "' undeclared\n";
        return 1;
    }
    int offset = symbolTable[name].offset;

    this->visit(ctx->val());

#ifdef __APPLE__
    std::cout << "    str w0, [sp, #" << offset << "]\n";
#else
    std::cout << "    movl %eax, " << offset << "(%rsp)\n";
#endif

    symbolTable[name].initialized = true;
    return 0;
}

// Load a constant or variable into the accumulator register
antlrcpp::Any CodeGenVisitor::visitVal(ifccParser::ValContext *ctx) {
    if (ctx->CONST()) {
        int val = stoi(ctx->CONST()->getText());
#ifdef __APPLE__
        std::cout << "    mov w0, #" << val << "\n";
#else
        std::cout << "    movl $" << val << ", %eax\n";
#endif
        return val;
    }
    std::string name = ctx->VAR_NAME()->getText();
    if (symbolTable.find(name) == symbolTable.end()) {
        std::cerr << "error: variable '" << name << "' undeclared\n";
        return 1;
    }
    if (!symbolTable[name].initialized) {
        std::cerr << "warning: variable '" << name << "' used before initialization\n";
    }
    int offset = symbolTable[name].offset;
#ifdef __APPLE__
    std::cout << "    ldr w0, [sp, #" << offset << "]\n";
#else
    std::cout << "    movl " << offset << "(%rsp), %eax\n";
#endif
    return offset;
}

antlrcpp::Any CodeGenVisitor::visitProg(ifccParser::ProgContext *ctx) {
#ifdef __APPLE__
    std::cout << ".globl _main\n";
    std::cout << "_main:\n";
    // Prologue: reserve space for local variables
    std::cout << "    sub sp, sp, #16\n";
#else
    std::cout << ".globl main\n";
    std::cout << "main:\n";
    // Prologue: reserve space for local variables
    std::cout << "    subq $16, %rsp\n";
#endif

    this->visitChildren(ctx);

    // Epilogue
#ifdef __APPLE__
    std::cout << "    add sp, sp, #16\n";
#else
    std::cout << "    addq $16, %rsp\n";
#endif
    std::cout << "    ret\n";
    return 0;
}

antlrcpp::Any CodeGenVisitor::visitReturn_stmt(ifccParser::Return_stmtContext *ctx) {
    // Evaluate the return value into the accumulator register
    this->visit(ctx->val());
    return 0;
}