#include "irgen/IRGen.h"

#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", "缺少语句块"));
  }
  for (auto* item : ctx->blockItem()) {
    if (item) {
      if (VisitBlockItemResult(*item) == BlockFlow::Terminated) {
        break;
      }
    }
  }
  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;
}

// ─── 工具：扁平化 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
}

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

std::any IRGenImpl::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  if (!ctx) return {};
  if (!ctx->btype() || !ctx->btype()->INT()) {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int const 声明"));
  }
  for (auto* def : ctx->constDef()) {
    if (def) def->accept(this);
  }
  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 标量声明缺少初始值"));
    }
    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;

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

  if (IsGlobalScope()) {
    // 全局 const 数组：创建全局数组变量（仅支持零初始化；非零初始化暂用零）
    // TODO: 支持全局 const 数组的非零初始化
    auto* gv = module_.CreateGlobalVar(name, 0, total);
    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() || !ctx->btype()->INT()) {
    throw std::runtime_error(FormatError("irgen", "当前仅支持 int 变量声明"));
  }
  for (auto* var_def : ctx->varDef()) {
    if (!var_def) {
      throw std::runtime_error(FormatError("irgen", "非法变量声明"));
    }
    var_def->accept(this);
  }
  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()) {
      auto* gv = module_.CreateGlobalVar(name, 0, total);
      storage_map_[ctx] = gv;
      global_storage_[name] = gv;
      global_array_dims_[name] = dims;
      // 全局数组：不支持运行时初始化（全零已足够）
    } else {
      auto* slot = CreateEntryAllocaArray(total, module_.GetContext().NextTemp());
      storage_map_[ctx] = slot;
      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(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());
            builder_.CreateStore(flat[i], ptr);
          }
        }
      }
    }
    return {};
  }

  // ── 标量变量 ──────────────────────────────────────────────────────────
  if (IsGlobalScope()) {
    int ival = 0;
    if (auto* init_value = ctx->initValue()) {
      if (!init_value->exp()) {
        throw std::runtime_error(
            FormatError("irgen", "全局标量变量仅支持表达式初始化"));
      }
      ival = EvalExpAsConst(init_value->exp());
    }
    auto* gv = module_.CreateGlobalVar(name, ival);
    storage_map_[ctx] = gv;
    global_storage_[name] = gv;
    return {};
  }

  // 局部标量
  if (storage_map_.find(ctx) != storage_map_.end()) {
    throw std::runtime_error(FormatError("irgen", "声明重复生成存储槽位"));
  }
  auto* 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 {
    init = builder_.CreateConstInt(0);
  }
  builder_.CreateStore(init, slot);
  return {};
}
