#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 (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] visitStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少语句"));
  }

  // return 语句 - 通过 Return() 关键字判断
  if (ctx->Return()) {
    return HandleReturnStmt(ctx);
  }

  // 赋值语句
  if (ctx->lVal() && ctx->Assign() && ctx->exp()) {
    return HandleAssignStmt(ctx);
  }

  // if 语句
  if (ctx->If()) {
    return HandleIfStmt(ctx);
  }

  // while 语句
  if (ctx->While()) {
    return HandleWhileStmt(ctx);
  }

  // break 语句
  if (ctx->Break()) {
    return HandleBreakStmt(ctx);
  }
  // continue 语句
  if (ctx->Continue()) {
    return HandleContinueStmt(ctx);   
  }
  // 块语句
  if (ctx->block()) {
    return ctx->block()->accept(this);
  }

  // 空语句或表达式语句（先计算表达式）
  if (ctx->exp()) {
    EvalExpr(*ctx->exp());
    return BlockFlow::Continue;
  }

  throw std::runtime_error(FormatError("irgen", "暂不支持的语句类型"));
}
// 修改 HandleReturnStmt 函数

IRGenImpl::BlockFlow IRGenImpl::HandleReturnStmt(SysYParser::StmtContext* ctx) {
  if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleReturnStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少 return 语句"));
  }
  
  // 检查函数是否存在
  if (!func_) {
    throw std::runtime_error(FormatError("irgen", "Return语句不在函数中"));
  }
  
  // 获取函数类型中的返回类型
  auto func_type = std::dynamic_pointer_cast<ir::FunctionType>(func_->GetType());
  if (!func_type) {
    throw std::runtime_error(FormatError("irgen", "函数类型无效"));
  }
  
  auto ret_type = func_type->GetReturnType();

  if (ret_type->IsVoid()) {
    if (ctx->exp()) {
      // 表达式被忽略（可计算但不使用）
      EvalExpr(*ctx->exp());
    }
    // 对于void函数，创建返回指令（不传参数）
    builder_.CreateRet(nullptr);
  } else {
    ir::Value* retValue = nullptr;
    if (ctx->exp()) {
      retValue = EvalExpr(*ctx->exp());
      if (!retValue) {
        throw std::runtime_error(FormatError("irgen", "返回值表达式求值失败"));
      }
      // 类型转换
      if (retValue->GetType() != ret_type) {
        if (ret_type->IsInt32() && retValue->GetType()->IsFloat()) {
          retValue = builder_.CreateFPToSI(retValue, ir::Type::GetInt32Type());
        } else if (ret_type->IsFloat() && retValue->GetType()->IsInt32()) {
          retValue = builder_.CreateSIToFP(retValue, ir::Type::GetFloatType());
        }
      }
    } else {
      // 无表达式，返回默认值
      if (ret_type->IsInt32()) {
        retValue = builder_.CreateConstInt(0);
      } else if (ret_type->IsFloat()) {
        retValue = builder_.CreateConstFloat(0.0f);
      } else {
        retValue = builder_.CreateConstInt(0); // fallback
      }
    }
    builder_.CreateRet(retValue);
  }
  return BlockFlow::Terminated;
}


