#include "irgen/IRGen.h"

#include <stdexcept>

#include "SysYParser.h"
#include "ir/IR.h"
#include "utils/Log.h"
#include <functional>

// helper functions removed; VarDef uses Ident() directly per current grammar.

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

std::any IRGenImpl::visitConstDef(SysYParser::ConstDefContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "缺少常量定义"));
  if (!ctx->Ident()) throw std::runtime_error(FormatError("irgen", "常量声明缺少名称"));
  if (!ctx->constInitVal()) throw std::runtime_error(FormatError("irgen", "常量必须初始化"));
  if (storage_map_.find(ctx) != storage_map_.end()) throw std::runtime_error(FormatError("irgen", "声明重复生成存储槽位"));

  auto* slot = (current_btype_ == "float") ?
                   static_cast<ir::AllocaInst*>(builder_.CreateAllocaFloat(module_.GetContext().NextTemp())) :
                   static_cast<ir::AllocaInst*>(builder_.CreateAllocaI32(module_.GetContext().NextTemp()));
  storage_map_[ctx] = slot;
  name_map_[ctx->Ident()->getText()] = slot;

  // Try to evaluate a scalar const initializer
  ir::ConstantValue* cinit = nullptr;
  try {
    auto* initval = ctx->constInitVal();
    if (initval && initval->constExp() && initval->constExp()->addExp()) {
      if (current_btype_ == "float") {
        auto* add = initval->constExp()->addExp();
        float fv = std::stof(add->getText());
        cinit = module_.GetContext().GetConstFloat(fv);
      } else {
        auto* add = initval->constExp()->addExp();
        int iv = std::stoi(add->getText());
        cinit = module_.GetContext().GetConstInt(iv);
      }
    }
  } catch(...) {
    // fallback: try evaluate via visitor
    try {
      auto* add = ctx->constInitVal()->constExp()->addExp();
      ir::Value* v = std::any_cast<ir::Value*>(add->accept(this));
      if (auto* cv = dynamic_cast<ir::ConstantValue*>(v)) cinit = cv;
    } catch(...) {}
  }
  if (cinit) builder_.CreateStore(cinit, slot);
  else builder_.CreateStore((current_btype_=="float"? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) : (ir::Value*)module_.GetContext().GetConstInt(0)), slot);
  // record simple integer consts for dimension evaluation
  try {
    if (auto* ci = dynamic_cast<ir::ConstantInt*>(cinit)) {
      const_values_[ctx->Ident()->getText()] = ci->GetValue();
    }
  } catch(...) {}
  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", "暂不支持的语句或声明"));
}

// 变量声明的 IR 生成目前也是最小实现：
// - 先检查声明的基础类型，当前仅支持局部 int；
// - 再把 Decl 中的变量定义交给 visitVarDef 继续处理。
//
// 和更完整的版本相比，这里还没有：
// - 一个 Decl 中多个变量定义的顺序处理；
// - const、数组、全局变量等不同声明形态；
// - 更丰富的类型系统。
std::any IRGenImpl::visitDecl(SysYParser::DeclContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少变量声明"));
  }
  if (ctx->varDecl()) {
    auto* vdecl = ctx->varDecl();
    if (!vdecl->bType() || vdecl->bType()->getText() != "int") {
      throw std::runtime_error(FormatError("irgen", "当前仅支持局部 int 变量声明"));
    }
    for (auto* var_def : vdecl->varDef()) {
      if (var_def) var_def->accept(this);
    }
    return {};
  }
  if (ctx->constDecl()) {
    auto* cdecl = ctx->constDecl();
    if (!cdecl->bType()) throw std::runtime_error(FormatError("irgen", "缺少常量基类型"));
    current_btype_ = cdecl->bType()->getText();
    if (current_btype_ != "int" && current_btype_ != "float") throw std::runtime_error(FormatError("irgen", "当前仅支持局部 int/float 常量声明"));
    for (auto* const_def : cdecl->constDef()) if (const_def) const_def->accept(this);
    current_btype_.clear();
    return {};
  }
  throw std::runtime_error(FormatError("irgen", "暂不支持的声明类型"));
}


