#include "irgen/IRGen.h"

#include <stdexcept>

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

// 表达式生成当前也只实现了很小的一个子集。
// 目前支持：
// - 整数字面量
// - 普通局部变量读取
// - 括号表达式
// - 二元加法
//
// 还未支持：
// - 减乘除与一元运算
// - 赋值表达式
// - 函数调用
// - 数组、指针、下标访问
// - 条件与比较表达式
// - ...

ir::Value* IRGenImpl::EvalExpr(SysYParser::ExpContext& expr) {
  return std::any_cast<ir::Value*>(expr.accept(this));
}

ir::Value* IRGenImpl::EvalCond(SysYParser::CondContext& cond) {
  return std::any_cast<ir::Value*>(cond.accept(this));
}

// 基本表达式：数字、变量、括号表达式
std::any IRGenImpl::visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "缺少基本表达式"));
  }
  
  // 处理数字字面量
  if (ctx->DECIMAL_INT()) {
    int value = std::stoi(ctx->DECIMAL_INT()->getText());
    return static_cast<ir::Value*>(builder_.CreateConstInt(value));
  }
  
  if (ctx->HEX_INT()) {
    std::string hex = ctx->HEX_INT()->getText();
    int value = std::stoi(hex, nullptr, 16);
    return static_cast<ir::Value*>(builder_.CreateConstInt(value));
  }
  
  if (ctx->OCTAL_INT()) {
    std::string oct = ctx->OCTAL_INT()->getText();
    int value = std::stoi(oct, nullptr, 8);
    return static_cast<ir::Value*>(builder_.CreateConstInt(value));
  }
  
  if (ctx->ZERO()) {
    return static_cast<ir::Value*>(builder_.CreateConstInt(0));
  }
  
  // 处理变量
  if (ctx->lVal()) {
    return ctx->lVal()->accept(this);
  }
  
  // 处理括号表达式
  if (ctx->L_PAREN() && ctx->exp()) {
    return EvalExpr(*ctx->exp());
  }
  
  throw std::runtime_error(FormatError("irgen", "不支持的基本表达式类型"));
}

// 左值（变量）处理
// 1. 先通过语义分析结果把变量使用绑定回声明；
// 2. 再通过 storage_map_ 找到该声明对应的栈槽位；
// 3. 最后生成 load，把内存中的值读出来。
std::any IRGenImpl::visitLVal(SysYParser::LValContext* ctx) {
  if (!ctx || !ctx->Ident()) {
    throw std::runtime_error(FormatError("irgen", "非法左值"));
  }
  
  std::string varName = ctx->Ident()->getText();
  
  // 从语义分析获取变量定义
  auto* decl = sema_.ResolveVarUse(ctx);
  if (!decl) {
    auto* const_def = sema_.ResolveConstUse(ctx);
    if (!const_def) {
        FormatError("irgen",
                    "使用缺少语义绑定: " + varName);
    }
  }
  
  auto it = storage_map_.find(decl);
  if (it == storage_map_.end()) {
    throw std::runtime_error(
        FormatError("irgen",
                    "变量声明缺少存储槽位: " + varName));
  }
  
  return static_cast<ir::Value*>(
      builder_.CreateLoad(it->second, module_.GetContext().NextTemp()));
}


// 加法表达式
std::any IRGenImpl::visitAddExp(SysYParser::AddExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法加法表达式"));

  if (!ctx->addExp()) {
    return ctx->mulExp()->accept(this);
  }

  ir::Value* left = std::any_cast<ir::Value*>(ctx->addExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->mulExp()->accept(this));
  if (ctx->AddOp()) {
    return builder_.CreateAdd(left, right, module_.GetContext().NextTemp());
  } else if (ctx->SubOp()) {
    return builder_.CreateSub(left, right, module_.GetContext().NextTemp());
  }
  throw std::runtime_error(FormatError("irgen", "未知的加法操作符"));
}

// 在 IRGenExp.cpp 中添加

// visitMulExp
std::any IRGenImpl::visitMulExp(SysYParser::MulExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法乘法表达式"));

  // 如果是基本形式 UnaryExp
  if (!ctx->mulExp()) {
    return ctx->unaryExp()->accept(this);
  }

  // 否则有左子节点和操作符
  ir::Value* left = std::any_cast<ir::Value*>(ctx->mulExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->unaryExp()->accept(this));
  std::string op;
  if (ctx->MulOp()) op = "*";
  else if (ctx->DivOp()) op = "/";
  else if (ctx->QuoOp()) op = "%";
  else throw std::runtime_error(FormatError("irgen", "缺少乘法类操作符"));

  if (op == "*") {
    return builder_.CreateMul(left, right, module_.GetContext().NextTemp());
  } else if (op == "/") {
    return builder_.CreateDiv(left, right, module_.GetContext().NextTemp());
  } else if (op == "%") {
    return builder_.CreateMod(left, right, module_.GetContext().NextTemp());
  }
  throw std::runtime_error(FormatError("irgen", "未知的乘法操作符"));
}