// if语句（待实现）
IRGenImpl::BlockFlow IRGenImpl::HandleIfStmt(SysYParser::StmtContext* ctx) {
    if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleIfStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
    
    auto* cond = ctx->cond();
    auto* thenStmt = ctx->stmt(0);
    auto* elseStmt = ctx->stmt(1);
    
    // 创建基本块
    auto* thenBlock = func_->CreateBlock("then");
    auto* elseBlock = (ctx->Else() && elseStmt) ? func_->CreateBlock("else") : nullptr;
    auto* mergeBlock = func_->CreateBlock("merge");
    
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] thenBlock: " << thenBlock->GetName() << std::endl;
    if (elseBlock) if (IRGen_DEBUG) std::cout << "[DEBUG IF] elseBlock: " << elseBlock->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] mergeBlock: " << mergeBlock->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] current insert block before cond: " 
              << builder_.GetInsertBlock()->GetName() << std::endl;
    
    // 生成条件
    auto* condValue = EvalCond(*cond);
    
    // 创建条件跳转
    if (elseBlock) {
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] Creating condbr: " << condValue->GetName() 
                  << " -> " << thenBlock->GetName() << ", " << elseBlock->GetName() << std::endl;
        builder_.CreateCondBr(condValue, thenBlock, elseBlock);
    } else {
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] Creating condbr: " << condValue->GetName() 
                  << " -> " << thenBlock->GetName() << ", " << mergeBlock->GetName() << std::endl;
        builder_.CreateCondBr(condValue, thenBlock, mergeBlock);
    }
    
    // 生成 then 分支
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] Generating then branch in block: " << thenBlock->GetName() << std::endl;
    builder_.SetInsertPoint(thenBlock);
    auto thenResult = thenStmt->accept(this);
    bool thenTerminated = (std::any_cast<BlockFlow>(thenResult) == BlockFlow::Terminated);
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] then branch terminated: " << thenTerminated << std::endl;
    
    if (!thenTerminated) {
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] Adding br to merge block from then" << std::endl;
        builder_.CreateBr(mergeBlock);
    }
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] then block has terminator: " << thenBlock->HasTerminator() << std::endl;
    
    // 生成 else 分支
    bool elseTerminated = false;
    if (elseBlock) {
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] Generating else branch in block: " << elseBlock->GetName() << std::endl;
        builder_.SetInsertPoint(elseBlock);
        auto elseResult = elseStmt->accept(this);
        elseTerminated = (std::any_cast<BlockFlow>(elseResult) == BlockFlow::Terminated);
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] else branch terminated: " << elseTerminated << std::endl;
        
        if (!elseTerminated) {
            if (IRGen_DEBUG) std::cout << "[DEBUG IF] Adding br to merge block from else" << std::endl;
            builder_.CreateBr(mergeBlock);
        }
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] else block has terminator: " << elseBlock->HasTerminator() << std::endl;
    }
    
    // 决定后续插入点
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] thenTerminated=" << thenTerminated 
              << ", elseTerminated=" << elseTerminated << std::endl;
    
    if (elseBlock) {
        if (thenTerminated && elseTerminated) {
            auto* afterIfBlock = func_->CreateBlock("after.if");
            if (IRGen_DEBUG) std::cout << "[DEBUG IF] Both branches terminated, creating new block: " 
                      << afterIfBlock->GetName() << std::endl;
            builder_.SetInsertPoint(afterIfBlock);
        } else {
            if (IRGen_DEBUG) std::cout << "[DEBUG IF] Setting insert point to merge block: " 
                      << mergeBlock->GetName() << std::endl;
            builder_.SetInsertPoint(mergeBlock);
        }
    } else {
        if (IRGen_DEBUG) std::cout << "[DEBUG IF] No else, setting insert point to merge block: " 
                  << mergeBlock->GetName() << std::endl;
        builder_.SetInsertPoint(mergeBlock);
    }
    
    if (IRGen_DEBUG) std::cout << "[DEBUG IF] Final insert block: " 
              << builder_.GetInsertBlock()->GetName() << std::endl;
    
    return BlockFlow::Continue;
}

