#include "irgen/IRGen.h"

#include <stdexcept>

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

namespace {
ir::ConstantValue* BuildConstantArray(ir::Context& ctx, std::shared_ptr<ir::Type> type, 
                                      const std::vector<ir::ConstantValue*>& flattened, 
                                      size_t& pos) {
  if (!type->IsArray()) {
    return flattened[pos++];
  }
  std::vector<ir::ConstantValue*> elements;
  for (int i = 0; i < type->GetNumElements(); ++i) {
    elements.push_back(BuildConstantArray(ctx, type->GetElementType(), flattened, pos));
  }
  return ctx.GetConstArray(type, elements);
}
}

std::any IRGenImpl::visitBlock(SysYParser::BlockContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少语句块"));
  }

  // 压入局部作用域
  storage_map_stack_.push_back({});
  const_values_stack_.push_back({});

  bool terminated = false;
  for (auto* item : ctx->blockItem()) {
    if (item) {
      if (VisitBlockItemResult(*item) == BlockFlow::Terminated) {
        terminated = true;
        break;
      }
    }
  }

  // 弹出局部作用域
  storage_map_stack_.pop_back();
  const_values_stack_.pop_back();

  return terminated ? BlockFlow::Terminated : BlockFlow::Continue;
}

IRGenImpl::BlockFlow IRGenImpl::VisitBlockItemResult(
    SysYParser::BlockItemContext& item) {
  return std::any_cast<BlockFlow>(item.accept(this));
}

std::any IRGenImpl::visitBlockItem(SysYParser::BlockItemContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少块内项"));
  }
  if (ctx->decl()) {
    ctx->decl()->accept(this);
    return BlockFlow::Continue;
  }
  if (ctx->stmt()) {
    return ctx->stmt()->accept(this);
  }
  throw std::runtime_error(FormatError("irgen", "暂不支持的语句或声明"));
}

std::any IRGenImpl::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  if (!ctx) return BlockFlow::Continue;
  
  if (!ctx->bType() || (!ctx->bType()->INT() && !ctx->bType()->FLOAT())) {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float 常量声明"));
  }
  
  for (auto* def : ctx->constDef()) {
    if (def) def->accept(this);
  }
  return BlockFlow::Continue;
}

void IRGenImpl::FlattenInitVal(SysYParser::InitValContext* ctx, 
                               const std::vector<int>& dims, 
                               const std::vector<int>& sub_sizes,
                               int dim_idx,
                               size_t& current_pos,
                               std::vector<ir::Value*>& results,
                               bool is_float) {
  if (ctx->exp()) {
    ir::Value* val = EvalExpr(*ctx->exp());
    // Implicit conversion
    if (is_float && !val->GetType()->IsFloat()) {
      val = builder_.CreateSIToFP(val, module_.GetContext().NextTemp());
    } else if (!is_float && val->GetType()->IsFloat()) {
      val = builder_.CreateFPToSI(val, module_.GetContext().NextTemp());
    }
    results[current_pos++] = val;
  } else {
    // Nested { ... }
    size_t start_pos = current_pos;
    for (auto* item : ctx->initVal()) {
      FlattenInitVal(item, dims, sub_sizes, dim_idx + 1, current_pos, results, is_float);
    }
    // Fill remaining with 0
    size_t end_pos = start_pos + sub_sizes[dim_idx];
    while (current_pos < end_pos) {
      results[current_pos++] = is_float ? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) 
                                        : (ir::Value*)module_.GetContext().GetConstInt(0);
    }
  }
}

