#include <algorithm>
#include <functional>
#include <iostream>
#include <sstream>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "ns3/applications-module.h"
#include "ns3/core-module.h"
#include "ns3/csma-module.h"
#include "ns3/internet-module.h"
#include "ns3/ipv4-global-routing-helper.h"
#include "ns3/network-module.h"
#include "ns3/point-to-point-module.h"
#include "ns3/tcp-l4-protocol.h"

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("InterCluster");

// node num
const int CORE_SWITCH_NUM = 2;
const int AGGREGATION_NUM = 2;
const int TOR_NUM = 4;
const int SERVER_NUM = 8;

const int CORE_THRESHOLD = 0;
const int AGGREGATION_THRESHOLD = CORE_THRESHOLD + CORE_SWITCH_NUM;
const int TOR_THRESHOLD = AGGREGATION_THRESHOLD + AGGREGATION_NUM;
const int SERVER_THRESHOLD = TOR_THRESHOLD + TOR_NUM;

int main(int argc, char* argv[]) {
  LogComponentEnable("PacketSink", LOG_LEVEL_INFO);
  LogComponentEnable("OnOffApplication", LOG_LEVEL_INFO);

  // Enable ECMP routing
  Config::SetDefault("ns3::Ipv4GlobalRouting::RandomEcmpRouting", BooleanValue(true));

  // define all nodes
  NodeContainer all_nodes;
  // core_switch | aggregation_switch | tor_switch | server
  // 0 - 1       | 2 - 3              | 4 - 7      | 8 - 15
  all_nodes.Create(CORE_SWITCH_NUM + AGGREGATION_NUM + TOR_NUM + SERVER_NUM);

  // netdevices
  NetDeviceContainer netdev_container[CORE_SWITCH_NUM * AGGREGATION_NUM +
                                      AGGREGATION_NUM + TOR_NUM];
  NetDeviceContainer* netdev_core_aggr = netdev_container;
  NetDeviceContainer* netdev_aggr_tor =
      netdev_core_aggr + CORE_SWITCH_NUM * AGGREGATION_NUM;
  NetDeviceContainer* netdev_tor_server = netdev_aggr_tor + AGGREGATION_NUM;

  // p2p (core_switch - aggregation_switch)
  NodeContainer p2p_nodes[CORE_SWITCH_NUM][AGGREGATION_NUM];
  // p2p link
  for (int i = 0; i < CORE_SWITCH_NUM; ++i) {
    for (int j = 0; j < AGGREGATION_NUM; ++j) {
      p2p_nodes[i][j].Add(all_nodes.Get(i + CORE_THRESHOLD));
      p2p_nodes[i][j].Add(all_nodes.Get(j + AGGREGATION_THRESHOLD));
    }
  }
  // p2p install
  PointToPointHelper p2p_helper[CORE_SWITCH_NUM][AGGREGATION_NUM];
  for (int i = 0; i < CORE_SWITCH_NUM; ++i) {
    for (int j = 0; j < AGGREGATION_NUM; ++j) {
      p2p_helper[i][j].SetDeviceAttribute("DataRate", StringValue("1.5Mbps"));
      p2p_helper[i][j].SetChannelAttribute("Delay", StringValue("500ns"));
      netdev_core_aggr[i * AGGREGATION_NUM + j] =
          p2p_helper[i][j].Install(p2p_nodes[i][j]);
    }
  }

  // csma
  NodeContainer csma_nodes[AGGREGATION_NUM + TOR_NUM];

  NodeContainer* csma_aggregation_nodes = csma_nodes;
  NodeContainer* csma_tor_nodes = csma_aggregation_nodes + AGGREGATION_NUM;
  // csma link (aggregation_switch - tor_switch)
  for (int i = 0; i < AGGREGATION_NUM; ++i) {
    csma_aggregation_nodes[i].Add(all_nodes.Get(AGGREGATION_THRESHOLD + i));
    csma_aggregation_nodes[i].Add(all_nodes.Get(TOR_THRESHOLD));
    csma_aggregation_nodes[i].Add(all_nodes.Get(TOR_THRESHOLD + 1));
    csma_aggregation_nodes[i].Add(all_nodes.Get(TOR_THRESHOLD + 2));
    csma_aggregation_nodes[i].Add(all_nodes.Get(TOR_THRESHOLD + 3));
  }
  // csma link (tor_switch - server)
  for (int i = 0; i < TOR_NUM; ++i) {
    csma_tor_nodes[i].Add(all_nodes.Get(TOR_THRESHOLD + i));
    csma_tor_nodes[i].Add(all_nodes.Get(SERVER_THRESHOLD + i * 2));
    csma_tor_nodes[i].Add(all_nodes.Get(SERVER_THRESHOLD + i * 2 + 1));
  }
  // csma install (aggregation_switch - tor_switch, tor_switch - server)
  CsmaHelper csma_helper[AGGREGATION_NUM + TOR_NUM];
  for (int i = 0; i < AGGREGATION_NUM + TOR_NUM; ++i) {
    csma_helper[i].SetChannelAttribute("DataRate", StringValue("1.0Mbps"));
    csma_helper[i].SetChannelAttribute("Delay", StringValue("500ns"));
    netdev_aggr_tor[i] = csma_helper[i].Install(csma_nodes[i]);
  }

  // install internet stack for all nodes
  InternetStackHelper stack;
  stack.Install(all_nodes);

  // assign ivp4 addresses
  Ipv4InterfaceContainer ipv4_container[CORE_SWITCH_NUM * AGGREGATION_NUM +
                                        AGGREGATION_NUM + TOR_NUM];
  Ipv4InterfaceContainer* core_aggregation = ipv4_container;
  Ipv4InterfaceContainer* aggregation_tor =
      core_aggregation + CORE_SWITCH_NUM * AGGREGATION_NUM;
  Ipv4InterfaceContainer* tor_server = aggregation_tor + AGGREGATION_NUM;

  // assign tor - server
  for (int i = 0; i < TOR_NUM; ++i) {
    Ipv4AddressHelper addr;
    std::stringstream ss;
    ss << "10.0." << i + 1 << ".0";
    std::string base_addr;
    ss >> base_addr;
    addr.SetBase(base_addr.data(), "255.255.255.0");
    tor_server[i] = addr.Assign(netdev_tor_server[i]);
  }
  // assign aggregation - tor
  for (int i = 0; i < AGGREGATION_NUM; ++i) {
    Ipv4AddressHelper addr;
    std::stringstream ss;
    ss << "10." << i + 1 << ".1.0";
    std::string base_addr;
    ss >> base_addr;
    addr.SetBase(base_addr.data(), "255.255.255.0");
    aggregation_tor[i] = addr.Assign(netdev_aggr_tor[i]);
  }
  // assign core - aggregation
  for (int i = 0; i < CORE_SWITCH_NUM * AGGREGATION_NUM; ++i) {
    Ipv4AddressHelper addr;
    std::stringstream ss;
    ss << "192.168." << i + 1 << ".0";
    std::string base_addr;
    ss >> base_addr;
    addr.SetBase(base_addr.data(), "255.255.255.0");
    core_aggregation[i] = addr.Assign(netdev_core_aggr[i]);
  }

  // set tcp sink app
  ApplicationContainer sinkOrClientApp[8];

  // sink - client settings idx: 1 - 8
  // sink app nodes:   2, 4, 5, 7
  //                   |  |  |  |
  // client app nodes: 6, 8, 1, 3
  const std::vector<std::pair<int, int>> sink_client_pairs = {
      {2, 6}, {4, 8}, {5, 1}, {7, 3}};

  std::for_each(
      sink_client_pairs.begin(), sink_client_pairs.end(),
      [&](const std::pair<int, int>& pair) {
        // c arrays start from 0 idx: 0 - 7
        int sink_idx = pair.first - 1;
        int client_idx = pair.second - 1;
        // sink app setup
        PacketSinkHelper packetSinkHelper(
            "ns3::TcpSocketFactory",
            InetSocketAddress(Ipv4Address::GetAny(), 8080 + sink_idx));
        sinkOrClientApp[sink_idx] = packetSinkHelper.Install(
            all_nodes.Get(SERVER_THRESHOLD + sink_idx));
        sinkOrClientApp[sink_idx].Start(Seconds(0.0));
        sinkOrClientApp[sink_idx].Stop(Seconds(60.0));

        // client app setup
        OnOffHelper onOffHelper(
            "ns3::TcpSocketFactory",
            InetSocketAddress(
                tor_server[sink_idx / 2].GetAddress(sink_idx % 2 + 1),
                8080 + sink_idx));
        onOffHelper.SetAttribute(
            "OnTime", StringValue("ns3::ConstantRandomVariable[Constant=50]"));
        onOffHelper.SetAttribute(
            "OffTime", StringValue("ns3::ConstantRandomVariable[Constant=0]"));
        onOffHelper.SetAttribute("PacketSize", UintegerValue(4096));
        sinkOrClientApp[client_idx] =
            onOffHelper.Install(all_nodes.Get(SERVER_THRESHOLD + client_idx));
        sinkOrClientApp[client_idx].Start(Seconds(0.0));
        sinkOrClientApp[client_idx].Stop(Seconds(60.0));
      });

  // populate routing tables
  Ipv4GlobalRoutingHelper::PopulateRoutingTables();

  // enable pcap at sink nodes
  std::for_each(sink_client_pairs.begin(), sink_client_pairs.end(),
                [&](const std::pair<int, int>& pair) {
                  // c arrays start from 0
                  int sink_idx = pair.first - 1;
                  csma_helper[AGGREGATION_NUM + sink_idx / 2].EnablePcap(
                      "inter-cluster",
                      netdev_tor_server[sink_idx / 2].Get(sink_idx % 2 + 1),
                      false);
                });
  Simulator::Run();
  Simulator::Destroy();

  return 0;
}
