#include "irgen/IRGen.h"

#include <stdexcept>
#include <vector>

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

namespace {

std::shared_ptr<ir::Type> BaseTypeFromDecl(SysYParser::BtypeContext* btype) {
  return (btype && btype->FLOAT()) ? ir::Type::GetFloatType() : ir::Type::GetInt32Type();
}

std::shared_ptr<ir::Type> StorageType(std::shared_ptr<ir::Type> ty) {
  if (ty->IsInt32()) return ir::Type::GetPtrInt32Type();
  if (ty->IsFloat()) return ir::Type::GetPtrFloatType();
  return ty;
}

size_t CountScalars(const std::shared_ptr<ir::Type>& ty) {
  if (!ty->IsArray()) return 1;
  auto arr_ty = ty->GetAsArrayType();
  return arr_ty->GetNumElements() * CountScalars(arr_ty->GetElementType());
}

}  // namespace

void IRGenImpl::ZeroInitializeLocal(ir::Value* ptr, std::shared_ptr<ir::Type> ty) {
  if (ty->IsArray()) {
    auto arr_ty = ty->GetAsArrayType();
    for (uint32_t i = 0; i < arr_ty->GetNumElements(); ++i) {
      auto* elem_ptr = builder_.CreateGEP(StorageType(arr_ty->GetElementType()), ptr,
                                          {builder_.CreateConstInt(0),
                                           builder_.CreateConstInt(static_cast<int>(i))},
                                          module_.GetContext().NextTemp());
      ZeroInitializeLocal(elem_ptr, arr_ty->GetElementType());
    }
    return;
  }

  ir::Value* zero = ty->IsFloat() ? static_cast<ir::Value*>(builder_.CreateConstFloat(0.0f))
                                  : static_cast<ir::Value*>(builder_.CreateConstInt(0));
  builder_.CreateStore(zero, ptr);
}

void IRGenImpl::EmitLocalInitValue(ir::Value* ptr, std::shared_ptr<ir::Type> ty,
                                   SysYParser::InitValueContext* init) {
  if (!init) return;

  auto build_flat_scalar_ptr = [&](ir::Value* base_ptr,
                                   const std::shared_ptr<ir::Type>& base_ty,
                                   size_t flat_index) -> ir::Value* {
    if (!base_ty->IsArray()) return base_ptr;

    std::vector<ir::Value*> indices;
    indices.push_back(builder_.CreateConstInt(0));
    auto cur_ty = base_ty;
    size_t offset = flat_index;
    while (cur_ty->IsArray()) {
      auto arr_ty = cur_ty->GetAsArrayType();
      const size_t step = CountScalars(arr_ty->GetElementType());
      const size_t idx = step == 0 ? 0 : offset / step;
      offset = step == 0 ? 0 : offset % step;
      indices.push_back(builder_.CreateConstInt(static_cast<int>(idx)));
      cur_ty = arr_ty->GetElementType();
    }
    return builder_.CreateGEP(StorageType(cur_ty), base_ptr, indices,
                              module_.GetContext().NextTemp());
  };

  if (ty->IsArray()) {
    auto arr_ty = ty->GetAsArrayType();
    const auto elem_ty = arr_ty->GetElementType();
    const size_t elem_step = CountScalars(elem_ty);

    if (init->exp()) {
      auto* elem_ptr = build_flat_scalar_ptr(ptr, ty, 0);
      EmitLocalInitValue(elem_ptr, elem_ty, init);
      return;
    }

    const auto& children = init->initValue();
    size_t scalar_cursor = 0;
    for (auto* child : children) {
      if (scalar_cursor >= CountScalars(ty)) break;
      if (child->exp()) {
        auto* elem_ptr = build_flat_scalar_ptr(ptr, ty, scalar_cursor);
        auto scalar_ty = elem_ty;
        while (scalar_ty->IsArray()) {
          scalar_ty = scalar_ty->GetAsArrayType()->GetElementType();
        }
        EmitLocalInitValue(elem_ptr, scalar_ty, child);
        ++scalar_cursor;
        continue;
      }

      const size_t elem_index = elem_step == 0 ? 0 : scalar_cursor / elem_step;
      if (elem_index >= arr_ty->GetNumElements()) break;
      auto* elem_ptr = builder_.CreateGEP(StorageType(elem_ty), ptr,
                                          {builder_.CreateConstInt(0),
                                           builder_.CreateConstInt(static_cast<int>(elem_index))},
                                          module_.GetContext().NextTemp());
      EmitLocalInitValue(elem_ptr, elem_ty, child);
      scalar_cursor = (elem_index + 1) * elem_step;
    }
    return;
  }

  if (!init->exp()) {
    return;
  }
  ir::Value* value = EvalExpr(*init->exp());
  if (ty->IsFloat() && value->GetType()->IsInt32()) {
    value = builder_.CreateSIToFP(value, ty, module_.GetContext().NextTemp());
  } else if (ty->IsInt32() && value->GetType()->IsFloat()) {
    value = builder_.CreateFPToSI(value, ty, module_.GetContext().NextTemp());
  }
  builder_.CreateStore(value, ptr);
}

