#pragma once

#include <any>
#include <memory>
#include <optional>
#include <string>
#include <vector>

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

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

  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override;

 private:
  enum class FlowState {
    Continue,
    Terminated,
  };

  struct TypedValue {
    ir::Value* value = nullptr;
    SemanticType type = SemanticType::Int;
    bool is_array = false;
    std::vector<int> dims;
  };

  struct LValueInfo {
    SymbolEntry* symbol = nullptr;
    ir::Value* addr = nullptr;
    SemanticType type = SemanticType::Int;
    bool is_array = false;
    std::vector<int> dims;
    bool root_param_array_no_index = false;
  };

  struct LoopContext {
    ir::BasicBlock* cond_block = nullptr;
    ir::BasicBlock* exit_block = nullptr;
  };

  struct InitExprSlot {
    size_t index = 0;
    SysYParser::ExpContext* expr = nullptr;
  };

  [[noreturn]] void ThrowError(const antlr4::ParserRuleContext* ctx,
                               const std::string& message) const;

  void RegisterBuiltinFunctions();
  void PredeclareTopLevel(SysYParser::CompUnitContext& ctx);
  void PredeclareFunction(SysYParser::FuncDefContext& ctx);
  void PredeclareGlobalDecl(SysYParser::DeclContext& ctx);

  void EmitGlobalDecl(SysYParser::DeclContext& ctx);
  void EmitFunction(SysYParser::FuncDefContext& ctx);
  void BindFunctionParams(SysYParser::FuncDefContext& ctx, ir::Function& func);

  void EmitBlock(SysYParser::BlockContext& ctx, bool create_scope = true);
  FlowState EmitBlockItem(SysYParser::BlockItemContext& ctx);
  FlowState EmitStmt(SysYParser::StmtContext& ctx);

  void EmitDecl(SysYParser::DeclContext& ctx, bool is_global);
  void EmitVarDecl(SysYParser::VarDeclContext* ctx, bool is_global, bool is_const);
  void EmitConstDecl(SysYParser::ConstDeclContext* ctx, bool is_global);
  void EmitGlobalVarDef(SysYParser::VarDefContext& ctx, SemanticType type);
  void EmitGlobalConstDef(SysYParser::ConstDefContext& ctx, SemanticType type);
  void EmitLocalVarDef(SysYParser::VarDefContext& ctx, SemanticType type, bool is_const);
  void EmitLocalConstDef(SysYParser::ConstDefContext& ctx, SemanticType type);

  std::string ExpectIdent(const antlr4::ParserRuleContext& ctx,
                          antlr4::tree::TerminalNode* ident) const;
  SemanticType ParseBType(SysYParser::BTypeContext* ctx) const;
  SemanticType ParseFuncType(SysYParser::FuncTypeContext* ctx) const;
  std::shared_ptr<ir::Type> GetIRScalarType(SemanticType type) const;
  std::shared_ptr<ir::Type> BuildArrayType(SemanticType base_type,
                                           const std::vector<int>& dims) const;
  std::vector<int> ParseArrayDims(const std::vector<SysYParser::ConstExpContext*>& dims_ctx);
  std::vector<int> ParseParamDims(SysYParser::FuncFParamContext& ctx);

  FunctionTypeInfo BuildFunctionTypeInfo(SysYParser::FuncDefContext& ctx);
  std::vector<std::shared_ptr<ir::Type>> BuildFunctionIRParamTypes(
      const FunctionTypeInfo& function_type) const;
  std::vector<std::string> BuildFunctionIRParamNames(SysYParser::FuncDefContext& ctx) const;

  TypedValue EmitExp(SysYParser::ExpContext& ctx);
  TypedValue EmitAddExp(SysYParser::AddExpContext& ctx);
  TypedValue EmitMulExp(SysYParser::MulExpContext& ctx);
  TypedValue EmitUnaryExp(SysYParser::UnaryExpContext& ctx);
  TypedValue EmitPrimaryExp(SysYParser::PrimaryExpContext& ctx);
  TypedValue EmitRelExp(SysYParser::RelExpContext& ctx);
  TypedValue EmitEqExp(SysYParser::EqExpContext& ctx);
  TypedValue EmitLValValue(SysYParser::LValContext& ctx);
  LValueInfo ResolveLVal(SysYParser::LValContext& ctx);
  ir::Value* GenLValAddr(SysYParser::LValContext& ctx);

  void EmitCond(SysYParser::CondContext& ctx, ir::BasicBlock* true_block,
                ir::BasicBlock* false_block);
  void EmitLOrCond(SysYParser::LOrExpContext& ctx, ir::BasicBlock* true_block,
                   ir::BasicBlock* false_block);
  void EmitLAndCond(SysYParser::LAndExpContext& ctx, ir::BasicBlock* true_block,
                    ir::BasicBlock* false_block);

  TypedValue CastScalar(TypedValue value, SemanticType target_type,
                        const antlr4::ParserRuleContext* ctx);
  ir::Value* CastToCondition(TypedValue value,
                             const antlr4::ParserRuleContext* ctx);
  TypedValue NormalizeLogicalValue(TypedValue value,
                                   const antlr4::ParserRuleContext* ctx);
  bool IsNumeric(const TypedValue& value) const;
  bool IsSameDims(const std::vector<int>& lhs, const std::vector<int>& rhs) const;

  ConstantValue ParseNumber(SysYParser::NumberContext& ctx) const;
  ConstantValue EvalConstExp(SysYParser::ExpContext& ctx);
  ConstantValue EvalConstAddExp(SysYParser::AddExpContext& ctx);
  ConstantValue EvalConstMulExp(SysYParser::MulExpContext& ctx);
  ConstantValue EvalConstUnaryExp(SysYParser::UnaryExpContext& ctx);
  ConstantValue EvalConstPrimaryExp(SysYParser::PrimaryExpContext& ctx);
  ConstantValue EvalConstLVal(SysYParser::LValContext& ctx);
  ConstantValue ConvertConst(ConstantValue value, SemanticType target_type) const;

  std::vector<ConstantValue> FlattenConstInitVal(SysYParser::ConstInitValContext* ctx,
                                                 SemanticType base_type,
                                                 const std::vector<int>& dims);
  std::vector<ConstantValue> FlattenInitVal(SysYParser::InitValContext* ctx,
                                            SemanticType base_type,
                                            const std::vector<int>& dims);
  std::vector<InitExprSlot> FlattenLocalInitVal(SysYParser::InitValContext* ctx,
                                                const std::vector<int>& dims);

  void FlattenConstInitValImpl(SysYParser::ConstInitValContext* ctx,
                               SemanticType base_type,
                               const std::vector<int>& dims, size_t depth,
                               size_t object_begin, size_t object_end,
                               size_t& cursor, std::vector<ConstantValue>& out);
  void FlattenInitValImpl(SysYParser::InitValContext* ctx, SemanticType base_type,
                          const std::vector<int>& dims, size_t depth,
                          size_t object_begin, size_t object_end,
                          size_t& cursor, std::vector<ConstantValue>& out);
  void FlattenLocalInitValImpl(SysYParser::InitValContext* ctx,
                               const std::vector<int>& dims, size_t depth,
                               size_t object_begin, size_t object_end,
                               size_t& cursor, std::vector<InitExprSlot>& out);

  size_t CountArrayElements(const std::vector<int>& dims, size_t start = 0) const;
  size_t AlignInitializerCursor(const std::vector<int>& dims, size_t depth,
                                size_t object_begin, size_t cursor) const;
  size_t FlattenIndices(const std::vector<int>& dims,
                        const std::vector<int>& indices) const;

  ConstantValue ZeroConst(SemanticType type) const;
  ir::Value* ZeroIRValue(SemanticType type);
  ir::Value* CreateTypedConstant(const ConstantValue& value);
  ir::AllocaInst* CreateEntryAlloca(std::shared_ptr<ir::Type> allocated_type,
                                    const std::string& name);
  void ZeroInitializeLocalArray(ir::Value* addr, SemanticType base_type,
                                const std::vector<int>& dims);
  void StoreLocalArrayElements(ir::Value* addr, SemanticType base_type,
                               const std::vector<int>& dims,
                               const std::vector<InitExprSlot>& init_slots);
  ir::Value* CreateArrayElementAddr(ir::Value* base_addr, bool is_param_array,
                                    SemanticType base_type,
                                    const std::vector<int>& dims,
                                    const std::vector<ir::Value*>& indices,
                                    const antlr4::ParserRuleContext* ctx);

  std::string NextTemp();
  std::string NextBlockName(const std::string& prefix);

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::IRBuilder builder_;
  SymbolTable symbols_;
  ir::Function* current_function_ = nullptr;
  SemanticType current_return_type_ = SemanticType::Void;
  std::vector<LoopContext> loop_stack_;
  bool builtins_registered_ = false;
};

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