#include "sem/Sema.h"

#include <stdexcept>
#include <string>

#include "sem/SymbolTable.h"
#include "utils/Log.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(FormatError("sema", "使用了未定义的变量: " + name));
    }
    return;
  }

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

  throw std::runtime_error(FormatError("sema", "暂不支持的表达式形式"));
}

void CheckExpr(SysYParser::ExpContext& exp, const SymbolTable& table) {
  if (!exp.addExp()) {
    throw std::runtime_error(FormatError("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(FormatError("sema", "缺少 main 函数定义"));
  }
  if (!func->Ident() || func->Ident()->getText() != "main") {
    throw std::runtime_error(FormatError("sema", "入口函数必须命名为 main"));
  }

  SymbolTable table;
  bool seen_return = false;

  const auto& items = func->block()->blockItem();
  if (items.empty()) {
    throw std::runtime_error(
        FormatError("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(
          FormatError("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(FormatError("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(
            FormatError("sema", "return 必须是 main 函数中的最后一条语句"));
      }
      continue;
    }
    throw std::runtime_error(FormatError("sema", "暂不支持的语句或声明"));
  }

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