#include <antlr4-runtime.h>
#include "Public/Public.h"
#include "ICG/GoTo3code.h"
#include "GoParser.h"
#include "GoLexer.h"
#include "GoParserListener.h"
#include "GoParserBaseListener.h"
#include "TCG/Translator.h"
#include <iostream>
#include <vector>
#include <string>
#include <filesystem>
#include <optional>
#include <memory>

using namespace std;

bool ends_with(const std::string& str, const std::string& suffix) {
    if (str.length() < suffix.length()) {
        return false;
    }
    return str.compare(str.length() - suffix.length(), suffix.length(), suffix) == 0;
}


void init_log(const std::string& log_file_name, const std::string& log_path) {
    system(("mkdir -p " + log_path).c_str());
    google::InitGoogleLogging(log_file_name.c_str());
    google::SetLogDestination(google::GLOG_INFO, (log_path + "/").c_str());
    google::InstallFailureSignalHandler(); // Capture signals
    FLAGS_log_prefix = false; // Log prefixes
    FLAGS_colorlogtostderr = true; // Color logs in stderr
    FLAGS_alsologtostderr = true; // Log to console
    FLAGS_logbufsecs = 0; // Disable log buffering
}

struct CmdParam {
    std::vector<std::string> input;
    std::string output;
    std::string output_3code = "3code.txt";
    std::string output_asm = "out.asm";
    std::string output_bin = "out.bin";
    std::string log_path = "./log";
    std::string log_file_name;
};

std::optional<std::string> parse_option(const std::string& option, int& i, int argc, char* argv[]) {
    if (i + 1 < argc) {
        return argv[++i];
    }
    std::cerr << "Option " << option << " requires a value." << std::endl;
    exit(EXIT_FAILURE);
}

CmdParam read_cmd_param(int argc, char* argv[]) {
    CmdParam cmd_param;
    cmd_param.log_file_name = argv[0];

    for (int i = 1; i < argc; ++i) {
        std::string argv_i = argv[i];
        if (argv_i[0] == '-') {
            if (argv_i.length() == 1) {
                std::cerr << "Invalid option." << std::endl;
                exit(EXIT_FAILURE);
            }
            if (auto value = parse_option(argv_i, i, argc, argv)) {
                if (argv_i == "-o" || argv_i == "--output") {
                    cmd_param.output = *value;
                } else if (argv_i == "-l" || argv_i == "--log") {
                    cmd_param.log_path = *value;
                } else {
                    std::cerr << "Unknown parameter " << argv_i << std::endl;
                    exit(EXIT_FAILURE);
                }
            }
        } else {
            cmd_param.input.emplace_back(argv[i]);
        }
    }

    if (cmd_param.input.empty()) {
        std::cerr << "No input .go files provided." << std::endl;
        exit(EXIT_FAILURE);
    }

    for (const auto& input_file : cmd_param.input) {
        std::string base_name = std::filesystem::path(input_file).stem().string();
        if (ends_with(base_name, ".go")) {
            base_name = base_name.substr(0, base_name.length() - 3);
        }
        cmd_param.output = base_name;
        std::filesystem::create_directories(cmd_param.output);
    }

    return cmd_param;
}

int main(int argc, char* argv[]) {
    try
    {
        CmdParam cmd_param = read_cmd_param(argc, argv);
    init_log(cmd_param.log_file_name, cmd_param.log_path);
    const string& filename = cmd_param.input.front();

    LOG(INFO) << "Starting translation for file: " << filename;
    antlr4::ANTLRFileStream file;
    file.loadFromFile(filename);
    antlr4::ANTLRInputStream inputStream(file.toString());

    LOG(INFO) << "Starting lexical analysis.";
    GoLexer lexer(&inputStream);
    antlr4::CommonTokenStream tokens(&lexer);

    LOG(INFO) << "Starting syntax analysis.";
    GoParser parser(&tokens);

    // 遍历AST，自定义生成三地址代码
    GoTo3code listener;
    antlr4::tree::ParseTreeWalker::DEFAULT.walk(&listener, parser.sourceFile());

    std::string path_3code = cmd_param.output + "/" + cmd_param.output_3code;
    listener.Go23file_(path_3code);

    // 三地址转汇编
    LOG(INFO) << "Translating three-address code to assembly.";
    auto tac_file = make_shared<TACFile>(std::move(listener.TACBlocks));
    Translator translator(tac_file, listener.globalScope);
    translator.Translate();
    
    std::string path_asm = cmd_param.output + "/" + cmd_param.output_asm;
    translator.OutputFile(path_asm);

    LOG(INFO) << "Translation completed successfully.";
    return 0;
    }
    catch(...)
    {
        LOG(ERROR) << "Translation completed unsuccessfully.";
        return -1;
    }
    
}