#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", "缺少语句"));
  }
  std::string text = ctx->getText();
  LogInfo("[irgen] visitStmt text='" + text + "' break_size=" + std::to_string(break_targets_.size()) + " cont_size=" + std::to_string(continue_targets_.size()), std::cerr);
  // return
  if (ctx->getStart()->getText() == "return") {
    if (ctx->exp()) {
      ir::Value* v = EvalExpr(*ctx->exp());
      builder_.CreateRet(v);
    } else {
      builder_.CreateRet(nullptr);
    }
    return BlockFlow::Terminated;
  }

  // assignment: lVal '=' exp
  if (ctx->lVal() && text.find("=") != std::string::npos) {
    ir::Value* val = EvalExpr(*ctx->exp());
    std::string name = ctx->lVal()->Ident()->getText();
    // 优先检查按名称的快速映射（支持全局变量）
    auto nit = name_map_.find(name);
    if (nit != name_map_.end()) {
      // 支持带索引的赋值
      if (ctx->lVal()->exp().size() > 0) {
        std::vector<ir::Value*> indices;
        indices.push_back(builder_.CreateConstInt(0));
        for (auto* e : ctx->lVal()->exp()) indices.push_back(EvalExpr(*e));
        auto* gep = builder_.CreateGEP(nit->second, indices, module_.GetContext().NextTemp());
        builder_.CreateStore(val, gep);
        return BlockFlow::Continue;
      }
      builder_.CreateStore(val, nit->second);
      return BlockFlow::Continue;
    }
    for (auto& kv : storage_map_) {
      if (!kv.first) continue;
      if (auto* vdef = dynamic_cast<SysYParser::VarDefContext*>(kv.first)) {
        if (vdef->Ident() && vdef->Ident()->getText() == name) {
          builder_.CreateStore(val, kv.second);
          return BlockFlow::Continue;
        }
      } else if (auto* fparam = dynamic_cast<SysYParser::FuncFParamContext*>(kv.first)) {
        if (fparam->Ident() && fparam->Ident()->getText() == name) {
          builder_.CreateStore(val, kv.second);
          return BlockFlow::Continue;
        }
      } else if (auto* cdef = dynamic_cast<SysYParser::ConstDefContext*>(kv.first)) {
        if (cdef->Ident() && cdef->Ident()->getText() == name) {
          builder_.CreateStore(val, kv.second);
          return BlockFlow::Continue;
        }
      }
    }
    throw std::runtime_error(FormatError("irgen", "变量未声明: " + name));
  }

  // if
  if (ctx->getStart()->getText() == "if" && ctx->cond()) {
    ir::Value* condv = std::any_cast<ir::Value*>(ctx->cond()->lOrExp()->accept(this));
    ir::BasicBlock* then_bb = func_->CreateBlock("if.then");
    ir::BasicBlock* else_bb = (ctx->stmt().size() > 1) ? func_->CreateBlock("if.else") : nullptr;
    ir::BasicBlock* merge_bb = func_->CreateBlock("if.merge");

    if (else_bb) builder_.CreateCondBr(condv, then_bb, else_bb);
    else builder_.CreateCondBr(condv, then_bb, merge_bb);

    // then
    builder_.SetInsertPoint(then_bb);
    ctx->stmt(0)->accept(this);
    if (!builder_.GetInsertBlock()->HasTerminator()) builder_.CreateBr(merge_bb);

    // else
    if (else_bb) {
      builder_.SetInsertPoint(else_bb);
      ctx->stmt(1)->accept(this);
      if (!builder_.GetInsertBlock()->HasTerminator()) builder_.CreateBr(merge_bb);
    }

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

  // while
  if (ctx->getStart()->getText() == "while" && ctx->cond()) {
    ir::BasicBlock* cond_bb = func_->CreateBlock("while.cond");
    ir::BasicBlock* body_bb = func_->CreateBlock("while.body");
    ir::BasicBlock* after_bb = func_->CreateBlock("while.after");

    builder_.CreateBr(cond_bb);
    // cond
    builder_.SetInsertPoint(cond_bb);
    ir::Value* condv = std::any_cast<ir::Value*>(ctx->cond()->lOrExp()->accept(this));
    builder_.CreateCondBr(condv, body_bb, after_bb);

    // body
    builder_.SetInsertPoint(body_bb);
    LogInfo("[irgen] while body about to push targets, before sizes: break=" + std::to_string(break_targets_.size()) + ", cont=" + std::to_string(continue_targets_.size()), std::cerr);
    break_targets_.push_back(after_bb);
    continue_targets_.push_back(cond_bb);
    LogInfo("[irgen] after push: break_targets size=" + std::to_string(break_targets_.size()) + ", continue_targets size=" + std::to_string(continue_targets_.size()), std::cerr);
    ctx->stmt(0)->accept(this);
    LogInfo("[irgen] before pop: break_targets size=" + std::to_string(break_targets_.size()) + ", continue_targets size=" + std::to_string(continue_targets_.size()), std::cerr);
    continue_targets_.pop_back();
    break_targets_.pop_back();
    LogInfo("[irgen] after pop: break_targets size=" + std::to_string(break_targets_.size()) + ", continue_targets size=" + std::to_string(continue_targets_.size()), std::cerr);
    if (!builder_.GetInsertBlock()->HasTerminator()) builder_.CreateBr(cond_bb);

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

  // break
  if (ctx->getStart()->getText() == "break") {
    if (break_targets_.empty()) {
      // fallback: 尝试通过函数块名找目标（不依赖 sema），兼容因栈丢失导致的情况
      ir::BasicBlock* fallback = nullptr;
      for (auto &bb_up : func_->GetBlocks()) {
        auto *bb = bb_up.get();
        if (!bb) continue;
        if (bb->GetName().find("while.after") != std::string::npos) fallback = bb;
      }
      if (fallback) {
        LogInfo("[irgen] emit break (fallback), target=" + fallback->GetName(), std::cerr);
        builder_.CreateBr(fallback);
        return BlockFlow::Terminated;
      }
      throw std::runtime_error(FormatError("irgen", "break 不在循环内"));
    }
    LogInfo("[irgen] emit break, break_targets size=" + std::to_string(break_targets_.size()), std::cerr);
    builder_.CreateBr(break_targets_.back());
    return BlockFlow::Terminated;
  }

  // continue
  if (ctx->getStart()->getText() == "continue") {
    if (continue_targets_.empty()) {
      ir::BasicBlock* fallback = nullptr;
      for (auto &bb_up : func_->GetBlocks()) {
        auto *bb = bb_up.get();
        if (!bb) continue;
        if (bb->GetName().find("while.cond") != std::string::npos) fallback = bb;
      }
      if (fallback) {
        LogInfo("[irgen] emit continue (fallback), target=" + fallback->GetName(), std::cerr);
        builder_.CreateBr(fallback);
        return BlockFlow::Terminated;
      }
      throw std::runtime_error(FormatError("irgen", "continue 不在循环内"));
    }
    LogInfo("[irgen] emit continue, continue_targets size=" + std::to_string(continue_targets_.size()), std::cerr);
    builder_.CreateBr(continue_targets_.back());
    return BlockFlow::Terminated;
  }

  // block
  if (ctx->block()) {
    return ctx->block()->accept(this);
  }

  // expression statement
  if (ctx->exp()) {
    EvalExpr(*ctx->exp());
    return BlockFlow::Continue;
  }

  throw std::runtime_error(FormatError("irgen", "暂不支持的语句类型"));
}


