#include "irgen/IRGen.h"

#include <stdexcept>

#include "SysYParser.h"
#include "ir/IR.h"
#include "utils/Log.h"

// 语句生成
std::any IRGenImpl::visitStmt(SysYParser::StmtContext* ctx) {
  if (!ctx) return BlockFlow::Continue;
  
  if (ctx->assignStmt()) return ctx->assignStmt()->accept(this);
  if (ctx->returnStmt()) return ctx->returnStmt()->accept(this);
  if (ctx->blockStmt()) return ctx->blockStmt()->accept(this);
  if (ctx->ifStmt()) return ctx->ifStmt()->accept(this);
  if (ctx->whileStmt()) return ctx->whileStmt()->accept(this);
  if (ctx->breakStmt()) return ctx->breakStmt()->accept(this);
  if (ctx->continueStmt()) return ctx->continueStmt()->accept(this);
  if (ctx->expStmt()) return ctx->expStmt()->accept(this);
  
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitBlockStmt(SysYParser::BlockStmtContext* ctx) {
  for (auto* item : ctx->blockItem()) {
    if (std::any_cast<BlockFlow>(item->accept(this)) == BlockFlow::Terminated) {
      return BlockFlow::Terminated;
    }
  }
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitBlockItem(SysYParser::BlockItemContext* ctx) {
  if (ctx->decl()) {
    ctx->decl()->accept(this);
    return BlockFlow::Continue;
  }
  if (ctx->stmt()) {
    return ctx->stmt()->accept(this);
  }
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitAssignStmt(SysYParser::AssignStmtContext* ctx) {
  ir::Value* ptr = GetLValuePtr(ctx->lValue());
  ir::Value* val = EvalExpr(*ctx->exp());
  if (ptr->GetType()->IsPtrFloat() && val->GetType()->IsInt32()) {
    val = builder_.CreateSIToFP(val, ir::Type::GetFloatType(),
                                module_.GetContext().NextTemp());
  } else if (ptr->GetType()->IsPtrInt32() && val->GetType()->IsFloat()) {
    val = builder_.CreateFPToSI(val, ir::Type::GetInt32Type(),
                                module_.GetContext().NextTemp());
  }
  builder_.CreateStore(val, ptr);
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitReturnStmt(SysYParser::ReturnStmtContext* ctx) {
  if (ctx->exp()) {
    ir::Value* v = EvalExpr(*ctx->exp());
    if (func_->GetType()->IsFloat() && v->GetType()->IsInt32()) {
      v = builder_.CreateSIToFP(v, ir::Type::GetFloatType(),
                                module_.GetContext().NextTemp());
    } else if (func_->GetType()->IsInt32() && v->GetType()->IsFloat()) {
      v = builder_.CreateFPToSI(v, ir::Type::GetInt32Type(),
                                module_.GetContext().NextTemp());
    }
    builder_.CreateRet(v);
  } else {
    builder_.CreateRet(nullptr);
  }
  return BlockFlow::Terminated;
}

std::any IRGenImpl::visitIfStmt(SysYParser::IfStmtContext* ctx) {
  const std::string suffix = module_.GetContext().NextTemp();
  ir::BasicBlock* true_bb = func_->CreateBlock("if.true." + suffix);
  ir::BasicBlock* false_bb =
      ctx->ELSE() ? func_->CreateBlock("if.false." + suffix) : nullptr;
  ir::BasicBlock* merge_bb = func_->CreateBlock("if.merge." + suffix);

  ir::Value* cond = EvalExpr(*ctx->exp());
  builder_.CreateCondBr(ToI1(cond), true_bb, false_bb ? false_bb : merge_bb);

  // True block
  builder_.SetInsertPoint(true_bb);
  if (std::any_cast<BlockFlow>(ctx->stmt(0)->accept(this)) == BlockFlow::Continue) {
    builder_.CreateBr(merge_bb);
  }

  // False block
  if (false_bb) {
    builder_.SetInsertPoint(false_bb);
    if (std::any_cast<BlockFlow>(ctx->stmt(1)->accept(this)) == BlockFlow::Continue) {
      builder_.CreateBr(merge_bb);
    }
  }

  builder_.SetInsertPoint(merge_bb);
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitWhileStmt(SysYParser::WhileStmtContext* ctx) {
  const std::string suffix = module_.GetContext().NextTemp();
  ir::BasicBlock* cond_bb = func_->CreateBlock("while.cond." + suffix);
  ir::BasicBlock* body_bb = func_->CreateBlock("while.body." + suffix);
  ir::BasicBlock* end_bb = func_->CreateBlock("while.end." + suffix);

  builder_.CreateBr(cond_bb);
  builder_.SetInsertPoint(cond_bb);
  ir::Value* cond = EvalExpr(*ctx->exp());
  builder_.CreateCondBr(ToI1(cond), body_bb, end_bb);

  break_stack_.push(end_bb);
  continue_stack_.push(cond_bb);

  builder_.SetInsertPoint(body_bb);
  if (std::any_cast<BlockFlow>(ctx->stmt()->accept(this)) == BlockFlow::Continue) {
    builder_.CreateBr(cond_bb);
  }

  break_stack_.pop();
  continue_stack_.pop();

  builder_.SetInsertPoint(end_bb);
  return BlockFlow::Continue;
}

std::any IRGenImpl::visitBreakStmt(SysYParser::BreakStmtContext* ctx) {
  if (break_stack_.empty()) {
    throw std::runtime_error(FormatError("irgen", "break 语句不在循环内"));
  }
  builder_.CreateBr(break_stack_.top());
  return BlockFlow::Terminated;
}

std::any IRGenImpl::visitContinueStmt(SysYParser::ContinueStmtContext* ctx) {
  if (continue_stack_.empty()) {
    throw std::runtime_error(FormatError("irgen", "continue 语句不在循环内"));
  }
  builder_.CreateBr(continue_stack_.top());
  return BlockFlow::Terminated;
}

std::any IRGenImpl::visitExpStmt(SysYParser::ExpStmtContext* ctx) {
  if (ctx->exp()) {
    EvalExpr(*ctx->exp());
  }
  return BlockFlow::Continue;
}