// 关系表达式
std::any IRGenImpl::visitRelExp(SysYParser::RelExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法关系表达式"));

  if (!ctx->relExp()) {
    return ctx->addExp()->accept(this);
  }

  ir::Value* left = std::any_cast<ir::Value*>(ctx->relExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->addExp()->accept(this));

  if (ctx->LOp()) {
    return builder_.CreateICmpSLT(left, right, module_.GetContext().NextTemp());
  } else if (ctx->GOp()) {
    return builder_.CreateICmpSGT(left, right, module_.GetContext().NextTemp());
  } else if (ctx->LeOp()) {
    return builder_.CreateICmpSLE(left, right, module_.GetContext().NextTemp());
  } else if (ctx->GeOp()) {
    return builder_.CreateICmpSGE(left, right, module_.GetContext().NextTemp());
  }
  throw std::runtime_error(FormatError("irgen", "未知的关系操作符"));
}

// 相等表达式
std::any IRGenImpl::visitEqExp(SysYParser::EqExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法相等表达式"));

  if (!ctx->eqExp()) {
    return ctx->relExp()->accept(this);
  }

  ir::Value* left = std::any_cast<ir::Value*>(ctx->eqExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->relExp()->accept(this));

  if (ctx->EqOp()) {
    return builder_.CreateICmpEQ(left, right, module_.GetContext().NextTemp());
  } else if (ctx->NeOp()) {
    return builder_.CreateICmpNE(left, right, module_.GetContext().NextTemp());
  }
  throw std::runtime_error(FormatError("irgen", "未知的相等操作符"));
}

// 逻辑与
std::any IRGenImpl::visitLAndExp(SysYParser::LAndExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法逻辑与表达式"));

  if (!ctx->lAndExp()) {
    return ctx->eqExp()->accept(this);
  }

  ir::Value* left = std::any_cast<ir::Value*>(ctx->lAndExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->eqExp()->accept(this));
  auto zero = builder_.CreateConstInt(0);
  auto left_bool = builder_.CreateICmpNE(left, zero, module_.GetContext().NextTemp());
  auto right_bool = builder_.CreateICmpNE(right, zero, module_.GetContext().NextTemp());
  return builder_.CreateAnd(left_bool, right_bool, module_.GetContext().NextTemp());
}

// 逻辑或
std::any IRGenImpl::visitLOrExp(SysYParser::LOrExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法逻辑或表达式"));

  if (!ctx->lOrExp()) {
    return ctx->lAndExp()->accept(this);
  }

  ir::Value* left = std::any_cast<ir::Value*>(ctx->lOrExp()->accept(this));
  ir::Value* right = std::any_cast<ir::Value*>(ctx->lAndExp()->accept(this));
  auto zero = builder_.CreateConstInt(0);
  auto left_bool = builder_.CreateICmpNE(left, zero, module_.GetContext().NextTemp());
  auto right_bool = builder_.CreateICmpNE(right, zero, module_.GetContext().NextTemp());
  return builder_.CreateOr(left_bool, right_bool, module_.GetContext().NextTemp());
}

std::any IRGenImpl::visitExp(SysYParser::ExpContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法表达式"));
  return ctx->addExp()->accept(this);
}

std::any IRGenImpl::visitCond(SysYParser::CondContext* ctx) {
  if (!ctx) throw std::runtime_error(FormatError("irgen", "非法条件表达式"));
  return ctx->lOrExp()->accept(this);
}

ir::Value* IRGenImpl::visitCallExp(SysYParser::UnaryExpContext* ctx) {
  if (!ctx || !ctx->Ident()) {
    throw std::runtime_error(FormatError("irgen", "非法函数调用"));
  }

  std::string funcName = ctx->Ident()->getText();

  // 语义检查（如果需要）
  // auto* funcDecl = sema_.ResolveFuncCall(ctx);
  // if (!funcDecl) throw ...

  // 收集实参
  std::vector<ir::Value*> args;
  if (ctx->funcRParams()) {
    auto argList = ctx->funcRParams()->accept(this);
    args = std::any_cast<std::vector<ir::Value*>>(argList);
  }

  // 查找函数对象
  ir::Function* callee = module_.FindFunction(funcName);
  if (!callee) {
    throw std::runtime_error(FormatError("irgen", "未找到函数: " + funcName));
  }

  // 生成调用指令
  return builder_.CreateCall(callee, args, module_.GetContext().NextTemp());
}

