#include "irgen/IRGen.h"

#include <cstdint>
#include <cstring>
#include <stdexcept>

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

std::any IRGenImpl::visitBlockStmt(SysYParser::BlockStmtContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少语句块"));
  }
  const auto saved_const_env = const_env_;
  const auto saved_const_float_env = const_float_env_;
  for (auto* item : ctx->blockItem()) {
    if (item) {
      if (VisitBlockItemResult(*item) == BlockFlow::Terminated) {
        break;
      }
    }
  }
  const_env_ = saved_const_env;
  const_float_env_ = saved_const_float_env;
  return {};
}

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::visitDecl(SysYParser::DeclContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少变量声明"));
  }
  if (ctx->constDecl()) {
    return ctx->constDecl()->accept(this);
  }
  if (ctx->varDecl()) {
    return ctx->varDecl()->accept(this);
  }
  return {};
}

// ─── 工具：扁平化 constInitValue ──────────────────────────────────────────
// 将嵌套的 const 初始化列表展开为长度 total 的整数数组。
// 遵循 C99 数组初始化规则：
//   - 标量直接填一格
//   - 大括号子列表对齐到 sub_size 边界，填满后补零

void IRGenImpl::FlattenConstInit(SysYParser::ConstInitValueContext* ctx,
                                  const std::vector<int>& dims, int dim_idx,
                                  std::vector<int>& out, int& pos) {
  if (!ctx) return;

  if (ctx->constExp()) {
    // 标量叶节点
    out[pos++] = EvalConstExpr(ctx->constExp());
    return;
  }

  // 大括号列表
  int sub_size = 1;
  for (int i = dim_idx + 1; i < (int)dims.size(); i++) sub_size *= dims[i];
  int agg_size = (dim_idx < (int)dims.size()) ? dims[dim_idx] * sub_size : 1;
  int start = pos;

  for (auto* item : ctx->constInitValue()) {
    if (!item || pos >= start + agg_size) break;
    if (item->constExp()) {
      // 标量：直接填当前位置
      out[pos++] = EvalConstExpr(item->constExp());
    } else {
      // 嵌套大括号：对齐到 sub_size 边界
      if (sub_size > 1) {
        int offset = pos - start;
        int rem = offset % sub_size;
        if (rem != 0) pos += sub_size - rem;
      }
      int sub_start = pos;
      FlattenConstInit(item, dims, dim_idx + 1, out, pos);
      // 补零到子聚合末尾
      int sub_end = sub_start + sub_size;
      while (pos < sub_end && pos < start + agg_size) out[pos++] = 0;
    }
  }
  // 剩余补零
  while (pos < start + agg_size) out[pos++] = 0;
}

void IRGenImpl::FlattenConstInitFloat(SysYParser::ConstInitValueContext* ctx,
                                      const std::vector<int>& dims, int dim_idx,
                                      std::vector<float>& out, int& pos) {
  if (!ctx) return;

  if (ctx->constExp()) {
    out[pos++] = EvalConstExprAsFloat(ctx->constExp());
    return;
  }

  int sub_size = 1;
  for (int i = dim_idx + 1; i < (int)dims.size(); i++) sub_size *= dims[i];
  int agg_size = (dim_idx < (int)dims.size()) ? dims[dim_idx] * sub_size : 1;
  int start = pos;

  for (auto* item : ctx->constInitValue()) {
    if (!item || pos >= start + agg_size) break;
    if (item->constExp()) {
      out[pos++] = EvalConstExprAsFloat(item->constExp());
    } else {
      if (sub_size > 1) {
        int offset = pos - start;
        int rem = offset % sub_size;
        if (rem != 0) pos += sub_size - rem;
      }
      int sub_start = pos;
      FlattenConstInitFloat(item, dims, dim_idx + 1, out, pos);
      int sub_end = sub_start + sub_size;
      while (pos < sub_end && pos < start + agg_size) out[pos++] = 0.0f;
    }
  }
  while (pos < start + agg_size) out[pos++] = 0.0f;
}