void IRGenImpl::FlattenConstInitVal(SysYParser::ConstInitValContext* ctx, 
                                    const std::vector<int>& dims, 
                                    const std::vector<int>& sub_sizes,
                                    int dim_idx,
                                    size_t& current_pos,
                                    std::vector<ir::ConstantValue*>& results,
                                    bool is_float) {
  if (ctx->constExp()) {
    ir::Value* val = std::any_cast<ir::Value*>(ctx->constExp()->accept(this));
    ir::ConstantValue* cval = dynamic_cast<ir::ConstantValue*>(val);
    if (!cval) throw std::runtime_error("Not a constant expression");
    
    // Constant conversion
    if (is_float && dynamic_cast<ir::ConstantInt*>(cval)) {
      cval = module_.GetContext().GetConstFloat((float)static_cast<ir::ConstantInt*>(cval)->GetValue());
    } else if (!is_float && dynamic_cast<ir::ConstantFloat*>(cval)) {
      cval = module_.GetContext().GetConstInt((int)static_cast<ir::ConstantFloat*>(cval)->GetValue());
    }
    results[current_pos++] = cval;
  } else {
    size_t start_pos = current_pos;
    for (auto* item : ctx->constInitVal()) {
      FlattenConstInitVal(item, dims, sub_sizes, dim_idx + 1, current_pos, results, is_float);
    }
    // Fill remaining with 0
    size_t end_pos = start_pos + sub_sizes[dim_idx];
    while (current_pos < end_pos) {
      results[current_pos++] = is_float ? (ir::ConstantValue*)module_.GetContext().GetConstFloat(0.0f) 
                                        : (ir::ConstantValue*)module_.GetContext().GetConstInt(0);
    }
  }
}

std::any IRGenImpl::visitConstDef(SysYParser::ConstDefContext* ctx) {
  if (!ctx || !ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "常量定义缺少名称"));
  }

  std::string var_name = ctx->ID()->getText();
  
  // Get dimensions
  std::vector<int> dims;
  for (auto* idx : ctx->constIndex()) {
    dims.push_back(EvaluateConstInt(idx->constExp()));
  }

  bool is_float = false;
  auto* parent_decl = dynamic_cast<SysYParser::ConstDeclContext*>(ctx->parent);
  if (parent_decl && parent_decl->bType() && parent_decl->bType()->FLOAT()) {
    is_float = true;
  }
  auto base_ty = is_float ? ir::Type::GetFloatType() : ir::Type::GetInt32Type();

  std::shared_ptr<ir::Type> var_ty = base_ty;
  for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
    var_ty = ir::Type::GetArrayType(var_ty, *it);
  }

  std::vector<int> sub_sizes(dims.size() + 1);
  sub_sizes[dims.size()] = 1;
  for (int i = (int)dims.size() - 1; i >= 0; --i) {
    sub_sizes[i] = sub_sizes[i+1] * dims[i];
  }

  ir::ConstantValue* init_const = nullptr;
  std::vector<ir::ConstantValue*> flattened;
  if (dims.empty()) {
    if (auto* init_val = ctx->constInitVal()) {
      if (init_val->constExp()) {
        ir::Value* val = std::any_cast<ir::Value*>(init_val->constExp()->accept(this));
        init_const = dynamic_cast<ir::ConstantValue*>(val);
        // Constant conversion
        if (is_float && dynamic_cast<ir::ConstantInt*>(init_const)) {
          init_const = module_.GetContext().GetConstFloat((float)static_cast<ir::ConstantInt*>(init_const)->GetValue());
        } else if (!is_float && dynamic_cast<ir::ConstantFloat*>(init_const)) {
          init_const = module_.GetContext().GetConstInt((int)static_cast<ir::ConstantFloat*>(init_const)->GetValue());
        }
      }
    }
  } else {
    flattened.resize(sub_sizes[0]);
    if (auto* init_val = ctx->constInitVal()) {
      size_t pos = 0;
      FlattenConstInitVal(init_val, dims, sub_sizes, 0, pos, flattened, is_float);
    } else {
      auto zero = is_float ? (ir::ConstantValue*)module_.GetContext().GetConstFloat(0.0f) : (ir::ConstantValue*)module_.GetContext().GetConstInt(0);
      for (auto& v : flattened) v = zero;
    }
    size_t pos = 0;
    init_const = BuildConstantArray(module_.GetContext(), var_ty, flattened, pos);
  }

  // 记录常量值供后续直接使用 (only for scalars for now)
  if (dims.empty() && !const_values_stack_.empty()) {
    const_values_stack_.back()[var_name] = init_const;
  }

  if (func_ == nullptr) {
    auto gv_ptr_ty = ir::Type::GetPointerType(var_ty);
    auto* gv = module_.CreateGlobalVariable(var_name, gv_ptr_ty, init_const);
    if (!storage_map_stack_.empty()) {
      storage_map_stack_.back()[var_name] = gv;
    }
  } else {
    // 局部作用域 - 确保 alloca 在入口块
    auto* current_bb = builder_.GetInsertBlock();
    builder_.SetInsertPoint(func_->GetEntry());
    ir::Value* slot = builder_.CreateAlloca(var_ty, module_.GetContext().NextTemp());
    builder_.SetInsertPoint(current_bb);
    
    if (!storage_map_stack_.empty()) {
      storage_map_stack_.back()[var_name] = slot;
    }
    if (dims.empty()) {
      if (init_const) builder_.CreateStore(init_const, slot);
    } else {
      for (size_t i = 0; i < flattened.size(); ++i) {
        std::vector<ir::Value*> indices;
        indices.push_back(builder_.CreateConstInt(0));
        size_t temp = i;
        for (size_t d = 0; d < dims.size(); ++d) {
          indices.push_back(builder_.CreateConstInt(temp / sub_sizes[d+1]));
          temp %= sub_sizes[d+1];
        }
        ir::Value* ptr = builder_.CreateGEP(ir::Type::GetPointerType(base_ty), slot, indices, module_.GetContext().NextTemp());
        builder_.CreateStore(flattened[i], ptr);
      }
    }
  }

  return BlockFlow::Continue;
}

