main.cpp

#include "src/app_query_handler.h"
#include "src/app_report_receiver.h"
#include "src/global_variables.h"
#include "src/probe_feedback_handler.h"
#include "src/rtt_report_sender.h"
#include "src/tcp_prober.h"
#include <atomic>
#include <chrono>
#include <iostream>
#include <string>
#include <thread>

std::atomic<bool> g_running(true);

void PrintUsage(const char *program_name) {
  std::cout << "Usage: " << program_name << " --index <number> [options]"
            << std::endl;
  std::cout << "Required:" << std::endl;
  std::cout << "  --index <number>   Instance index for unique socket naming "
               "(required)"
            << std::endl;
  std::cout << "Options:" << std::endl;
  std::cout
      << "  --ip <address>     Destination IP address (default: 192.168.2.3)"
      << std::endl;
  std::cout << "  --port <number>    Destination port (default: 20000)"
            << std::endl;
  std::cout << "  --rtt-port <number>       RTT report port (default: 20004)"
            << std::endl;
  std::cout << "  --help             Display this help message" << std::endl;
}

int main(int argc, char *argv[]) {
  // Default parameters
  std::string dest_ip = "192.168.2.3";
  int dest_port = 20000;
  int rtt_port = 20004;
  int index = -1; // Required parameter, -1 indicates not set

  // Parse command line arguments
  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];

    if (arg == "--help") {
      PrintUsage(argv[0]);
      return 0;
    } else if (arg == "--index" && i + 1 < argc) {
      try {
        index = std::stoi(argv[++i]);
        if (index < 0) {
          std::cerr << "Error: Index must be non-negative" << std::endl;
          return 1;
        }
      } catch (const std::exception &e) {
        std::cerr << "Error: Invalid index number" << std::endl;
        return 1;
      }
    } else if (arg == "--ip" && i + 1 < argc) {
      dest_ip = argv[++i];
    } else if (arg == "--port" && i + 1 < argc) {
      try {
        dest_port = std::stoi(argv[++i]);
        if (dest_port <= 0 || dest_port > 65535) {
          std::cerr << "Error: Port must be between 1 and 65535" << std::endl;
          return 1;
        }
      } catch (const std::exception &e) {
        std::cerr << "Error: Invalid port number" << std::endl;
        return 1;
      }
    } else if (arg == "--rtt-port" && i + 1 < argc) {
      try {
        rtt_port = std::stoi(argv[++i]);
        if (rtt_port <= 0 || rtt_port > 65535) {
          std::cerr << "Error: RTT port must be between 1 and 65535"
                    << std::endl;
          return 1;
        }
      } catch (const std::exception &e) {
        std::cerr << "Error: Invalid RTT port number" << std::endl;
        return 1;
      }
    } else {
      std::cerr << "Unknown option: " << arg << std::endl;
      PrintUsage(argv[0]);
      return 1;
    }
  }

  // Check if required index parameter was provided
  if (index == -1) {
    std::cerr << "Error: --index parameter is required" << std::endl;
    PrintUsage(argv[0]);
    return 1;
  }

  // Create unique socket paths using index
  std::string response_socket_path =
      "/tmp/edge_client_response_" + std::to_string(index) + ".sock";
  std::string rtt_socket_path =
      "/tmp/edge_client_rtt_" + std::to_string(index) + ".sock";

  // Set the global socket paths
  RESPONSE_SOCKET_PATH = response_socket_path;
  RTT_SOCKET_PATH = rtt_socket_path;

  std::cout << "Starting TCP prober to " << dest_ip << ":" << dest_port
            << std::endl;
  std::cout << "RTT port: " << rtt_port << std::endl;
  std::cout << "Response socket path: " << response_socket_path << std::endl;
  std::cout << "RTT socket path: " << rtt_socket_path << std::endl;

  // Create TCP prober and handlers
  TCPProber prober;
  AppQueryHandler query_handler;
  ProbeFeedbackHandler feedback_handler;

  // Create and start TCP prober in a separate thread
  std::thread probe_thread(
      [&prober, &query_handler, &feedback_handler, &dest_ip, dest_port, index]() {
        // Start TCP prober
        if (prober.Start(dest_ip, dest_port)) {
          // Get the client key and set it in query handler
          ClientKey client_key = prober.GetClientKey();
          query_handler.SetClientKey(client_key);
          // Start the query handler
          query_handler.Start(static_cast<uint32_t>(index));
          
          // Start the probe feedback handler
          feedback_handler.Start(static_cast<uint32_t>(index), &prober);

          std::cout << "TCP prober and handlers started successfully"
                    << std::endl;
        }

        // Keep the prober running
        while (g_running) {
          std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }

        // Stop all services
        query_handler.Stop();
        feedback_handler.Stop();
        prober.Stop();
      });

  // Create and start receiver in a separate thread
  std::thread receiver_thread(ReceiveAppMessages);

  // Create and start RTT receiver in a separate thread
  std::thread rtt_receiver_thread(ReceiveRTTReports);

  // Create and start RTT sender in a separate thread
  std::thread rtt_sender_thread([&dest_ip, rtt_port]() {
    RunRTTSender(dest_ip, rtt_port, g_running);
  });

  probe_thread.detach();
  receiver_thread.detach();
  rtt_receiver_thread.detach();
  rtt_sender_thread.detach();

  // Wait until receiver is stopped
  while (g_running) {
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
  }

  // Stop RTT sender
  StopRTTSender();

  // Give the threads a moment to clean up
  std::this_thread::sleep_for(std::chrono::seconds(1));

  std::cout << "Server scheduler application exited" << std::endl;
  return 0;
}