// ─── 工具：扁平化 initValue ───────────────────────────────────────────────
void IRGenImpl::FlattenInit(SysYParser::InitValueContext* ctx,
                             const std::vector<int>& dims, int dim_idx,
                             std::vector<ir::Value*>& out, int& pos) {
  if (!ctx) return;

  if (ctx->exp()) {
    out[pos++] = EvalExpr(*ctx->exp());
    return;
  }

  int sub_size = 1;
  for (int i = dim_idx + 1; i < (int)dims.size(); i++) sub_size *= dims[i];
  int agg_size = (dim_idx < (int)dims.size()) ? dims[dim_idx] * sub_size : 1;
  int start = pos;

  for (auto* item : ctx->initValue()) {
    if (!item || pos >= start + agg_size) break;
    if (item->exp()) {
      out[pos++] = EvalExpr(*item->exp());
    } else {
      if (sub_size > 1) {
        int offset = pos - start;
        int rem = offset % sub_size;
        if (rem != 0) pos += sub_size - rem;  // zeros already in out
      }
      int sub_start = pos;
      FlattenInit(item, dims, dim_idx + 1, out, pos);
      int sub_end = sub_start + sub_size;
      while (pos < sub_end && pos < start + agg_size) pos++;  // zeros
    }
  }
  while (pos < start + agg_size) pos++;  // zeros
}

void IRGenImpl::FlattenGlobalInitInt(SysYParser::InitValueContext* ctx,
                                     const std::vector<int>& dims, int dim_idx,
                                     std::vector<int>& out, int& pos) {
  if (!ctx) return;

  if (ctx->exp()) {
    out[pos++] = EvalExpAsConst(ctx->exp());
    return;
  }

  int sub_size = 1;
  for (int i = dim_idx + 1; i < (int)dims.size(); i++) sub_size *= dims[i];
  int agg_size = (dim_idx < (int)dims.size()) ? dims[dim_idx] * sub_size : 1;
  int start = pos;

  for (auto* item : ctx->initValue()) {
    if (!item || pos >= start + agg_size) break;
    if (item->exp()) {
      out[pos++] = EvalExpAsConst(item->exp());
    } else {
      if (sub_size > 1) {
        int offset = pos - start;
        int rem = offset % sub_size;
        if (rem != 0) pos += sub_size - rem;
      }
      int sub_start = pos;
      FlattenGlobalInitInt(item, dims, dim_idx + 1, out, pos);
      int sub_end = sub_start + sub_size;
      while (pos < sub_end && pos < start + agg_size) out[pos++] = 0;
    }
  }
  while (pos < start + agg_size) out[pos++] = 0;
}

void IRGenImpl::FlattenGlobalInitFloat(SysYParser::InitValueContext* ctx,
                                       const std::vector<int>& dims,
                                       int dim_idx, std::vector<float>& out,
                                       int& pos) {
  if (!ctx) return;

  if (ctx->exp()) {
    out[pos++] = EvalExpAsConstFloat(ctx->exp());
    return;
  }

  int sub_size = 1;
  for (int i = dim_idx + 1; i < (int)dims.size(); i++) sub_size *= dims[i];
  int agg_size = (dim_idx < (int)dims.size()) ? dims[dim_idx] * sub_size : 1;
  int start = pos;

  for (auto* item : ctx->initValue()) {
    if (!item || pos >= start + agg_size) break;
    if (item->exp()) {
      out[pos++] = EvalExpAsConstFloat(item->exp());
    } else {
      if (sub_size > 1) {
        int offset = pos - start;
        int rem = offset % sub_size;
        if (rem != 0) pos += sub_size - rem;
      }
      int sub_start = pos;
      FlattenGlobalInitFloat(item, dims, dim_idx + 1, out, pos);
      int sub_end = sub_start + sub_size;
      while (pos < sub_end && pos < start + agg_size) out[pos++] = 0.0f;
    }
  }
  while (pos < start + agg_size) out[pos++] = 0.0f;
}

// ─── const 声明 ───────────────────────────────────────────────────────────

