#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));
}


std::any IRGenImpl::visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "非法基本表达式"));
  }
  // 处理括号表达式：LPAREN exp RPAREN
  if (ctx->exp()) {
    return EvalExpr(*ctx->exp());
  }
  // 处理 lVal（变量使用）- 交给 visitLVal 处理
  if (ctx->lVal()) {
    // 直接在这里处理变量读取，避免 accept 调用可能导致的问题
    auto* lval_ctx = ctx->lVal();
    if (!lval_ctx || !lval_ctx->ID()) {
      throw std::runtime_error(FormatError("irgen", "当前仅支持普通整型变量"));
    }
    const auto* decl = sema_.ResolveObjectUse(lval_ctx);
    if (!decl) {
      throw std::runtime_error(
          FormatError("irgen",
                      "变量使用缺少语义绑定：" + lval_ctx->ID()->getText()));
    }
    std::string var_name = lval_ctx->ID()->getText();
    auto it = storage_map_.find(var_name);
    if (it == storage_map_.end()) {
      throw std::runtime_error(
          FormatError("irgen",
                      "变量声明缺少存储槽位：" + var_name));
    }
    return static_cast<ir::Value*>(
        builder_.CreateLoad(it->second, module_.GetContext().NextTemp()));
  }
  // 处理 number
  if (ctx->number()) {
    return ctx->number()->accept(this);
  }
  throw std::runtime_error(FormatError("irgen", "不支持的基本表达式类型"));
}


std::any IRGenImpl::visitNumber(SysYParser::NumberContext* ctx) {
  if (!ctx || !ctx->intConst()) {
    throw std::runtime_error(FormatError("irgen", "当前仅支持整数字面量"));
  }
  return static_cast<ir::Value*>(
      builder_.CreateConstInt(std::stoi(ctx->intConst()->getText())));
}

// 变量使用的处理流程：
// 1. 先通过语义分析结果把变量使用绑定回声明；
// 2. 再通过 storage_map_ 找到该声明对应的栈槽位；
// 3. 最后生成 load，把内存中的值读出来。
//
// 因此当前 IRGen 自己不再做名字查找，而是直接消费 Sema 的绑定结果。
std::any IRGenImpl::visitLVal(SysYParser::LValContext* ctx) {
  if (!ctx || !ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "当前仅支持普通整型变量"));
  }
  const auto* decl = sema_.ResolveObjectUse(ctx);
  if (!decl) {
    throw std::runtime_error(
        FormatError("irgen",
                    "变量使用缺少语义绑定：" + ctx->ID()->getText()));
  }
  // 使用变量名查找存储槽位
  std::string var_name = ctx->ID()->getText();
  auto it = storage_map_.find(var_name);
  if (it == storage_map_.end()) {
    throw std::runtime_error(
        FormatError("irgen",
                    "变量声明缺少存储槽位：" + var_name));
  }
  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", "非法加减法表达式"));
  }
  
  // 如果是 mulExp 直接返回（addExp : mulExp）
  if (ctx->mulExp() && ctx->addExp() == nullptr) {
    return ctx->mulExp()->accept(this);
  }
  
  // 处理 addExp op mulExp 的递归形式
  if (!ctx->addExp() || !ctx->mulExp()) {
    throw std::runtime_error(FormatError("irgen", "非法加减法表达式结构"));
  }
  
  ir::Value* lhs = std::any_cast<ir::Value*>(ctx->addExp()->accept(this));
  ir::Value* rhs = std::any_cast<ir::Value*>(ctx->mulExp()->accept(this));
  
  ir::Opcode op = ir::Opcode::Add;
  if (ctx->ADD()) {
    op = ir::Opcode::Add;
  } else if (ctx->SUB()) {
    op = ir::Opcode::Sub;
  } else {
    throw std::runtime_error(FormatError("irgen", "未知的加减运算符"));
  }
  
  return static_cast<ir::Value*>(
      builder_.CreateBinary(op, lhs, rhs, module_.GetContext().NextTemp()));
}


std::any IRGenImpl::visitUnaryExp(SysYParser::UnaryExpContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "非法一元表达式"));
  }
  
  // 如果是 primaryExp 直接返回（unaryExp : primaryExp）
  if (ctx->primaryExp()) {
    return ctx->primaryExp()->accept(this);
  }
  
  // 处理函数调用（unaryExp : ID LPAREN funcRParams? RPAREN）
  // 当前暂不支持，留给后续扩展
  if (ctx->ID()) {
    throw std::runtime_error(FormatError("irgen", "暂不支持函数调用"));
  }
  
  // 处理一元运算符（unaryExp : addUnaryOp unaryExp）
  if (ctx->addUnaryOp() && ctx->unaryExp()) {
    ir::Value* operand = std::any_cast<ir::Value*>(ctx->unaryExp()->accept(this));
    
    // 判断是正号还是负号
    if (ctx->addUnaryOp()->SUB()) {
      // 负号：生成 sub 0, operand（LLVM IR 中没有 neg 指令）
      ir::Value* zero = builder_.CreateConstInt(0);
      return static_cast<ir::Value*>(
          builder_.CreateSub(zero, operand, module_.GetContext().NextTemp()));
    } else if (ctx->addUnaryOp()->ADD()) {
      // 正号：直接返回操作数（+x 等价于 x）
      return operand;
    } else {
      throw std::runtime_error(FormatError("irgen", "未知的一元运算符"));
    }
  }
  
  throw std::runtime_error(FormatError("irgen", "不支持的一元表达式类型"));
}

std::any IRGenImpl::visitMulExp(SysYParser::MulExpContext* ctx) {
  if (!ctx) {
    throw std::runtime_error(FormatError("irgen", "非法乘除法表达式"));
  }
  
  // 如果是 unaryExp 直接返回（mulExp : unaryExp）
  if (ctx->unaryExp() && ctx->mulExp() == nullptr) {
    return ctx->unaryExp()->accept(this);
  }
  
  // 处理 mulExp op unaryExp 的递归形式
  if (!ctx->mulExp() || !ctx->unaryExp()) {
    throw std::runtime_error(FormatError("irgen", "非法乘除法表达式结构"));
  }
  
  ir::Value* lhs = std::any_cast<ir::Value*>(ctx->mulExp()->accept(this));
  ir::Value* rhs = std::any_cast<ir::Value*>(ctx->unaryExp()->accept(this));
  
  ir::Opcode op = ir::Opcode::Mul;
  if (ctx->MUL()) {
    op = ir::Opcode::Mul;
  } else if (ctx->DIV()) {
    op = ir::Opcode::Div;
  } else if (ctx->MOD()) {
    op = ir::Opcode::Mod;
  } else {
    throw std::runtime_error(FormatError("irgen", "未知的乘除运算符"));
  }
  
  return static_cast<ir::Value*>(
      builder_.CreateBinary(op, lhs, rhs, module_.GetContext().NextTemp()));
}