// while语句（待实现）IRGenImpl::BlockFlow IRGenImpl::HandleWhileStmt(SysYParser::StmtContext* ctx) {
IRGenImpl::BlockFlow IRGenImpl::HandleWhileStmt(SysYParser::StmtContext* ctx) {
    if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleWhileStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
    
    if (!ctx || !ctx->cond() || !ctx->stmt(0)) {
        throw std::runtime_error(FormatError("irgen", "非法 while 语句"));
    }
    
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Current insert block before while: " 
              << builder_.GetInsertBlock()->GetName() << std::endl;
    
    auto* condBlock = func_->CreateBlock("while.cond");
    auto* bodyBlock = func_->CreateBlock("while.body");
    auto* exitBlock = func_->CreateBlock("while.exit");
    
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] condBlock: " << condBlock->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] bodyBlock: " << bodyBlock->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] exitBlock: " << exitBlock->GetName() << std::endl;
    
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Adding br to condBlock from current block" << std::endl;
    builder_.CreateBr(condBlock);
    
    loopStack_.push_back({condBlock, bodyBlock, exitBlock});
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] loopStack size: " << loopStack_.size() << std::endl;
    
    // 条件块
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Generating condition in block: " << condBlock->GetName() << std::endl;
    builder_.SetInsertPoint(condBlock);
    auto* condValue = EvalCond(*ctx->cond());
    builder_.CreateCondBr(condValue, bodyBlock, exitBlock);
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] condBlock has terminator: " << condBlock->HasTerminator() << std::endl;
    
    // 循环体
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Generating body in block: " << bodyBlock->GetName() << std::endl;
    builder_.SetInsertPoint(bodyBlock);
    auto bodyResult = ctx->stmt(0)->accept(this);
    bool bodyTerminated = (std::any_cast<BlockFlow>(bodyResult) == BlockFlow::Terminated);
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] body terminated: " << bodyTerminated << std::endl;
    
    if (!bodyTerminated) {
        if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Adding br to condBlock from body" << std::endl;
        builder_.CreateBr(condBlock);
    }
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] bodyBlock has terminator: " << bodyBlock->HasTerminator() << std::endl;
    
    loopStack_.pop_back();
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] loopStack size after pop: " << loopStack_.size() << std::endl;
    
    // 设置插入点为 exitBlock
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] Setting insert point to exitBlock: " << exitBlock->GetName() << std::endl;
    builder_.SetInsertPoint(exitBlock);
    if (IRGen_DEBUG) std::cout << "[DEBUG WHILE] exitBlock has terminator before return: " 
              << exitBlock->HasTerminator() << std::endl;
    
    return BlockFlow::Continue;
}

// break语句（待实现）
IRGenImpl::BlockFlow IRGenImpl::HandleBreakStmt(SysYParser::StmtContext* ctx) {
    if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleBreakStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
    
    if (loopStack_.empty()) {
        throw std::runtime_error(FormatError("irgen", "break 语句不在循环中"));
    }
    
    if (IRGen_DEBUG) std::cout << "[DEBUG BREAK] Current insert block before break: " 
              << builder_.GetInsertBlock()->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG BREAK] Breaking to exitBlock: " 
              << loopStack_.back().exitBlock->GetName() << std::endl;
    
    // 跳转到循环退出块
    builder_.CreateBr(loopStack_.back().exitBlock);
    
    // break 本身就是终止当前路径，后续 unreachable 代码不需要继续生成。
    return BlockFlow::Terminated;
}

IRGenImpl::BlockFlow IRGenImpl::HandleContinueStmt(SysYParser::StmtContext* ctx) {
    if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleContinueStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
    
    if (loopStack_.empty()) {
        throw std::runtime_error(FormatError("irgen", "continue 语句不在循环中"));
    }
    
    if (IRGen_DEBUG) std::cout << "[DEBUG CONTINUE] Current insert block before continue: " 
              << builder_.GetInsertBlock()->GetName() << std::endl;
    if (IRGen_DEBUG) std::cout << "[DEBUG CONTINUE] Continuing to condBlock: " 
              << loopStack_.back().condBlock->GetName() << std::endl;
    
    // 跳转到循环条件块
    builder_.CreateBr(loopStack_.back().condBlock);
    
    // continue 本身就是终止当前路径，后续 unreachable 代码不需要继续生成。
    return BlockFlow::Terminated;
}

