#include "irgen/IRGen.h"

#include <stdexcept>
#include <string>

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

// ─── visitStmt ────────────────────────────────────────────────────────────────
std::any IRGenImpl::visitStmt(SysYParser::StmtContext* ctx) {
  return VisitStmt(*ctx);
}

IRGenImpl::BlockFlow IRGenImpl::VisitStmt(SysYParser::StmtContext& s) {
  auto* ctx = &s;

  // 若当前块已经终结，跳过死代码
  {
    auto* cur = builder_.GetInsertBlock();
    if (cur && cur->HasTerminator()) return BlockFlow::Terminated;
  }

  // ── return ──────────────────────────────────────────────────────────────────
  if (ctx->Return()) {
    if (ctx->exp()) {
      ir::Value* v = EvalExpr(*ctx->exp());
      // Coerce return value to function return type
      if (func_->GetType()->IsInt32() && v->IsFloat32()) {
        v = ToInt(v);
      } else if (func_->GetType()->IsFloat32() && v->IsInt32()) {
        v = ToFloat(v);
      } else if (func_->GetType()->IsInt32() && v->IsInt1()) {
        v = ToI32(v);
      }
      builder_.CreateRet(v);
    } else {
      builder_.CreateRetVoid();
    }
    return BlockFlow::Terminated;
  }

  // ── break ────────────────────────────────────────────────────────────────────
  if (ctx->Break()) {
    if (loop_stack_.empty()) {
      throw std::runtime_error(FormatError("irgen", "break 在循环外"));
    }
    builder_.CreateBr(loop_stack_.back().after_bb);
    return BlockFlow::Terminated;
  }

  // ── continue ─────────────────────────────────────────────────────────────────
  if (ctx->Continue()) {
    if (loop_stack_.empty()) {
      throw std::runtime_error(FormatError("irgen", "continue 在循环外"));
    }
    builder_.CreateBr(loop_stack_.back().cond_bb);
    return BlockFlow::Terminated;
  }

  // ── 赋值 lVar = exp ──────────────────────────────────────────────────────────
  if (ctx->lVar() && ctx->Assign()) {
    ir::Value* rhs = EvalExpr(*ctx->exp());
    ir::Value* addr = EvalLVarAddr(ctx->lVar());
    // Coerce rhs to match slot type
    if (addr->IsPtrInt32() && rhs->IsFloat32()) {
      rhs = ToInt(rhs);
    } else if (addr->IsPtrFloat32() && rhs->IsInt32()) {
      rhs = ToFloat(rhs);
    } else if (addr->IsPtrInt32() && rhs->IsInt1()) {
      rhs = ToI32(rhs);
    }
    builder_.CreateStore(rhs, addr);
    return BlockFlow::Continue;
  }

  // ── 纯表达式语句 (exp)? ; ────────────────────────────────────────────────────
  if (ctx->Semi() && !ctx->Return() && !ctx->If() && !ctx->While() &&
      !ctx->Break() && !ctx->Continue() && !ctx->Assign()) {
    if (ctx->exp()) {
      EvalExpr(*ctx->exp());
    }
    return BlockFlow::Continue;
  }

  // ── if ───────────────────────────────────────────────────────────────────────
  if (ctx->If()) {
    if (!ctx->cond()) {
      throw std::runtime_error(FormatError("irgen", "if 缺少条件"));
    }

    auto stmts = ctx->stmt();

    // Step 1: evaluate condition (may create short-circuit blocks with lower
    // SSA numbers — must happen before any branch-target blocks are created).
    ir::Value* cond_val = EvalCond(*ctx->cond());

    // Step 2: create then_bb now (its label number will be >= all short-circuit
    // block numbers allocated during EvalCond).
    ir::BasicBlock* then_bb = func_->CreateBlock(
        module_.GetContext().NextLabel() + ".if.then");

    // Step 3: create else_bb/merge_bb as placeholders. They will be moved to
    // the end of the block list after their predecessors are filled in, so the
    // block ordering in the output will be correct even though their label
    // numbers are allocated here (before then-body sub-blocks are created).
    ir::BasicBlock* else_bb = nullptr;
    ir::BasicBlock* merge_bb = nullptr;

    if (stmts.size() >= 2) {
      else_bb  = func_->CreateBlock(module_.GetContext().NextLabel() + ".if.else");
      merge_bb = func_->CreateBlock(module_.GetContext().NextLabel() + ".if.end");
    } else {
      merge_bb = func_->CreateBlock(module_.GetContext().NextLabel() + ".if.end");
    }

    // Check if current block already terminated (short-circuit may do this)
    if (builder_.GetInsertBlock() && builder_.GetInsertBlock()->HasTerminator()) {
      func_->MoveBlockToEnd(then_bb);
      if (else_bb) func_->MoveBlockToEnd(else_bb);
      func_->MoveBlockToEnd(merge_bb);
      builder_.SetInsertPoint(merge_bb);
      return BlockFlow::Continue;
    }

    if (stmts.size() >= 2) {
      builder_.CreateCondBr(cond_val, then_bb, else_bb);
    } else {
      builder_.CreateCondBr(cond_val, then_bb, merge_bb);
    }

    // then 分支 — visit body (may create many sub-blocks with higher numbers)
    func_->MoveBlockToEnd(then_bb);
    builder_.SetInsertPoint(then_bb);
    auto then_flow = VisitStmt(*stmts[0]);
    if (then_flow != BlockFlow::Terminated) {
      builder_.CreateBr(merge_bb);
    }

    // else 分支
    if (else_bb) {
      func_->MoveBlockToEnd(else_bb);
      builder_.SetInsertPoint(else_bb);
      auto else_flow = VisitStmt(*stmts[1]);
      if (else_flow != BlockFlow::Terminated) {
        builder_.CreateBr(merge_bb);
      }
    }

    func_->MoveBlockToEnd(merge_bb);
    builder_.SetInsertPoint(merge_bb);
    return BlockFlow::Continue;
  }

  // ── while ────────────────────────────────────────────────────────────────────
  if (ctx->While()) {
    if (!ctx->cond()) {
      throw std::runtime_error(FormatError("irgen", "while 缺少条件"));
    }

    ir::BasicBlock* cond_bb = func_->CreateBlock(
        module_.GetContext().NextLabel() + ".while.cond");

    // 跳转到条件块
    if (!builder_.GetInsertBlock()->HasTerminator()) {
      builder_.CreateBr(cond_bb);
    }

    // EvalCond MUST come before creating body_bb/after_bb so that
    // short-circuit blocks get lower SSA numbers than the loop body blocks.
    builder_.SetInsertPoint(cond_bb);
    ir::Value* cond_val = EvalCond(*ctx->cond());

    ir::BasicBlock* body_bb = func_->CreateBlock(
        module_.GetContext().NextLabel() + ".while.body");
    ir::BasicBlock* after_bb = func_->CreateBlock(
        module_.GetContext().NextLabel() + ".while.end");

    // 检查条件求值后是否已终结
    if (!builder_.GetInsertBlock()->HasTerminator()) {
      builder_.CreateCondBr(cond_val, body_bb, after_bb);
    }

    // 循环体（压入循环栈）
    func_->MoveBlockToEnd(body_bb);
    loop_stack_.push_back({cond_bb, after_bb});
    builder_.SetInsertPoint(body_bb);
    auto stmts = ctx->stmt();
    if (!stmts.empty()) {
      auto body_flow = VisitStmt(*stmts[0]);
      if (body_flow != BlockFlow::Terminated) {
        builder_.CreateBr(cond_bb);  // 循环回跳
      }
    } else {
      builder_.CreateBr(cond_bb);
    }
    loop_stack_.pop_back();

    func_->MoveBlockToEnd(after_bb);
    builder_.SetInsertPoint(after_bb);
    return BlockFlow::Continue;
  }

  // ── 块语句 ───────────────────────────────────────────────────────────────────
  if (ctx->block()) {
    auto result = ctx->block()->accept(this);
    if (result.has_value()) {
      try {
        auto flow = std::any_cast<BlockFlow>(result);
        if (flow == BlockFlow::Terminated) return BlockFlow::Terminated;
      } catch (...) {}
    }
    // 检查 builder 的实际状态
    if (builder_.GetInsertBlock() && builder_.GetInsertBlock()->HasTerminator()) {
      return BlockFlow::Terminated;
    }
    return BlockFlow::Continue;
  }

  // ── 空语句 ; ─────────────────────────────────────────────────────────────────
  if (ctx->Semi()) {
    return BlockFlow::Continue;
  }

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