#include <iostream>
#include "ASTPrinter.h"
#include "SysYParser.h"

using namespace std;

std::any ASTPrinter::visitNumber(SysYParser::NumberContext *ctx) {
    cout << ctx->ILITERAL()->getText();
    return nullptr;
}

std::any ASTPrinter::visitString(SysYParser::StringContext *ctx) {
    cout << ctx->STRING()->getText();
    return nullptr;
}

std::any ASTPrinter::visitModule(SysYParser::ModuleContext *ctx) {
    for (auto dcl : ctx->dcl()) dcl->accept(this);
    for (auto func : ctx->funcDef()) func->accept(this);
    if (ctx->funcRParams()) ctx->funcRParams()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitFuncRParams(SysYParser::FuncRParamsContext *ctx) {
    bool first = true;
    for (auto param : ctx->funcRParam()) {
        if (!first) cout << ", ";
        param->accept(this);
        first = false;
    }
    return nullptr;
}

std::any ASTPrinter::visitExpAsRParam(SysYParser::ExpAsRParamContext *ctx) {
    ctx->number()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitStringAsRParam(SysYParser::StringAsRParamContext *ctx) {
    ctx->string()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitExpsAsRParam(SysYParser::ExpsAsRParamContext *ctx) {
    bool first = true;
    for (auto exp : ctx->exp()) {
        if (!first) cout << ", ";
        exp->accept(this);
        first = false;
    }
    return nullptr;
}

std::any ASTPrinter::visitConstDecl(SysYParser::ConstDeclContext *ctx) {
    cout << getIndent() << "const " << ctx->bType()->getText() << " ";
    bool first = true;
    for (auto def : ctx->constDef()) {
        if (!first) cout << ", ";
        def->accept(this);
        first = false;
    }
    cout << ";" << endl;
    return nullptr;
}

std::any ASTPrinter::visitVarDecl(SysYParser::VarDeclContext *ctx) {
    cout << getIndent() << ctx->bType()->getText() << " ";
    bool first = true;
    for (auto def : ctx->varDef()) {
        if (!first) cout << ", ";
        def->accept(this);
        first = false;
    }
    cout << ";" << endl;
    return nullptr;
}

std::any ASTPrinter::visitVarDef(SysYParser::VarDefContext *ctx) {
    cout << ctx->IDENT()->getText();
    for (auto exp : ctx->constExp()) {
        cout << "[";
        exp->accept(this);
        cout << "]";
    }
    if (ctx->initVal()) {
        cout << " = ";
        ctx->initVal()->accept(this);
    }
    return nullptr;
}

std::any ASTPrinter::visitConstDef(SysYParser::ConstDefContext *ctx) {
    cout << ctx->IDENT()->getText();
    for (auto exp : ctx->constExp()) {
        cout << "[";
        exp->accept(this);
        cout << "]";
    }
    cout << " = ";
    ctx->constInitVal()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitFuncDef(SysYParser::FuncDefContext *ctx) {
    cout << getIndent() << ctx->funcType()->getText() << " " << ctx->IDENT()->getText() << "(";
    if (ctx->funcFParams()) ctx->funcFParams()->accept(this);
    cout << ")" << endl;
    ctx->block()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitFuncFParams(SysYParser::FuncFParamsContext *ctx) {
    bool first = true;
    for (auto param : ctx->funcFParam()) {
        if (!first) cout << ", ";
        param->accept(this);
        first = false;
    }
    return nullptr;
}

std::any ASTPrinter::visitFuncFParam(SysYParser::FuncFParamContext *ctx) {
    cout << ctx->bType()->getText() << " " << ctx->IDENT()->getText();
    if (!ctx->exp().empty()) {
        cout << "[]";
        for (auto exp : ctx->exp()) {
            cout << "[";
            exp->accept(this);
            cout << "]";
        }
    }
    return nullptr;
}

std::any ASTPrinter::visitExp(SysYParser::ExpContext *ctx) {
    ctx->addExp()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitCond(SysYParser::CondContext *ctx) {
    ctx->lorExp()->accept(this);
    return nullptr;
}

std::any ASTPrinter::visitLVal(SysYParser::LValContext *ctx) {
    cout << ctx->IDENT()->getText();
    for (auto exp : ctx->exp()) {
        cout << "[";
        exp->accept(this);
        cout << "]";
    }
    return nullptr;
}

std::any ASTPrinter::visitAddExp(SysYParser::AddExpContext *ctx) {
    if (ctx->addExp()) {
        ctx->addExp()->accept(this);
        cout << " " << ctx->ADD()->getText() << " ";
        ctx->mulExp()->accept(this);
    } else {
        ctx->mulExp()->accept(this);
    }
    return nullptr;
}

std::any ASTPrinter::visitMulExp(SysYParser::MulExpContext *ctx) {
    auto unaryExps = ctx->unaryExp();
    if (unaryExps.size() == 1) {
        unaryExps[0]->accept(this);
    } else {
        for (size_t i = 0; i < unaryExps.size() - 1; ++i) {
            auto opNode = dynamic_cast<antlr4::tree::TerminalNode *>(ctx->children[2 * i + 1]);
            if (opNode) {
                unaryExps[i]->accept(this);
                cout << " " << opNode->getText() << " ";
            }
        }
        unaryExps.back()->accept(this);
    }
    return nullptr;
}

std::any ASTPrinter::visitUnaryExp(SysYParser::UnaryExpContext *ctx) {
    if (ctx->primaryExp()) {
        ctx->primaryExp()->accept(this);
    } else if (ctx->IDENT()) {
        cout << ctx->IDENT()->getText() << "(";
        if (ctx->funcRParams()) ctx->funcRParams()->accept(this);
        cout << ")";
    } else if (ctx->unaryOp()) {
        cout << ctx->unaryOp()->getText();
        ctx->unaryExp()->accept(this);
    }
    return nullptr;
}


// std::any ASTPrinter::visitLorExp(SysYParser::LorExpContext *ctx) {
//     if (ctx->lorExp()) {
//         // 左递归部分
//         ctx->lorExp()->accept(this);
//         cout << " || ";
//         ctx->landExp()->accept(this);
//     } else {
//         // 基础部分
//         ctx->landExp()->accept(this);
//     }
//     return nullptr;
// }

std::any ASTPrinter::visitStmt(SysYParser::StmtContext *ctx) {
    if (ctx->lVal() && ctx->exp()) {
        cout << getIndent();
        ctx->lVal()->accept(this);
        cout << " = ";
        ctx->exp()->accept(this);
        cout << ";" << endl;
    }
    // else if (ctx->exp()) {
    //     cout << getIndent();
    //     ctx->exp()->accept(this);
    //     cout << ";" << endl;
    // }
     else if (ctx->block()) {
        ctx->block()->accept(this);
    } else {
        for (auto child : ctx->children) {
            if (auto terminal = dynamic_cast<antlr4::tree::TerminalNode *>(child)) {
                if (terminal->getText() == "if") {
                    cout << getIndent() << "if ("; 
                    ctx->cond()->accept(this);
                    cout << ")" << endl;
                    ctx->stmt(0)->accept(this);
                    if (ctx->stmt().size() > 1) {
                        cout << getIndent() << "else" << endl;
                        ctx->stmt(1)->accept(this);
                    }
                } else if (terminal->getText() == "while") {
                    cout << getIndent() << "while (";
                    ctx->cond()->accept(this);
                    cout << ")" << endl;
                    ctx->stmt(0)->accept(this);
                } else if (terminal->getText() == "return") {
                    cout << getIndent() << "return";
                    if (ctx->exp()) {
                        cout << " ";
                        ctx->exp()->accept(this);
                    }
                    cout << ";" << endl;
                } else if (terminal->getText() == "break") {
                    cout << getIndent() << "break;" << endl;
                } else if (terminal->getText() == "continue") {
                    cout << getIndent() << "continue;" << endl;
                }
            }
        }
    }
    return nullptr;
}

std::any ASTPrinter::visitBlock(SysYParser::BlockContext *ctx) {
    cout << getIndent() << "{" << endl;
    indentLevel++;
    for (auto item : ctx->blockItem()) item->accept(this);
    indentLevel--;
    cout << getIndent() << "}" << endl;
    return nullptr;
}