// 当前仍是教学用的最小版本，因此这里只支持：
// - 局部 int 变量；
// - 标量初始化；
// - 一个 VarDef 对应一个槽位。
std::any IRGenImpl::visitVarDef(SysYParser::VarDefContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少变量定义"));
  }
  if (!ctx->Ident()) {
    throw std::runtime_error(FormatError("irgen", "变量声明缺少名称"));
  }
  if (storage_map_.find(ctx) != storage_map_.end()) {
    throw std::runtime_error(FormatError("irgen", "声明重复生成存储槽位"));
  }
  // check if this is an array declaration (has constExp dimensions)
  if (!ctx->constExp().empty()) {
    // parse dims
    std::vector<int> dims;
    for (auto* ce : ctx->constExp()) {
      try {
        int v = 0;
        auto anyv = sema_.GetConstVal(ce);
        if (anyv.has_value()) {
          if (anyv.type() == typeid(int)) v = std::any_cast<int>(anyv);
          else if (anyv.type() == typeid(long)) v = (int)std::any_cast<long>(anyv);
          else throw std::runtime_error("not-const-int");
        } else {
          // try simple patterns like NUM or IDENT+NUM or NUM+IDENT
          std::string s = ce->addExp()->getText();
          s.erase(std::remove_if(s.begin(), s.end(), ::isspace), s.end());
          auto pos = s.find('+');
          if (pos == std::string::npos) {
            // plain number or identifier
            try { v = std::stoi(s); }
            catch(...) {
              // try lookup identifier in recorded consts or symbol table
              auto it = const_values_.find(s);
              if (it != const_values_.end()) v = (int)it->second;
              else {
                VarInfo vi; void* declctx = nullptr;
                if (sema_.GetSymbolTable().LookupVar(s, vi, declctx) && vi.const_val.has_value()) {
                  if (vi.const_val.type() == typeid(int)) v = std::any_cast<int>(vi.const_val);
                  else if (vi.const_val.type() == typeid(long)) v = (int)std::any_cast<long>(vi.const_val);
                  else throw std::runtime_error("not-const-int");
                } else throw std::runtime_error("not-const-int");
              }
            }
          } else {
            // form A+B where A or B may be ident or number
            std::string L = s.substr(0, pos);
            std::string R = s.substr(pos + 1);
            int lv = 0, rv = 0; bool ok = false;
            // try left
            try { lv = std::stoi(L); ok = true; } catch(...) {
              auto it = const_values_.find(L);
              if (it != const_values_.end()) { lv = (int)it->second; ok = true; }
              else {
                VarInfo vi; void* declctx = nullptr;
                if (sema_.GetSymbolTable().LookupVar(L, vi, declctx) && vi.const_val.has_value()) {
                  if (vi.const_val.type() == typeid(int)) lv = std::any_cast<int>(vi.const_val);
                  else if (vi.const_val.type() == typeid(long)) lv = (int)std::any_cast<long>(vi.const_val);
                  ok = true;
                }
              }
            }
            // try right
            try { rv = std::stoi(R); ok = ok && true; } catch(...) {
              auto it2 = const_values_.find(R);
              if (it2 != const_values_.end()) { rv = (int)it2->second; ok = ok && true; }
              else {
                VarInfo vi2; void* declctx2 = nullptr;
                if (sema_.GetSymbolTable().LookupVar(R, vi2, declctx2) && vi2.const_val.has_value()) {
                  if (vi2.const_val.type() == typeid(int)) rv = std::any_cast<int>(vi2.const_val);
                  else if (vi2.const_val.type() == typeid(long)) rv = (int)std::any_cast<long>(vi2.const_val);
                  ok = ok && true;
                } else ok = false;
              }
            }
            if (!ok) throw std::runtime_error("not-const-int");
            v = lv + rv;
          }
        }
        dims.push_back(v);
      } catch (...) {
        throw std::runtime_error(FormatError("irgen", "数组维度必须为常量整数"));
      }
    }

    std::shared_ptr<ir::Type> elemTy = (current_btype_ == "float") ? ir::Type::GetFloat32Type() : ir::Type::GetInt32Type();

    // build nested array type
    std::function<std::shared_ptr<ir::Type>(size_t)> makeArrayType = [&](size_t level) -> std::shared_ptr<ir::Type> {
      if (level + 1 >= dims.size()) return ir::Type::GetArrayType(elemTy, dims[level]);
      auto sub = makeArrayType(level + 1);
      return ir::Type::GetArrayType(sub, dims[level]);
    };
    auto fullArrayTy = makeArrayType(0);
    auto arr_ptr_ty = ir::Type::GetPointerType(fullArrayTy);
    auto* array_slot = builder_.CreateAlloca(arr_ptr_ty, module_.GetContext().NextTemp());
    storage_map_[ctx] = array_slot;
    name_map_[ctx->Ident()->getText()] = array_slot;

    // compute spans and total scalar slots
    int nlevels = (int)dims.size();
    std::vector<int> span(nlevels);
    int total = 1;
    for (int i = nlevels - 1; i >= 0; --i) {
      if (i == nlevels - 1) span[i] = 1;
      else span[i] = span[i + 1] * dims[i + 1];
      total *= dims[i];
    }

    ir::Value* zero = elemTy->IsFloat32() ? (ir::Value*)module_.GetContext().GetConstFloat(0.0f) : (ir::Value*)module_.GetContext().GetConstInt(0);
    std::vector<ir::Value*> slots(total, zero);

    // process initializer (if any) into linear slots
    if (auto* init_value = ctx->initVal()) {
      std::function<void(SysYParser::InitValContext*, int, int&)> process_group;
      process_group = [&](SysYParser::InitValContext* init, int level, int& pos) {
        if (level >= nlevels) return;
        int sub_span = span[level];
        int elems = dims[level];
        if (!init) { pos += elems * sub_span; return; }
        for (auto* child : init->initVal()) {
          if (pos >= total) break;
          if (!child) { pos += 1; continue; }
          if (!child->initVal().empty()) {
            int subpos = pos;
            int inner = subpos;
            process_group(child, level + 1, inner);
            pos = subpos + sub_span;
          } else if (child->exp()) {
            try { ir::Value* v = EvalExpr(*child->exp()); if (pos < total) slots[pos] = v; } catch(...) {}
            pos += 1;
          } else { pos += 1; }
        }
      };
      int pos0 = 0;
      process_group(init_value, 0, pos0);
    }

    // emit stores for each scalar slot in row-major order
    for (int idx = 0; idx < total; ++idx) {
      std::vector<int> indices;
      int rem = idx;
      for (int L = 0; L < nlevels; ++L) {
        int ind = rem / span[L];
        indices.push_back(ind % dims[L]);
        rem = rem % span[L];
      }
      std::vector<ir::Value*> gep_inds;
      gep_inds.push_back(module_.GetContext().GetConstInt(0));
      for (int v : indices) gep_inds.push_back(module_.GetContext().GetConstInt(v));
      while (gep_inds.size() < (size_t)(1 + nlevels)) gep_inds.push_back(module_.GetContext().GetConstInt(0));
      auto* gep = builder_.CreateGEP(array_slot, gep_inds, module_.GetContext().NextTemp());
      builder_.CreateStore(slots[idx], gep);
    }

    return {};
  }

  // scalar variable
  auto* slot = builder_.CreateAllocaI32(module_.GetContext().NextTemp());
  storage_map_[ctx] = slot;
  ir::Value* init = nullptr;
  if (auto* init_value = ctx->initVal()) {
    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 {};
}