std::any IRGenImpl::visitDecl(SysYParser::DeclContext* ctx) {
  if (ctx->constDecl()) return ctx->constDecl()->accept(this);
  if (ctx->varDecl()) return ctx->varDecl()->accept(this);
  return {};
}

std::any IRGenImpl::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  for (auto* def : ctx->constDef()) {
    def->accept(this);
  }
  return {};
}

std::any IRGenImpl::visitVarDecl(SysYParser::VarDeclContext* ctx) {
  for (auto* def : ctx->varDef()) {
    def->accept(this);
  }
  return {};
}

std::any IRGenImpl::visitConstDef(SysYParser::ConstDefContext* ctx) {
  const std::string name = ctx->ID()->getText();
  auto ty = BaseTypeFromDecl(
      dynamic_cast<SysYParser::ConstDeclContext*>(ctx->parent)->btype());
  const auto dims = ctx->exp();
  for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
    auto* dim = EvalConstExpr(**it);
    if (auto* ci = dynamic_cast<ir::ConstantInt*>(dim)) {
      ty = ir::ArrayType::Get(ty, ci->GetValue());
      continue;
    }
    throw std::runtime_error(FormatError("irgen", "数组维度必须是整型常量"));
  }
  ir::Value* slot = nullptr;

  if (is_global_scope_) {
    ir::ConstantValue* init = nullptr;
    if (ctx->initValue() && ctx->initValue()->exp()) {
      init = EvalConstExpr(*ctx->initValue()->exp());
      if (ty->IsInt32() && init->GetType()->IsFloat()) {
        init = module_.GetContext().GetConstInt(
            static_cast<int>(static_cast<ir::ConstantFloat*>(init)->GetValue()));
      } else if (ty->IsFloat() && init->GetType()->IsInt32()) {
        init = module_.GetContext().GetConstFloat(
            static_cast<float>(static_cast<ir::ConstantInt*>(init)->GetValue()));
      }
    }
    slot = module_.CreateGlobalValue(name, StorageType(ty), init);
  } else {
    slot = builder_.CreateAlloca(StorageType(ty), name);
    ZeroInitializeLocal(slot, ty);
    EmitLocalInitValue(slot, ty, ctx->initValue());
  }

  storage_map_[ctx] = slot;
  return {};
}

std::any IRGenImpl::visitVarDef(SysYParser::VarDefContext* ctx) {
  const std::string name = ctx->ID()->getText();
  auto ty = BaseTypeFromDecl(
      dynamic_cast<SysYParser::VarDeclContext*>(ctx->parent)->btype());
  const auto dims = ctx->exp();
  for (auto it = dims.rbegin(); it != dims.rend(); ++it) {
    auto* dim = EvalConstExpr(**it);
    if (auto* ci = dynamic_cast<ir::ConstantInt*>(dim)) {
      ty = ir::ArrayType::Get(ty, ci->GetValue());
      continue;
    }
    throw std::runtime_error(FormatError("irgen", "数组维度必须是整型常量"));
  }
  ir::Value* slot = nullptr;

  if (is_global_scope_) {
    ir::ConstantValue* init = nullptr;
    if (ctx->initValue() && ctx->initValue()->exp()) {
      init = EvalConstExpr(*ctx->initValue()->exp());
      if (ty->IsInt32() && init->GetType()->IsFloat()) {
        init = module_.GetContext().GetConstInt(
            static_cast<int>(static_cast<ir::ConstantFloat*>(init)->GetValue()));
      } else if (ty->IsFloat() && init->GetType()->IsInt32()) {
        init = module_.GetContext().GetConstFloat(
            static_cast<float>(static_cast<ir::ConstantInt*>(init)->GetValue()));
      }
    }
    slot = module_.CreateGlobalValue(name, StorageType(ty), init);
  } else {
    slot = builder_.CreateAlloca(StorageType(ty), name);
    if (ctx->initValue()) {
      ZeroInitializeLocal(slot, ty);
      EmitLocalInitValue(slot, ty, ctx->initValue());
    }
  }

  storage_map_[ctx] = slot;
  return {};
}