std::any IRGenImpl::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  if (!ctx) return {};
  if (!ctx->btype()) {
    throw std::runtime_error(FormatError("irgen", "缺少类型声明"));
  }
  if (ctx->btype()->INT()) {
    current_decl_type_ = ir::Type::GetInt32Type();
  } else if (ctx->btype()->FLOAT()) {
    current_decl_type_ = ir::Type::GetFloat32Type();
  } else {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float const 声明"));
  }
  for (auto* def : ctx->constDef()) {
    if (def) def->accept(this);
  }
  current_decl_type_ = nullptr;
  return {};
}

std::any IRGenImpl::visitConstDef(SysYParser::ConstDefContext* ctx) {
  if (!ctx || !ctx->ID()) return {};
  const std::string name = ctx->ID()->getText();

  // ── 标量 const ────────────────────────────────────────────────────────
  if (ctx->LBRACK().empty()) {
    if (!ctx->constInitValue() || !ctx->constInitValue()->constExp()) {
      throw std::runtime_error(FormatError("irgen", "const 标量声明缺少初始值"));
    }
    const bool is_float_const = current_decl_type_ && current_decl_type_->IsFloat32();
    if (is_float_const) {
      float fval = EvalConstExprAsFloat(ctx->constInitValue()->constExp());
      const_float_env_[name] = fval;

      if (IsGlobalScope()) {
        std::int32_t bits = 0;
        std::memcpy(&bits, &fval, sizeof(bits));
        auto* gv = module_.CreateGlobalVar(
            name, static_cast<int>(bits), 1, ir::Type::GetPtrFloat32Type());
        global_storage_[name] = gv;
      } else {
        auto* slot = CreateEntryAllocaF32(module_.GetContext().NextTemp());
        named_storage_[name] = slot;
        builder_.CreateStore(module_.GetContext().GetConstFloat(fval), slot);
      }
    } else {
      int ival = EvalConstExpr(ctx->constInitValue()->constExp());
      const_env_[name] = ival;  // 存入编译期环境

      if (IsGlobalScope()) {
        auto* gv = module_.CreateGlobalVar(name, ival);
        global_storage_[name] = gv;
      } else {
        auto* slot = CreateEntryAllocaI32(module_.GetContext().NextTemp());
        named_storage_[name] = slot;
        builder_.CreateStore(builder_.CreateConstInt(ival), slot);
      }
    }
    return {};
  }

  // ── 数组 const ────────────────────────────────────────────────────────
  std::vector<int> dims;
  for (auto* ce : ctx->constExp()) {
    dims.push_back(EvalConstExpr(ce));
  }
  int total = 1;
  for (int d : dims) total *= d;

  const bool is_float_const = current_decl_type_ && current_decl_type_->IsFloat32();

  if (is_float_const) {
    std::vector<float> flat(total, 0.0f);
    if (ctx->constInitValue()) {
      int pos = 0;
      FlattenConstInitFloat(ctx->constInitValue(), dims, 0, flat, pos);
    }
    std::vector<int> init_bits;
    init_bits.reserve(flat.size());
    for (float v : flat) {
      std::int32_t bits = 0;
      std::memcpy(&bits, &v, sizeof(bits));
      init_bits.push_back(static_cast<int>(bits));
    }

    if (IsGlobalScope()) {
      auto* gv = module_.CreateGlobalVar(
          name, 0, total, ir::Type::GetPtrFloat32Type(), std::move(init_bits));
      global_storage_[name] = gv;
      global_array_dims_[name] = dims;
    } else {
      auto* slot = CreateEntryAllocaF32Array(total, module_.GetContext().NextTemp());
      named_storage_[name] = slot;
      local_array_dims_[name] = dims;
      for (int i = 0; i < total; i++) {
        auto* idx = builder_.CreateConstInt(i);
        auto* ptr = builder_.CreateGep(slot, idx, module_.GetContext().NextTemp());
        builder_.CreateStore(module_.GetContext().GetConstFloat(flat[i]), ptr);
      }
    }
    return {};
  }

  // 扁平化初始化值
  std::vector<int> flat(total, 0);
  if (ctx->constInitValue()) {
    int pos = 0;
    FlattenConstInit(ctx->constInitValue(), dims, 0, flat, pos);
  }

  if (IsGlobalScope()) {
    auto* gv = module_.CreateGlobalVar(name, 0, total,
                                       ir::Type::GetPtrInt32Type(),
                                       std::move(flat));
    global_storage_[name] = gv;
    global_array_dims_[name] = dims;
  } else {
    // 局部 const 数组：alloca + 逐元素 store
    auto* slot = CreateEntryAllocaArray(total, module_.GetContext().NextTemp());
    named_storage_[name] = slot;
    local_array_dims_[name] = dims;
    for (int i = 0; i < total; i++) {
      auto* idx = builder_.CreateConstInt(i);
      auto* ptr = builder_.CreateGep(slot, idx, module_.GetContext().NextTemp());
      builder_.CreateStore(builder_.CreateConstInt(flat[i]), ptr);
    }
  }
  return {};
}

