// 将 parse tree 转换为 AST。
#include "frontend/AstBuilder.h"

#include <any>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "SysYBaseVisitor.h"
#include "SysYParser.h"
#include "ast/AstNodes.h"
#include "antlr4-runtime.h"

namespace {

using ast::BinaryExpr;
using ast::BinaryOp;
using ast::Block;
using ast::CompUnit;
using ast::FuncDef;
using ast::NumberExpr;
using ast::ReturnStmt;
using ast::VarDecl;
using ast::VarExpr;

template <typename T>
T Take(std::any&& value) {
  if (auto* ptr = std::any_cast<T>(&value)) {
    return std::move(*ptr);
  }
  throw std::runtime_error("AST 构建失败：类型不匹配");
}

class Builder : public SysYBaseVisitor {
 public:
  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override {
    auto func = Take<std::shared_ptr<FuncDef>>(visit(ctx->funcDef()));
    return std::make_shared<CompUnit>(std::move(func));
  }

  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override {
    auto body = Take<std::shared_ptr<Block>>(visit(ctx->block()));
    return std::make_shared<FuncDef>("main", std::move(body));
  }

  std::any visitBlock(SysYParser::BlockContext* ctx) override {
    auto block = std::make_shared<Block>();
    for (auto stmtCtx : ctx->stmt()) {
      if (stmtCtx->varDecl()) {
        block->varDecls.emplace_back(
            Take<std::shared_ptr<VarDecl>>(visit(stmtCtx->varDecl())));
      } else if (stmtCtx->returnStmt()) {
        block->stmts.emplace_back(
            Take<std::shared_ptr<ReturnStmt>>(visit(stmtCtx->returnStmt())));
      }
    }
    return block;
  }

  std::any visitVarDecl(SysYParser::VarDeclContext* ctx) override {
    std::shared_ptr<ast::Expr> init;
    if (ctx->exp()) {
      init = Take<std::shared_ptr<ast::Expr>>(visit(ctx->exp()));
    }
    return std::make_shared<VarDecl>(ctx->Ident()->getText(), std::move(init));
  }

  std::any visitReturnStmt(SysYParser::ReturnStmtContext* ctx) override {
    auto expr = Take<std::shared_ptr<ast::Expr>>(visit(ctx->exp()));
    return std::make_shared<ReturnStmt>(std::move(expr));
  }

  std::any visitExp(SysYParser::ExpContext* ctx) override {
    return visit(ctx->addExp());
  }

  std::any visitAddExp(SysYParser::AddExpContext* ctx) override {
    auto node = Take<std::shared_ptr<ast::Expr>>(visit(ctx->primary(0)));
    for (size_t i = 1; i < ctx->primary().size(); ++i) {
      auto rhs = Take<std::shared_ptr<ast::Expr>>(visit(ctx->primary(i)));
      auto opToken = ctx->AddOp(i - 1);
      BinaryOp op = BinaryOp::Add;
      if (opToken->getText() == "-") op = BinaryOp::Sub;
      node = std::make_shared<BinaryExpr>(op, std::move(node), std::move(rhs));
    }
    return node;
  }

  std::any visitPrimary(SysYParser::PrimaryContext* ctx) override {
    if (ctx->Number()) {
      std::shared_ptr<ast::Expr> expr =
          std::make_shared<NumberExpr>(std::stoi(ctx->Number()->getText()));
      return expr;
    }
    if (ctx->Ident()) {
      std::shared_ptr<ast::Expr> expr =
          std::make_shared<VarExpr>(ctx->Ident()->getText());
      return expr;
    }
    return visit(ctx->exp());
  }
};

}  // namespace

std::shared_ptr<ast::CompUnit> BuildAst(antlr4::tree::ParseTree* tree) {
  if (!tree) {
    throw std::runtime_error("parse tree 为空");
  }
  Builder visitor;
  auto result = visitor.visit(tree);
  return Take<std::shared_ptr<ast::CompUnit>>(std::move(result));
}