#include "3rdparty/easylogging++.h"
#include "BaseErrorListener.h"
#include "Exceptions.h"
#include "antlr4-runtime.h"
#include "antlrgen/SysyLexer.h"
#include "antlrgen/SysyParser.h"
#include "common.h"
#include "llir.h"
#include "passes.h"
#include "visitor.h"
#include <3rdparty/argparse.hpp>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>

INITIALIZE_EASYLOGGINGPP

using namespace antlrSysY;
using namespace antlr4;

class AbortErrorListener : public BaseErrorListener {
public:
  void syntaxError(
    Recognizer *recognizer,
    Token *offendingSymbol,
    size_t line,
    size_t charPositionInLine,
    const std::string &msg,
    std::exception_ptr e
  ) override {
    throw ParseCancellationException(
      "line " + std::to_string(line) + ":" + std::to_string(charPositionInLine) + " " + msg
    );
  }
};

int main(int argc, const char **argv) {
#pragma region ArgParse
  argparse::ArgumentParser arg_parser("Catfood's SysY Compiler");
  arg_parser.add_argument("source").help("A single SysY source file");
  arg_parser.add_argument("-S").implicit_value(true).help("Useless but required by the contest").required();
  arg_parser.add_argument("-o").help("Output file name").required();
  arg_parser.add_argument("-O1").implicit_value(true).default_value(false).help("Performance mode");
  arg_parser.add_argument("-emit-llvm").implicit_value(true).default_value(false).help("Generate llvm ir");
  arg_parser.add_argument("--v").default_value(0);
  try {
    arg_parser.parse_args(argc, argv);
  } catch (const std::runtime_error &err) {
    std::cerr << err.what() << std::endl;
    std::cerr << arg_parser;
    std::exit(1);
  }
  auto source_file = arg_parser.get<std::string>("source");
  auto output_file = arg_parser.get<std::string>("-o");
  auto flg_O1      = arg_parser["-O1"] == true;
  auto emit_llvm   = arg_parser["-emit-llvm"] == true;
  // std::cout << source_file << " " << output_file << " " << flg_O1 <<
  // std::endl;
#pragma endregion
  START_EASYLOGGINGPP(argc, argv);
#pragma region Logger
  el::Configurations defaultConf;
  defaultConf.setToDefault();
  defaultConf.set(el::Level::Debug, el::ConfigurationType::Format, "%level %loc %msg");
  defaultConf.set(el::Level::Warning, el::ConfigurationType::Format, "%level %loc %msg");
  defaultConf.set(el::Level::Error, el::ConfigurationType::Format, "%level %loc %msg");
  defaultConf.set(el::Level::Info, el::ConfigurationType::Format, "%level %loc %msg");
  defaultConf.set(el::Level::Verbose, el::ConfigurationType::Format, "%level %loc %msg");
  el::Loggers::addFlag(el::LoggingFlag::ColoredTerminalOutput);
  el::Loggers::reconfigureLogger("default", defaultConf);
#pragma endregion

  std::ifstream ifs_source_file(source_file);
  if (!ifs_source_file.good()) {
    LOG(ERROR) << "Source file " << source_file << " not found.";
    return 1;
  }
  ANTLRInputStream input(ifs_source_file);
  SysyLexer lexer(&input);
  CommonTokenStream tokens(&lexer);

  // tokens.fill();
  // for (auto token : tokens.getTokens()) {
  //   std::cout << token->toString() << std::endl;
  // }

  SysyParser parser(&tokens);
  parser.removeErrorListeners();
  parser.addErrorListener(new AbortErrorListener());
  auto tree = parser.program();
  Visitor visitor(lexer);
  visitor.visitProgram(tree);

  std::vector<std::shared_ptr<Pass>> passes = {
    std::make_shared<PassBuildCFG>(),
    std::make_shared<PassMem2Reg>(),
  };

  for (auto pass : passes) {
    pass->run(visitor.module);
  }

  if (emit_llvm) {
    auto llir_file = output_file.substr(0, output_file.rfind(".")) + ".ll";
    std::ofstream ofs_llir_file(llir_file);
    if (!ofs_llir_file.good()) {
      LOG(ERROR) << "Failed to create/overwrite LLIR output file " << llir_file;
      return 1;
    }
    visitor.llir_gen(ofs_llir_file);
  }

  // std::cout << tree->toStringTree(&parser) << std::endl << std::endl;

  return 0;
}