// 实现一元表达式
std::any IRGenImpl::visitUnaryExp(SysYParser::UnaryExpContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "非法一元表达式"));
  }
  
  // 基本表达式
  if (ctx->primaryExp()) {
    return ctx->primaryExp()->accept(this);
  }
  
  // 函数调用
  if (ctx->Ident()) {
    return visitCallExp(ctx);
  }
  
  // 一元运算
  if (ctx->unaryOp() && ctx->unaryExp()) {
    auto* operand = std::any_cast<ir::Value*>(ctx->unaryExp()->accept(this));
    std::string op = ctx->unaryOp()->getText();
    
    if (op == "+") {
      // +x 等价于 x
      return operand;
    } else if (op == "-") {
      // -x 等价于 0 - x
      ir::Value* zero = builder_.CreateConstInt(0);
      return static_cast<ir::Value*>(
          builder_.CreateBinary(ir::Opcode::Sub, zero, operand,
                               module_.GetContext().NextTemp()));
    } else if (op == "!") {
      return builder_.CreateNot(operand, module_.GetContext().NextTemp());
    }
  }
  
  throw std::runtime_error(FormatError("irgen", "暂不支持的一元表达式形式"));
}

// 实现函数调用
std::any IRGenImpl::visitFuncRParams(SysYParser::FuncRParamsContext* ctx) {
  if (!ctx) return std::vector<ir::Value*>{};
  std::vector<ir::Value*> args;
  for (auto* exp : ctx->exp()) {
    args.push_back(EvalExpr(*exp));
  }
  return args;
}

std::any IRGenImpl::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  // 常量声明在编译时求值，不分配存储。只需遍历常量定义，将常量值存入符号表即可。
  // 由于 IRGen 阶段不需要常量的存储槽位，这里只触发子节点访问（常量定义会处理值）。
  for (auto* constDef : ctx->constDef()) {
    constDef->accept(this);
  }
  return {};
}

std::any IRGenImpl::visitConstInitVal(SysYParser::ConstInitValContext* ctx) {
  // 常量初始化值可能是单一常量表达式，也可能是大括号列表（聚合）。
  // 为简化，我们仅支持单一表达式（标量常量），数组常量暂未实现。
  if (ctx->constExp()) {
    // 常量表达式求值：需要语义分析支持常量折叠，这里我们返回求值后的常量值。
    // 假设 sema_ 有方法 EvaluateConstExp 返回 int。
    auto* constExp = ctx->constExp();
    // 通过语义分析求值（暂未实现，先抛异常）
    // int val = sema_.EvaluateConstExp(constExp);
    // return builder_.CreateConstInt(val);
    throw std::runtime_error(FormatError("irgen", "常量表达式求值暂未实现"));
  } else {
    // 聚合初始化：返回数组值，暂未实现
    throw std::runtime_error(FormatError("irgen", "常量聚合初始化暂未实现"));
  }
}

std::any IRGenImpl::visitInitVal(SysYParser::InitValContext* ctx) {
  // 返回初始化的值（可能是单个表达式或聚合列表）
  if (ctx->exp()) {
    return EvalExpr(*ctx->exp());
  } else {
    // 聚合初始化：需要返回一个列表，暂未实现
    throw std::runtime_error(FormatError("irgen", "聚合初始化暂未实现"));
  }
}

std::any IRGenImpl::visitConstExp(SysYParser::ConstExpContext* ctx) {
  // 常量表达式由 AddExp 构成，我们直接复用表达式求值（但需要常量折叠）。
  // 由于常量表达式在编译时求值，IR 生成阶段其实不需要计算值，语义分析会提供。
  // 这里我们只用来获取 IR 值（例如数组大小），但数组大小需要在 IR 生成前就知道。
  // 暂时复用 EvalExpr（但 EvalExpr 生成 IR 指令，这不合适）。实际上常量表达式应在语义分析阶段求值，IRGen 直接使用结果。
  // 为避免复杂，我们暂时返回表达式的 IR 值（但会产生计算指令，不适合数组大小）。
  // 更好的做法是调用语义分析求值。
  // 这里先抛异常，等待语义团队提供接口。
  throw std::runtime_error(FormatError("irgen", "常量表达式求值暂未实现"));
}