#include "sem/Sema.h"

#include <stdexcept>
#include <string>

#include "sem/SymbolTable.h"

namespace {

void CheckExpr(SysYParser::ExpContext& exp, const SymbolTable& table);

void CheckPrimary(SysYParser::PrimaryContext& primary,
                  const SymbolTable& table) {
  if (primary.Number()) {
    return;
  }

  if (primary.Ident()) {
    const std::string name = primary.Ident()->getText();
    if (!table.Contains(name)) {
      throw std::runtime_error("[sema] 使用了未定义的变量: " + name);
    }
    return;
  }

  if (primary.exp()) {
    CheckExpr(*primary.exp(), table);
    return;
  }

  throw std::runtime_error("[sema] 暂不支持的 primary 形式");
}

void CheckExpr(SysYParser::ExpContext& exp, const SymbolTable& table) {
  if (!exp.addExp()) {
    throw std::runtime_error("[sema] 非法表达式");
  }
  const auto& terms = exp.addExp()->primary();
  for (auto* term : terms) {
    CheckPrimary(*term, table);
  }
}

}  // namespace

void RunSema(SysYParser::CompUnitContext& comp_unit) {
  auto* func = comp_unit.funcDef();
  if (!func || !func->block()) {
    throw std::runtime_error("[sema] 缺少 main 函数定义");
  }
  if (!func->Ident() || func->Ident()->getText() != "main") {
    throw std::runtime_error("[sema] 入口函数必须命名为 main");
  }

  SymbolTable table;
  bool seen_return = false;

  const auto& items = func->block()->blockItem();
  if (items.empty()) {
    throw std::runtime_error("[sema] main 函数不能为空，且必须以 return 结束");
  }

  for (size_t i = 0; i < items.size(); ++i) {
    auto* item = items[i];
    if (!item) {
      continue;
    }
    if (seen_return) {
      throw std::runtime_error("[sema] return 必须是 main 函数中的最后一条语句");
    }
    if (auto* decl = item->decl() ? item->decl()->varDecl() : nullptr) {
      const std::string name = decl->Ident()->getText();
      if (table.Contains(name)) {
        throw std::runtime_error("[sema] 重复定义变量: " + name);
      }
      if (decl->exp()) {
        CheckExpr(*decl->exp(), table);
      }
      table.Add(name);
      continue;
    }
    if (auto* stmt = item->stmt(); stmt && stmt->returnStmt()) {
      auto* ret = stmt->returnStmt();
      CheckExpr(*ret->exp(), table);
      seen_return = true;
      if (i + 1 != items.size()) {
        throw std::runtime_error("[sema] return 必须是 main 函数中的最后一条语句");
      }
      continue;
    }
    throw std::runtime_error("[sema] 暂不支持的 blockItem 类型");
  }

  if (!seen_return) {
    throw std::runtime_error("[sema] main 函数必须包含 return 语句");
  }
}