#include <iostream>
#include <ostream>
#include "antlr4-runtime.h"
#include "ProtoParserBaseListener.h"
#include "ProtoLexer.h"
#include "ProtoParser.h"
#include "message.h"
#include "service.h"
#include "plugins/c/c_message_plugin.h"
#include "plugins/c/c_service_server_plugin.h"
#include "plugins/c/c_service_client_plugin.h"
#include <stack>
#include <map>
#include <unistd.h>

class Listener : public ProtoParserBaseListener {
public:
    void enterMessageStatement(ProtoParser::MessageStatementContext *context) override
    {
        message_stack.push(std::make_shared<tqcq::Message>(context->ID()->getText()));
    }

    void exitMessageStatement(ProtoParser::MessageStatementContext *context) override
    {
        message_stack.top()->SetItems(std::move(items));
        assert(items.empty());
        if (message_map.find(message_stack.top()->id()) != message_map.end()) {
            std::cout << "message " << message_stack.top()->id() << " already exists" << std::endl;
            exit(1);
        }

        message_map.insert(std::make_pair(message_stack.top()->id(), message_stack.top()));
        message_stack.pop();
    }

    void exitItemStatement(ProtoParser::ItemStatementContext *context) override
    {
        int idx = std::stoi(context->NUMBER()->getText());
        tqcq::Item::Ptr item = nullptr;
        if (context->INT32()) {
            item = std::make_shared<tqcq::Item>(tqcq::Item::Type::kInt32, context->ID()->getText(), idx);
        } else if (context->STRING()){
            item = std::make_shared<tqcq::Item>(tqcq::Item::Type::kString, context->ID()->getText(), idx);
        } else {
            assert(false);
        }
        items.push_back(std::move(item));
    }

    void exitServiceStatement(ProtoParser::ServiceStatementContext *context) override
    {
        std::string id = context->ID()->getText();
        auto service = std::make_shared<tqcq::Service>(id);
        service->SetRPCs(std::move(rpcs));
        assert(rpcs.empty());
        services.push_back(std::move(service));
    }

    void exitRpcStatement(ProtoParser::RpcStatementContext *context) override {
        std::string id = context->ID(0)->getText();
        std::string request_id = context->ID(1)->getText();
        std::string response_id = context->ID(2)->getText();
        auto rpc = std::make_shared<tqcq::RPC>(id, request_id, response_id);
        rpcs.push_back(std::move(rpc));
    }

    std::map<std::string, tqcq::Message::Ptr> message_map;
    std::vector<tqcq::Item::Ptr> items;
    std::stack<tqcq::Message::Ptr> message_stack;

    std::vector<tqcq::RPC::Ptr> rpcs;
    std::vector<tqcq::Service::Ptr> services;
};

int
main(int argc, char *argv[])
{
    if (argc < 3) {
        std::cout << "usage: " << argv[0] << " <proto file> <output dir>" << std::endl;
        return 1;
    }
    std::string proto_file = argv[1];
    std::string output_dir = argv[2];

    // show current working directory
    char cwd[1024];
    if (getcwd(cwd, sizeof(cwd)) != NULL) {
        std::cout << "Current working dir: " << cwd << std::endl;
    } else {
        perror("getcwd() error");
        return 1;
    }

    std::ifstream ifs;
    ifs.open(proto_file);

    antlr4::ANTLRInputStream input(ifs);
    ProtoLexer lexer(&input);
    antlr4::CommonTokenStream tokens(&lexer);
    ProtoParser parser(&tokens);
    // parser.removeErrorListeners();
    // Visitor visitor;
    // visitor.visit(parser.program());
    Listener listener;
    Listener listener2;
    antlr4::tree::ParseTreeWalker::DEFAULT.walk(&listener, parser.program());
//    for (const auto& item : listener.message_map) {
//        std::cout << item.second->ToString() << std::endl;
//    }

//    for (const auto& service : listener.services) {
//        std::cout << service->ToString() << std::endl;
//    }

    tqcq::CMessagePlugin message_plugin(output_dir);
    message_plugin.Generate(listener.message_map);

    tqcq::CServiceServerPlugin service_server_plugin(output_dir);
    service_server_plugin.Generate(listener.services);
    // tqcq::CServiceClientPlugin service_client_plugin("./gen/");

    return 0;
}