// 变量声明的 IR 生成目前也是最小实现：
// - 先检查声明的基础类型，支持 int 和 float；
// - 再把 Decl 中的变量定义交给 visitVarDef 继续处理。
std::any IRGenImpl::visitDecl(SysYParser::DeclContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少变量声明"));
  }
  // 当前语法中 decl 包含 constDecl 或 varDecl
  if (auto* var_decl = ctx->varDecl()) {
    if (!var_decl->bType() || (!var_decl->bType()->INT() && !var_decl->bType()->FLOAT())) {
      throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float 变量声明"));
    }
    for (auto* var_def : var_decl->varDef()) {
      if (var_def) {
        var_def->accept(this);
      }
    }
  } else if (auto* const_decl = ctx->constDecl()) {
    return const_decl->accept(this);
  } else {
    throw std::runtime_error(FormatError("irgen", "当前仅支持变量声明"));
  }
  return BlockFlow::Continue;
}


// 当前仍是教学用的最小版本，因此这里只支持：
// - 局部 int 变量；
// - 标量初始化；
// - 一个 VarDef 对应一个槽位。
std::any IRGenImpl::visitVarDef(SysYParser::VarDefContext* ctx) {
  if (!ctx || !ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "变量声明缺少名称"));
  }
  std::string var_name = ctx->ID()->getText();
  if (!storage_map_stack_.empty() && storage_map_stack_.back().find(var_name) != storage_map_stack_.back().end()) {
    throw std::runtime_error(FormatError("irgen", "声明重复生成存储槽位"));
  }

  // Get dimensions
  std::vector<int> dims;
  for (auto* idx : ctx->constIndex()) {
    dims.push_back(EvaluateConstInt(idx->constExp()));
  }

  // Determine base type
  bool is_float = false;
  auto* parent_decl = dynamic_cast<SysYParser::VarDeclContext*>(ctx->parent);
  if (parent_decl && parent_decl->bType() && parent_decl->bType()->FLOAT()) {
      is_float = true;
  }
  auto base_ty = is_float ? ir::Type::GetFloatType() : ir::Type::GetInt32Type();
  
  std::shared_ptr<ir::Type> var_ty = base_ty;
  for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
    var_ty = ir::Type::GetArrayType(var_ty, *it);
  }

  std::vector<int> sub_sizes(dims.size() + 1);
  sub_sizes[dims.size()] = 1;
  for (int i = (int)dims.size() - 1; i >= 0; --i) {
    sub_sizes[i] = sub_sizes[i+1] * dims[i];
  }

  if (func_ == nullptr) {
    // 全局作用域
    ir::ConstantValue* init_const = nullptr;
    if (dims.empty()) {
      if (auto* init_val = ctx->initVal()) {
        if (init_val->exp()) {
          auto* val = EvalExpr(*init_val->exp());
          init_const = dynamic_cast<ir::ConstantValue*>(val);
          // Constant conversion
          if (is_float && dynamic_cast<ir::ConstantInt*>(init_const)) {
            init_const = module_.GetContext().GetConstFloat((float)static_cast<ir::ConstantInt*>(init_const)->GetValue());
          } else if (!is_float && dynamic_cast<ir::ConstantFloat*>(init_const)) {
            init_const = module_.GetContext().GetConstInt((int)static_cast<ir::ConstantFloat*>(init_const)->GetValue());
          }
        }
      } else {
        init_const = is_float ? (ir::ConstantValue*)module_.GetContext().GetConstFloat(0.0f) : (ir::ConstantValue*)module_.GetContext().GetConstInt(0);
      }
    } else {
      if (auto* init_val = ctx->initVal()) {
      std::vector<ir::ConstantValue*> flattened(sub_sizes[0]);
      // VarDef's InitVal can be an expression or { ... }
      if (init_val->exp()) {
           auto* val = EvalExpr(*init_val->exp());
           auto* cval = dynamic_cast<ir::ConstantValue*>(val);
           flattened[0] = cval;
           auto zero = is_float ? (ir::ConstantValue*)module_.GetContext().GetConstFloat(0.0f) : (ir::ConstantValue*)module_.GetContext().GetConstInt(0);
           for (size_t i = 1; i < flattened.size(); ++i) {
             flattened[i] = zero;
           }
           size_t bpos = 0;
           init_const = BuildConstantArray(module_.GetContext(), var_ty, flattened, bpos);
        } else {
            size_t fpos = 0;
            std::vector<ir::Value*> flat_vals(sub_sizes[0]);
            auto zero = is_float ? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) : (ir::Value*)module_.GetContext().GetConstInt(0);
            for (auto& v : flat_vals) v = zero;
            
            FlattenInitVal(init_val, dims, sub_sizes, 0, fpos, flat_vals, is_float);
            for (size_t i = 0; i < flat_vals.size(); ++i) {
              flattened[i] = dynamic_cast<ir::ConstantValue*>(flat_vals[i]);
            }
            size_t bpos = 0;
            init_const = BuildConstantArray(module_.GetContext(), var_ty, flattened, bpos);
         }
      } else {
        init_const = module_.GetContext().GetConstZero(var_ty);
      }
    }

    auto gv_ptr_ty = ir::Type::GetPointerType(var_ty);
    auto* gv = module_.CreateGlobalVariable(var_name, gv_ptr_ty, init_const);
    if (!storage_map_stack_.empty()) {
      storage_map_stack_.back()[var_name] = gv;
    }
  } else {
    // 局部作用域 - 确保 alloca 在入口块
    auto* current_bb = builder_.GetInsertBlock();
    builder_.SetInsertPoint(func_->GetEntry());
    ir::Value* slot = builder_.CreateAlloca(var_ty, module_.GetContext().NextTemp());
    builder_.SetInsertPoint(current_bb);

    if (!storage_map_stack_.empty()) {
      storage_map_stack_.back()[var_name] = slot;
    }

    if (auto* init_val = ctx->initVal()) {
      if (dims.empty()) {
        if (init_val->exp()) {
          ir::Value* init = EvalExpr(*init_val->exp());
          if (is_float && !init->GetType()->IsFloat()) {
            init = builder_.CreateSIToFP(init, module_.GetContext().NextTemp());
          } else if (!is_float && init->GetType()->IsFloat()) {
            init = builder_.CreateFPToSI(init, module_.GetContext().NextTemp());
          }
          builder_.CreateStore(init, slot);
        }
      } else {
        std::vector<ir::Value*> flattened(sub_sizes[0]);
        auto zero = is_float ? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) : (ir::Value*)module_.GetContext().GetConstInt(0);
        for (auto& v : flattened) v = zero;
        
        size_t pos = 0;
        FlattenInitVal(init_val, dims, sub_sizes, 0, pos, flattened, is_float);
        for (size_t i = 0; i < flattened.size(); ++i) {
          // Optimization: only store non-zero? 
          // For now, store all to be safe.
          std::vector<ir::Value*> indices;
          indices.push_back(builder_.CreateConstInt(0));
          size_t temp = i;
          for (size_t d = 0; d < dims.size(); ++d) {
            indices.push_back(builder_.CreateConstInt(temp / sub_sizes[d+1]));
            temp %= sub_sizes[d+1];
          }
          ir::Value* ptr = builder_.CreateGEP(ir::Type::GetPointerType(base_ty), slot, indices, module_.GetContext().NextTemp());
          builder_.CreateStore(flattened[i], ptr);
        }
      }
    } else {
      // Initialize scalar locals to 0
      if (dims.empty()) {
        ir::Value* zero = is_float ? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) : (ir::Value*)module_.GetContext().GetConstInt(0);
        builder_.CreateStore(zero, slot);
      }
    }
  }
  
  return BlockFlow::Continue;
}
