#pragma once

#include <memory>
#include <string>
#include <unordered_map>
#include <vector>

#include "SysYParser.h"
#include "ir/IR.h"
#include "sem/Sema.h"

class IRGenImpl {
 public:
  IRGenImpl(ir::Module& module, const SemanticContext& sema);

  void Gen(SysYParser::CompUnitContext& cu);

 private:
  struct StorageEntry {
    ir::Value* storage = nullptr;
    std::shared_ptr<ir::Type> declared_type;
    bool is_array_param = false;
    bool is_global = false;
    bool is_const = false;
  };

  void DeclareBuiltins();
  void GenGlobals(SysYParser::CompUnitContext& cu);
  void GenFunctionDecls(SysYParser::CompUnitContext& cu);
  void GenFunctionBodies(SysYParser::CompUnitContext& cu);

  void GenFuncDef(SysYParser::FuncDefContext& func);
  void GenBlock(SysYParser::BlockContext& block);
  void GenBlockItem(SysYParser::BlockItemContext& item);
  void GenDecl(SysYParser::DeclContext& decl);
  void GenConstDecl(SysYParser::ConstDeclContext& decl);
  void GenVarDecl(SysYParser::VarDeclContext& decl);
  void GenStmt(SysYParser::StmtContext& stmt);

  ir::Value* GenExpr(SysYParser::ExpContext& expr);
  ir::Value* GenAddExpr(SysYParser::AddExpContext& add);
  ir::Value* GenMulExpr(SysYParser::MulExpContext& mul);
  ir::Value* GenUnaryExpr(SysYParser::UnaryExpContext& unary);
  ir::Value* GenPrimary(SysYParser::PrimaryContext& primary);
  ir::Value* GenRelExpr(SysYParser::RelExpContext& rel);
  ir::Value* GenEqExpr(SysYParser::EqExpContext& eq);

  ir::Value* GenLValueAddress(SysYParser::LValContext& lval);
  ir::Value* GenLValueValue(SysYParser::LValContext& lval);

  void GenCond(SysYParser::CondContext& cond, ir::BasicBlock* true_block,
               ir::BasicBlock* false_block);
  void GenLOrCond(SysYParser::LOrExpContext& expr, ir::BasicBlock* true_block,
                  ir::BasicBlock* false_block);
  void GenLAndCond(SysYParser::LAndExpContext& expr, ir::BasicBlock* true_block,
                   ir::BasicBlock* false_block);

  ir::Value* CastValue(ir::Value* value, const std::shared_ptr<ir::Type>& dst_type);
  ir::Value* ToBool(ir::Value* value);
  ir::Value* DecayArrayPointer(ir::Value* array_ptr);

  void EnterScope();
  void ExitScope();
  void EnsureInsertableBlock();
  void DeclareLocal(const std::string& name, StorageEntry entry);
  StorageEntry* LookupStorage(const std::string& name);
  const StorageEntry* LookupStorage(const std::string& name) const;

  size_t CountScalars(const std::shared_ptr<ir::Type>& type) const;
  std::vector<int> FlatIndexToIndices(const std::shared_ptr<ir::Type>& type,
                                      size_t flat_index) const;
  void EmitArrayStore(ir::Value* base_ptr, const std::shared_ptr<ir::Type>& array_type,
                      size_t flat_index, ir::Value* value);
  void ZeroInitializeLocalArray(ir::Value* base_ptr,
                                const std::shared_ptr<ir::Type>& array_type);
  void EmitLocalArrayInit(ir::Value* base_ptr, const std::shared_ptr<ir::Type>& array_type,
                          SysYParser::InitValContext& init);
  void EmitLocalConstArrayInit(ir::Value* base_ptr,
                               const std::shared_ptr<ir::Type>& array_type,
                               SysYParser::ConstInitValContext& init);

  ir::ConstantValue* BuildGlobalInitializer(const std::shared_ptr<ir::Type>& type,
                                            SysYParser::InitValContext* init);
  ir::ConstantValue* BuildGlobalConstInitializer(
      const std::shared_ptr<ir::Type>& type, SysYParser::ConstInitValContext* init);

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::Function* current_function_ = nullptr;
  std::shared_ptr<ir::Type> current_return_type_;
  ir::IRBuilder builder_;
  std::vector<std::unordered_map<std::string, StorageEntry>> local_scopes_;
  std::unordered_map<std::string, StorageEntry> globals_;
  std::vector<ir::BasicBlock*> break_targets_;
  std::vector<ir::BasicBlock*> continue_targets_;
  std::unordered_map<std::string, ConstantData> global_const_values_;
};

std::unique_ptr<ir::Module> GenerateIR(SysYParser::CompUnitContext& tree,
                                       const SemanticContext& sema);