#include "irgen/IRGen.h"

#include <stdexcept>

#include "utils/Log.h"

namespace {

using ir::Type;

std::shared_ptr<Type> BuiltinFn(std::shared_ptr<Type> ret,
                                std::vector<std::shared_ptr<Type>> params) {
  return Type::GetFunctionType(std::move(ret), std::move(params));
}

bool SameType(const std::shared_ptr<Type>& lhs, const std::shared_ptr<Type>& rhs) {
  return lhs && rhs && lhs->Equals(*rhs);
}

}  // namespace

IRGenImpl::IRGenImpl(ir::Module& module, const SemanticContext& sema)
    : module_(module),
      sema_(sema),
      current_return_type_(Type::GetVoidType()),
      builder_(module.GetContext(), nullptr) {}

void IRGenImpl::Gen(SysYParser::CompUnitContext& cu) {
  DeclareBuiltins();
  GenGlobals(cu);
  GenFunctionDecls(cu);
  GenFunctionBodies(cu);
}

void IRGenImpl::DeclareBuiltins() {
  const auto i32 = Type::GetInt32Type();
  const auto f32 = Type::GetFloatType();
  const auto void_ty = Type::GetVoidType();

  const struct {
    const char* name;
    std::shared_ptr<Type> type;
  } builtins[] = {
      {"getint", BuiltinFn(i32, {})},
      {"getch", BuiltinFn(i32, {})},
      {"getfloat", BuiltinFn(f32, {})},
      {"getarray", BuiltinFn(i32, {Type::GetPointerType(i32)})},
      {"getfarray", BuiltinFn(i32, {Type::GetPointerType(f32)})},
      {"putint", BuiltinFn(void_ty, {i32})},
      {"putch", BuiltinFn(void_ty, {i32})},
      {"putfloat", BuiltinFn(void_ty, {f32})},
      {"putarray", BuiltinFn(void_ty, {i32, Type::GetPointerType(i32)})},
      {"putfarray", BuiltinFn(void_ty, {i32, Type::GetPointerType(f32)})},
      {"starttime", BuiltinFn(void_ty, {})},
      {"stoptime", BuiltinFn(void_ty, {})},
  };

  for (const auto& builtin : builtins) {
    if (!module_.FindFunction(builtin.name)) {
      module_.CreateFunction(builtin.name, builtin.type, true);
    }
  }
}

void IRGenImpl::GenFunctionDecls(SysYParser::CompUnitContext& cu) {
  for (auto* func : cu.funcDef()) {
    if (!func || !func->Ident()) {
      continue;
    }
    auto* symbol = sema_.ResolveFuncDef(func);
    if (!symbol) {
      throw std::runtime_error(FormatError("irgen", "缺少函数语义信息"));
    }
    auto* ir_func = module_.FindFunction(symbol->name);
    if (ir_func) {
      continue;
    }
    ir_func = module_.CreateFunction(symbol->name, symbol->type, false);
    const auto& params = symbol->type->GetParamTypes();
    for (size_t i = 0; i < params.size(); ++i) {
      ir_func->AddArgument(params[i], "%arg" + std::to_string(i));
    }
  }
}

void IRGenImpl::GenFunctionBodies(SysYParser::CompUnitContext& cu) {
  for (auto* func : cu.funcDef()) {
    if (func) {
      GenFuncDef(*func);
    }
  }
}

void IRGenImpl::GenFuncDef(SysYParser::FuncDefContext& func) {
  auto* symbol = sema_.ResolveFuncDef(&func);
  if (!symbol) {
    throw std::runtime_error(FormatError("irgen", "函数缺少语义绑定"));
  }

  current_function_ = module_.FindFunction(symbol->name);
  if (!current_function_) {
    throw std::runtime_error(FormatError("irgen", "函数声明缺失: " + symbol->name));
  }
  current_return_type_ = symbol->type->GetReturnType();
  auto* entry = current_function_->CreateBlock("entry");
  auto* body = current_function_->CreateBlock("entry.body");
  builder_.SetInsertPoint(body);
  local_scopes_.clear();
  break_targets_.clear();
  continue_targets_.clear();

  EnterScope();
  if (auto* params = func.funcFParams()) {
    const auto& args = current_function_->GetArguments();
    for (size_t i = 0; i < params->funcFParam().size(); ++i) {
      auto* param = params->funcFParam(i);
      const auto* arg = args.at(i).get();
      const std::string name = param->Ident()->getText();
      auto* slot = builder_.CreateAlloca(arg->GetType(), module_.GetContext().NextTemp());
      builder_.CreateStore(const_cast<ir::Argument*>(arg), slot);
      DeclareLocal(name, {slot, arg->GetType(), !param->L_BRACK().empty(), false, false});
    }
  }

  GenBlock(*func.block());
  ExitScope();

  ir::IRBuilder entry_builder(module_.GetContext(), entry);
  entry_builder.CreateBr(body);

  if (builder_.GetInsertBlock() && !builder_.GetInsertBlock()->HasTerminator()) {
    if (current_return_type_->IsVoid()) {
      builder_.CreateRetVoid();
    } else if (current_return_type_->IsFloat32()) {
      builder_.CreateRet(builder_.CreateConstFloat(0.0f));
    } else {
      builder_.CreateRet(builder_.CreateConstInt(0));
    }
  }
}

void IRGenImpl::EnterScope() { local_scopes_.emplace_back(); }

void IRGenImpl::ExitScope() {
  if (!local_scopes_.empty()) {
    local_scopes_.pop_back();
  }
}

void IRGenImpl::EnsureInsertableBlock() {
  if (!builder_.GetInsertBlock()) {
    auto* block = current_function_->CreateBlock(module_.GetContext().NextBlock("dead"));
    builder_.SetInsertPoint(block);
    return;
  }
  if (builder_.GetInsertBlock()->HasTerminator()) {
    auto* block = current_function_->CreateBlock(module_.GetContext().NextBlock("dead"));
    builder_.SetInsertPoint(block);
  }
}

void IRGenImpl::DeclareLocal(const std::string& name, StorageEntry entry) {
  if (local_scopes_.empty()) {
    EnterScope();
  }
  local_scopes_.back()[name] = std::move(entry);
}

IRGenImpl::StorageEntry* IRGenImpl::LookupStorage(const std::string& name) {
  for (auto it = local_scopes_.rbegin(); it != local_scopes_.rend(); ++it) {
    auto found = it->find(name);
    if (found != it->end()) {
      return &found->second;
    }
  }
  auto global = globals_.find(name);
  return global == globals_.end() ? nullptr : &global->second;
}

const IRGenImpl::StorageEntry* IRGenImpl::LookupStorage(const std::string& name) const {
  for (auto it = local_scopes_.rbegin(); it != local_scopes_.rend(); ++it) {
    auto found = it->find(name);
    if (found != it->end()) {
      return &found->second;
    }
  }
  auto global = globals_.find(name);
  return global == globals_.end() ? nullptr : &global->second;
}

size_t IRGenImpl::CountScalars(const std::shared_ptr<Type>& type) const {
  if (!type->IsArray()) {
    return 1;
  }
  return type->GetArraySize() * CountScalars(type->GetElementType());
}

std::vector<int> IRGenImpl::FlatIndexToIndices(const std::shared_ptr<Type>& type,
                                               size_t flat_index) const {
  if (!type->IsArray()) {
    return {};
  }
  size_t inner = CountScalars(type->GetElementType());
  int current = static_cast<int>(flat_index / inner);
  auto tail = FlatIndexToIndices(type->GetElementType(), flat_index % inner);
  tail.insert(tail.begin(), current);
  return tail;
}