#include "irgen/IRGen.h"

#include <stdexcept>

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

// 语句生成当前只实现了最小子集。
// 目前支持：
// - return <exp>;
//
// 还未支持：
// - 赋值语句
// - if / while 等控制流
// - 空语句、块语句嵌套分发之外的更多语句形态

std::any IRGenImpl::visitStmt(SysYParser::StmtContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少语句"));
  }
  if (ctx->lValue() && ctx->ASSIGN() && ctx->exp()) {
    ir::Value* rhs = EvalExpr(*ctx->exp());
    ir::Value* slot = ResolveStorage(ctx->lValue());
    if (!slot) {
      throw std::runtime_error(
          FormatError("irgen", "赋值目标未找到存储槽位: " +
                                   (ctx->lValue()->ID()
                                        ? ctx->lValue()->ID()->getText()
                                        : "?")));
    }
    builder_.CreateStore(rhs, slot);
    return BlockFlow::Continue;
  }
  if (ctx->blockStmt()) {
    ctx->blockStmt()->accept(this);
    return builder_.GetInsertBlock() && builder_.GetInsertBlock()->HasTerminator()
               ? BlockFlow::Terminated
               : BlockFlow::Continue;
  }
  if (ctx->IF()) {
    if (!ctx->cond() || ctx->stmt().empty()) {
      throw std::runtime_error(FormatError("irgen", "if 语句不完整"));
    }
    auto* then_bb = func_->CreateBlock(NextBlockName());
    auto* merge_bb = func_->CreateBlock(NextBlockName());
    auto* else_bb = ctx->ELSE() ? func_->CreateBlock(NextBlockName()) : merge_bb;

    ir::Value* cond = ToBoolValue(EvalCond(*ctx->cond()));
    builder_.CreateCondBr(cond, then_bb, else_bb);

    builder_.SetInsertPoint(then_bb);
    auto then_flow = std::any_cast<BlockFlow>(ctx->stmt(0)->accept(this));
    bool then_term = (then_flow == BlockFlow::Terminated);
    if (then_flow != BlockFlow::Terminated) {
      builder_.CreateBr(merge_bb);
    }

    bool else_term = false;
    if (ctx->ELSE()) {
      builder_.SetInsertPoint(else_bb);
      auto else_flow = std::any_cast<BlockFlow>(ctx->stmt(1)->accept(this));
      else_term = (else_flow == BlockFlow::Terminated);
      if (else_flow != BlockFlow::Terminated) {
        builder_.CreateBr(merge_bb);
      }
    }

    if (ctx->ELSE() && then_term && else_term) {
      // 两个分支都终结时，merge 块不可达；补一个自环 terminator 以满足结构校验。
      builder_.SetInsertPoint(merge_bb);
      builder_.CreateBr(merge_bb);
      return BlockFlow::Terminated;
    }

    builder_.SetInsertPoint(merge_bb);
    return BlockFlow::Continue;
  }
  if (ctx->WHILE()) {
    if (!ctx->cond() || ctx->stmt().empty()) {
      throw std::runtime_error(FormatError("irgen", "while 语句不完整"));
    }
    auto* cond_bb = func_->CreateBlock(NextBlockName());
    auto* body_bb = func_->CreateBlock(NextBlockName());
    auto* exit_bb = func_->CreateBlock(NextBlockName());

    builder_.CreateBr(cond_bb);
    builder_.SetInsertPoint(cond_bb);
    ir::Value* cond = ToBoolValue(EvalCond(*ctx->cond()));
    builder_.CreateCondBr(cond, body_bb, exit_bb);

    loop_stack_.push_back({cond_bb, exit_bb});
    builder_.SetInsertPoint(body_bb);
    auto body_flow = std::any_cast<BlockFlow>(ctx->stmt(0)->accept(this));
    if (body_flow != BlockFlow::Terminated) {
      builder_.CreateBr(cond_bb);
    }
    loop_stack_.pop_back();

    builder_.SetInsertPoint(exit_bb);
    return BlockFlow::Continue;
  }
  if (ctx->BREAK()) {
    if (loop_stack_.empty()) {
      throw std::runtime_error(FormatError("irgen", "break 不在循环中"));
    }
    builder_.CreateBr(loop_stack_.back().break_target);
    return BlockFlow::Terminated;
  }
  if (ctx->CONTINUE()) {
    if (loop_stack_.empty()) {
      throw std::runtime_error(FormatError("irgen", "continue 不在循环中"));
    }
    builder_.CreateBr(loop_stack_.back().continue_target);
    return BlockFlow::Terminated;
  }
  if (ctx->returnStmt()) {
    return ctx->returnStmt()->accept(this);
  }
  if (ctx->exp()) {
    EvalExpr(*ctx->exp());
    return BlockFlow::Continue;
  }
  if (ctx->SEMICOLON()) {
    return BlockFlow::Continue;
  }
  throw std::runtime_error(FormatError("irgen", "暂不支持的语句类型"));
}


std::any IRGenImpl::visitReturnStmt(SysYParser::ReturnStmtContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少 return 语句"));
  }
  if (!ctx->exp()) {
    // void 函数的 return;
    builder_.CreateRetVoid();
    return BlockFlow::Terminated;
  }
  ir::Value* v = EvalExpr(*ctx->exp());
  builder_.CreateRet(v);
  return BlockFlow::Terminated;
}