// ─── var 声明 ─────────────────────────────────────────────────────────────

std::any IRGenImpl::visitVarDecl(SysYParser::VarDeclContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少变量声明"));
  }
  if (!ctx->btype()) {
    throw std::runtime_error(FormatError("irgen", "缺少类型声明"));
  }
  // 设置当前声明类型
  if (ctx->btype()->INT()) {
    current_decl_type_ = ir::Type::GetInt32Type();
  } else if (ctx->btype()->FLOAT()) {
    current_decl_type_ = ir::Type::GetFloat32Type();
  } else {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int/float 变量声明"));
  }

  for (auto* var_def : ctx->varDef()) {
    if (!var_def) {
      throw std::runtime_error(FormatError("irgen", "非法变量声明"));
    }
    var_def->accept(this);
  }
  current_decl_type_ = nullptr; // 清理
  return {};
}

std::any IRGenImpl::visitVarDef(SysYParser::VarDefContext* ctx) {
  if (!ctx || !ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "变量定义缺少名称"));
  }
  const std::string name = ctx->ID()->getText();

  // ── 数组变量 ──────────────────────────────────────────────────────────
  if (!ctx->LBRACK().empty()) {
    std::vector<int> dims;
    for (auto* ce : ctx->constExp()) {
      dims.push_back(EvalConstExpr(ce));
    }
    int total = 1;
    for (int d : dims) total *= d;

    if (IsGlobalScope()) {
      std::vector<int> init_elems;
      if (auto* init_val = ctx->initValue()) {
        if (current_decl_type_->IsFloat32()) {
          std::vector<float> flat(total, 0.0f);
          int pos = 0;
          FlattenGlobalInitFloat(init_val, dims, 0, flat, pos);
          init_elems.reserve(flat.size());
          for (float v : flat) {
            std::int32_t bits = 0;
            std::memcpy(&bits, &v, sizeof(bits));
            init_elems.push_back(static_cast<int>(bits));
          }
        } else {
          init_elems.assign(total, 0);
          int pos = 0;
          FlattenGlobalInitInt(init_val, dims, 0, init_elems, pos);
        }
      }
      auto* gv = module_.CreateGlobalVar(
          name, 0, total,
          current_decl_type_->IsFloat32() ? ir::Type::GetPtrFloat32Type()
                                          : ir::Type::GetPtrInt32Type(),
          std::move(init_elems));
      storage_map_[ctx] = gv;
      global_storage_[name] = gv;
      global_array_dims_[name] = dims;
    } else {
      // 根据当前声明类型创建数组alloca
      ir::AllocaInst* slot;
      if (current_decl_type_->IsFloat32()) {
        slot = CreateEntryAllocaF32Array(total, module_.GetContext().NextTemp());
      } else {
        slot = CreateEntryAllocaArray(total, module_.GetContext().NextTemp());
      }
      storage_map_[ctx] = slot;
      named_storage_[name] = slot;
      local_array_dims_[name] = dims;

      // 先零初始化：float 数组走 memset，int 数组维持逐元素 store。
      if (current_decl_type_->IsFloat32()) {
        if (total > 0) {
          auto* memset_fn = module_.FindFunction("memset");
          if (!memset_fn) {
            memset_fn = module_.CreateFunction(
                "memset", ir::Type::GetVoidType(),
                {ir::Type::GetPtrFloat32Type(), ir::Type::GetInt32Type(),
                 ir::Type::GetInt32Type()});
            memset_fn->SetExternal(true);
          }
          builder_.CreateCall(
              memset_fn,
              {slot, builder_.CreateConstInt(0), builder_.CreateConstInt(total * 4)},
              module_.GetContext().NextTemp());
        }
      } else {
        for (int i = 0; i < total; i++) {
          auto* idx = builder_.CreateConstInt(i);
          auto* ptr = builder_.CreateGep(slot, idx, module_.GetContext().NextTemp());
          builder_.CreateStore(builder_.CreateConstInt(0), ptr);
        }
      }
      // 如果有初始化列表，覆盖零
      if (auto* init_val = ctx->initValue()) {
        std::vector<ir::Value*> flat(total, nullptr);
        int pos = 0;
        FlattenInit(init_val, dims, 0, flat, pos);
        for (int i = 0; i < total; i++) {
          if (flat[i] != nullptr) {
            auto* idx = builder_.CreateConstInt(i);
            auto* ptr = builder_.CreateGep(slot, idx, module_.GetContext().NextTemp());
            ir::Value* val = flat[i];
            if (ptr->GetType()->IsPtrFloat32()) {
              val = CastToFloat(val);
            } else {
              val = CastToInt(val);
            }
            builder_.CreateStore(val, ptr);
          }
        }
      }
    }
    return {};
  }

  // ── 标量变量 ──────────────────────────────────────────────────────────
  if (IsGlobalScope()) {
    int init_bits_or_int = 0;
    if (current_decl_type_->IsFloat32()) {
      float fval = 0.0f;
      if (auto* init_value = ctx->initValue()) {
        if (!init_value->exp()) {
          throw std::runtime_error(
              FormatError("irgen", "全局标量变量仅支持表达式初始化"));
        }
        fval = EvalExpAsConstFloat(init_value->exp());
      }
      std::int32_t bits = 0;
      std::memcpy(&bits, &fval, sizeof(bits));
      init_bits_or_int = static_cast<int>(bits);
    } else {
      if (auto* init_value = ctx->initValue()) {
        if (!init_value->exp()) {
          throw std::runtime_error(
              FormatError("irgen", "全局标量变量仅支持表达式初始化"));
        }
        init_bits_or_int = EvalExpAsConst(init_value->exp());
      }
    }
    auto* gv = module_.CreateGlobalVar(
        name, init_bits_or_int, 1,
        current_decl_type_->IsFloat32() ? ir::Type::GetPtrFloat32Type()
                                        : ir::Type::GetPtrInt32Type());
    storage_map_[ctx] = gv;
    global_storage_[name] = gv;
    return {};
  }

  // 局部标量
  if (storage_map_.find(ctx) != storage_map_.end()) {
    throw std::runtime_error(FormatError("irgen", "声明重复生成存储槽位"));
  }

  // 根据当前声明类型创建alloca
  ir::AllocaInst* slot;
  if (current_decl_type_->IsFloat32()) {
    slot = CreateEntryAllocaF32(module_.GetContext().NextTemp());
  } else {
    slot = CreateEntryAllocaI32(module_.GetContext().NextTemp());
  }
  storage_map_[ctx] = slot;
  named_storage_[name] = slot;

  ir::Value* init = nullptr;
  if (auto* init_value = ctx->initValue()) {
    if (!init_value->exp()) {
      throw std::runtime_error(FormatError("irgen", "当前不支持聚合初始化到标量"));
    }
    init = EvalExpr(*init_value->exp());
  } else {
    if (current_decl_type_->IsFloat32()) {
      init = module_.GetContext().GetConstFloat(0.0f);
    } else {
      init = builder_.CreateConstInt(0);
    }
  }
  if (current_decl_type_->IsFloat32()) {
    init = CastToFloat(init);
  } else {
    init = CastToInt(init);
  }
  builder_.CreateStore(init, slot);
  return {};
}