#include "CodeGenVisitor.h"
#include "generated/ifccParser.h"

std::any CodeGenVisitor::visitProg(ifccParser::ProgContext *ctx) {
    scopeStack.push_back("main");
    int size = symbolTable->stackSize(currentScope());
#ifdef __APPLE__
    std::cout << ".globl _main\n_main:\n";
    std::cout << "    sub sp, sp, #" << size << "\n";
#else
    std::cout << ".globl main\nmain:\n";
    std::cout << "    subq $" << size << ", %rsp\n";
#endif

    this->visitChildren(ctx);

#ifdef __APPLE__
    std::cout << "    add sp, sp, #" << size << "\n";
#else
    std::cout << "    addq $" << size << ", %rsp\n";
#endif
    std::cout << "    ret\n";
    scopeStack.pop_back();
    return 0;
}

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

}
std::any CodeGenVisitor::visitReturn_stmt(ifccParser::Return_stmtContext *ctx) {
    this->visit(ctx->expr());
    return 0;
}

// Declaration pass already filled the symbol table: nothing to do here.
std::any CodeGenVisitor::visitDecl_stmt(ifccParser::Decl_stmtContext *ctx) {
    return 0;
}

// Assign a value to a variable: evaluate val, store on stack, mark as initialized.
std::any CodeGenVisitor::visitSet_stmt(ifccParser::Set_stmtContext *ctx) {
    std::string name = ctx->VAR_NAME()->getText();
    int offset = symbolTable->getOffset(currentScope(), name) + stackOffset;

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

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

    symbolTable->markInitialized(currentScope(), name);
    return 0;
}

// Load a constant or variable into the accumulator register.
std::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();
    int offset = symbolTable->getOffset(currentScope(), name) + stackOffset;
    symbolTable->isInitialized(currentScope(), name); // emits warning if needed

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

    return offset;
}

std::any CodeGenVisitor::visitExpr(ifccParser::ExprContext *ctx) {
    // Parenthesized expression or single val: delegate directly
    if (ctx->expr().size() == 1) {
        return this->visit(ctx->expr(0));
    }
    if (ctx->expr().empty()) {
        return this->visit(ctx->val());
    }

    // Evaluate left, push result onto stack to survive nested sub-expressions
    this->visit(ctx->expr(0));
#ifdef __APPLE__
    std::cout << "    str w0, [sp, #-16]!\n"; // sp -= 16, push -> [sp]
#else
    std::cout << "    pushq %rax\n";
#endif
    stackOffset += 16;

    // Evaluate right, result in w0 / eax
    this->visit(ctx->expr(1));

#ifdef __APPLE__
    std::cout << "    ldr w8, [sp], #16\n";  // pop -> w8, sp += 16
    if (ctx->MUL())      std::cout << "    mul w0, w8, w0\n";
    else if (ctx->DIV()) std::cout << "    sdiv w0, w8, w0\n";
    else if (ctx->ADD()) std::cout << "    add w0, w8, w0\n";
    else if (ctx->SUB()) std::cout << "    sub w0, w8, w0\n";
#else

#endif

    stackOffset -= 16;
    return 0;
}