// 赋值语句（待实现）
// 赋值语句
// 赋值语句
IRGenImpl::BlockFlow IRGenImpl::HandleAssignStmt(SysYParser::StmtContext* ctx) {
  if (IRGen_DEBUG) std::cout << "[DEBUG IRGEN] HandleAssignStmt: " << (ctx ? ctx->getText() : "<null>") << std::endl;
  
  if (!ctx || !ctx->lVal() || !ctx->exp()) {
    throw std::runtime_error(FormatError("irgen", "非法赋值语句"));
  }

  // 计算右值
  ir::Value* rhs = EvalExpr(*ctx->exp());
  if (!rhs) {
    throw std::runtime_error(FormatError("irgen", "赋值 RHS 计算失败"));
  }
  
  auto* lval = ctx->lVal();
  std::string varName = lval->Ident()->getText();
  if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: assigning to " << varName << std::endl;
  
  // 1. 检查是否为常量（不能给常量赋值）
  auto* const_decl = sema_.ResolveConstUse(lval);
  if (const_decl) {
    throw std::runtime_error(
        FormatError("irgen", "不能给常量赋值: " + varName));
  }
  
  // 2. 查找存储位置
  ir::Value* base_ptr = nullptr;
  
  // 2.1 尝试通过语义分析获取变量定义，并从 storage_map_ 查找
  auto* var_decl = sema_.ResolveVarUse(lval);
  if (var_decl) {
    auto it = storage_map_.find(var_decl);
    if (it != storage_map_.end()) {
      base_ptr = it->second;
      if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: found in storage_map_ for " << varName 
                << ", ptr = " << (void*)base_ptr << std::endl;
    }
  }
  
  // 2.2 从参数映射查找（关键！）
  if (!base_ptr) {
    auto it2 = param_map_.find(varName);
    if (it2 != param_map_.end()) {
      base_ptr = it2->second;
      if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: found in param_map_ for " << varName 
                << ", ptr = " << (void*)base_ptr << std::endl;
    }
  }
  
  // 2.3 从全局变量映射查找
  if (!base_ptr) {
    auto it3 = global_map_.find(varName);
    if (it3 != global_map_.end()) {
      base_ptr = it3->second;
      if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: found in global_map_ for " << varName 
                << ", ptr = " << (void*)base_ptr << std::endl;
    }
  }
  
  // 2.4 从局部变量映射查找（fallback）
  if (!base_ptr) {
    auto it4 = local_var_map_.find(varName);
    if (it4 != local_var_map_.end()) {
      base_ptr = it4->second;
      if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: found in local_var_map_ for " << varName 
                << ", ptr = " << (void*)base_ptr << std::endl;
    }
  }
  
  // 如果还是找不到，才报错
  if (!base_ptr) {
    throw std::runtime_error(
        FormatError("irgen", "变量声明缺少存储槽位: " + varName));
  }
  
  // 3. 检查是否有数组下标
  auto exp_list = lval->exp();
  if (!exp_list.empty()) {
    // 数组元素赋值
    std::vector<ir::Value*> indices;
    indices.push_back(builder_.CreateConstInt(0));
    
    for (auto* exp : exp_list) {
      ir::Value* index = EvalExpr(*exp);
      indices.push_back(index);
    }
    
    ir::Value* elem_ptr = builder_.CreateGEP(
        base_ptr, indices, module_.GetContext().NextTemp());
    
    builder_.CreateStore(rhs, elem_ptr);
  } else {
    // 普通标量赋值
    if (IRGen_DEBUG) std::cerr << "[DEBUG] HandleAssignStmt: scalar assignment to " << varName 
              << ", ptr = " << (void*)base_ptr 
              << ", rhs = " << (void*)rhs << std::endl;
    // 在 HandleAssignStmt 中，存储前添加类型调试
    if (base_ptr && base_ptr->GetType()) {
        
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Is int32: " << base_ptr->GetType()->IsInt32() << std::endl;
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Is float: " << base_ptr->GetType()->IsFloat() << std::endl;
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Is ptr int32: " << base_ptr->GetType()->IsPtrInt32() << std::endl;
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Is ptr float: " << base_ptr->GetType()->IsPtrFloat() << std::endl;
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Is array: " << base_ptr->GetType()->IsArray() << std::endl;
    }
    if (rhs && rhs->GetType()) {
        
        if (IRGen_DEBUG) std::cerr << "[DEBUG] Value is int32: " << rhs->GetType()->IsInt32() << std::endl;
    }
        builder_.CreateStore(rhs, base_ptr);
      }
  
  return BlockFlow::Continue;
}