#include "ns3/core-module.h"
#include "ns3/applications-module.h"
#include "ns3/command-line.h"
#include "ns3/simulator.h"
#include "ns3/node.h"
#include "ns3/drop-tail-queue.h"
#include "ns3/timestamp-tag.h"

#include "ns3/tsn-multidrop-net-device.h"
#include "ns3/tsn-multidrop-channel.h"
#include "ns3/ethernet-generator.h"
#include "ns3/ethernet-header2.h"


/**
 * \file
 *
 * Example of the use of tsn-net-device.cc tsn-multidrop-channel.cc on a network
 * composed of four end-stations connected by a 10Base-T1S
 *  ES1 ====== ES2
 *        ||
 *        ||== ES3
 *        ||
 *        ||== ES4
 */

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("Example");


//A callback to log pkt fifo entry time
static void
MacTxCallback(std::string context, Ptr<const Packet> p)
{
  Ptr<Packet> originalPacket = p->Copy();
  EthernetHeader2 ethHeader;
  originalPacket->RemoveHeader(ethHeader);
  NS_LOG_INFO((Simulator::Now()).As(Time::S) << " \t" << context << " : Pkt #" << p->GetUid() << "(VID:" << ethHeader.GetVid() << ") entering the FIFO on the producer");
}

//A callback to log the pkt latency
static void
LatencyCallback(std::string context, Ptr<const Packet> p)
{
  TimestampTag tag;
  if (!p->FindFirstMatchingByteTag(tag))
  {
      return;
  }
  Time arrival = Simulator::Now();
  Time latency = arrival - tag.GetTimestamp();

  Ptr<Packet> originalPacket = p->Copy();
  EthernetHeader2 ethHeader;
  originalPacket->RemoveHeader(ethHeader);
  NS_LOG_INFO((Simulator::Now()).As(Time::S) << " \t" << context << " : Pkt #" << p->GetUid() << " received from " << ethHeader.GetSrc() << "(VID:" << ethHeader.GetVid() << ") with a latency=" << latency.GetNanoSeconds() <<"ns");
}

//A callback to log the PLCA state
static void
PLCAStateCallback(std::string context, TsnMultidropNetDevice::PLCAState state)
{
  std::string stateStr = "";
  if ( state == TsnMultidropNetDevice::PLCAState::DISABLE)
  {
    stateStr = "DISABLE";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::RECOVER)
  {
    stateStr = "RECOVER";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::RESYNC)
  {
    stateStr = "RESYNC";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::SEND_BEACON)
  {
    stateStr = "SEND_BEACON";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::SYNCING)
  {
    stateStr = "SYNCING";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::WAIT_TO)
  {
    stateStr = "WAIT_TO";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::COMMIT)
  {
    stateStr = "COMMIT";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::TRANSMIT)
  {
    stateStr = "TRANSMIT";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::BURST)
  {
    stateStr = "BURST";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::RECEIVE)
  {
    stateStr = "RECEIVE";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::YIELD)
  {
    stateStr = "YIELD";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::ABORT)
  {
    stateStr = "ABORT";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::NEXT_TX_OPPORTUNITY)
  {
    stateStr = "NEXT_TX_OPPORTUNITY";
  }
  else if ( state == TsnMultidropNetDevice::PLCAState::EARLY_RECEIVE)
  {
    stateStr = "EARLY_RECEIVE";
  }
  else
  {
    stateStr = "UNKNOWN_STATE";
  }

  NS_LOG_INFO((Simulator::Now()).As(Time::S) << " \t" << context << " : PLCAState => " << stateStr);
}


int
main(int argc, char* argv[])
{
    //Enable logging
    LogComponentEnable("Example", LOG_LEVEL_INFO);
    LogComponentEnable("EthernetGenerator", LOG_LEVEL_INFO);
    // LogComponentEnable("TsnMultidropNetDevice", LOG_LEVEL_INFO);
    // LogComponentEnable("TsnMultidropChannel", LOG_LEVEL_INFO);

    CommandLine cmd(__FILE__);
    cmd.Parse(argc, argv);

    //Create four nodes
    Ptr<TsnNode> n0 = CreateObject<TsnNode>();
    Names::Add("ES1", n0);
    Ptr<TsnNode> n1 = CreateObject<TsnNode>();
    Names::Add("ES2", n1);
    Ptr<TsnNode> n2 = CreateObject<TsnNode>();
    Names::Add("ES3", n2);
    Ptr<TsnNode> n3 = CreateObject<TsnNode>();
    Names::Add("ES4", n3);


    //Create and add a netDevice to each node
    Ptr<TsnMultidropNetDevice> net0 = CreateObject<TsnMultidropNetDevice>();
    net0->SetAttribute("PLCALocalNodeId", UintegerValue(0));
    net0->SetAttribute("PLCANodeCount", UintegerValue(4));
    // net0->SetAttribute("PLCAMaxBurstCount", UintegerValue(1));
    n0->AddDevice(net0);
    Names::Add("ES1#01", net0);
    Ptr<TsnMultidropNetDevice> net1 = CreateObject<TsnMultidropNetDevice>();
    net1->SetAttribute("PLCALocalNodeId", UintegerValue(1));
    net1->SetAttribute("PLCANodeCount", UintegerValue(4));
    n1->AddDevice(net1);
    Names::Add("ES2#01", net1);
    Ptr<TsnMultidropNetDevice> net2 = CreateObject<TsnMultidropNetDevice>();
    net2->SetAttribute("PLCALocalNodeId", UintegerValue(2));
    net2->SetAttribute("PLCANodeCount", UintegerValue(4));
    n2->AddDevice(net2);
    Names::Add("ES3#01", net2);
    Ptr<TsnMultidropNetDevice> net3 = CreateObject<TsnMultidropNetDevice>();
    net3->SetAttribute("PLCALocalNodeId", UintegerValue(3));
    net3->SetAttribute("PLCANodeCount", UintegerValue(4));
    n3->AddDevice(net3);
    Names::Add("ES4#01", net3);


    //Create a 10Base-T1S Channel and attach it two the netDevices
    Ptr<TsnMultidropChannel> channel = CreateObject<TsnMultidropChannel>();
    net0->Attach(channel);
    net1->Attach(channel);
    net2->Attach(channel);
    net3->Attach(channel);


    //Allocate a Mac address and create a FIFO (for the output port)
    //for each netDevice.
    net0->SetAddress(Mac48Address::Allocate());
    net0->SetQueue(CreateObject<DropTailQueue<Packet>>());

    net1->SetAddress(Mac48Address::Allocate());
    net1->SetQueue(CreateObject<DropTailQueue<Packet>>());

    net2->SetAddress(Mac48Address::Allocate());
    net2->SetQueue(CreateObject<DropTailQueue<Packet>>());

    net3->SetAddress(Mac48Address::Allocate());
    net3->SetQueue(CreateObject<DropTailQueue<Packet>>());


    //Application description
    Ptr<EthernetGenerator> app0 = CreateObject<EthernetGenerator>();
    app0->Setup(net0);
    app0->SetAttribute("BurstSize", UintegerValue(2));
    app0->SetAttribute("PayloadSize", UintegerValue(1400));
    app0->SetAttribute("Period", TimeValue(Seconds(5)));
    app0->SetAttribute("VlanID", UintegerValue(1));

    n0->AddApplication(app0);
    app0->SetStartTime(Seconds(0));
    app0->SetStopTime(Seconds(10));


    Ptr<EthernetGenerator> app1 = CreateObject<EthernetGenerator>();
    app1->Setup(net1);
    app1->SetAttribute("BurstSize", UintegerValue(2));
    app1->SetAttribute("PayloadSize", UintegerValue(1400));
    app1->SetAttribute("Period", TimeValue(Seconds(5)));
    app1->SetAttribute("VlanID", UintegerValue(2));

    n1->AddApplication(app1);
    app1->SetStartTime(Seconds(0));
    app1->SetStopTime(Seconds(10));

    //Callback to log pkt fifo entry time
    std::string context = Names::FindName(n0) + ":" + Names::FindName(net0);
    net0->TraceConnectWithoutContext("MacTx", MakeBoundCallback(&MacTxCallback, context));
    context = Names::FindName(n1) + ":" + Names::FindName(net1);
    net1->TraceConnectWithoutContext("MacTx", MakeBoundCallback(&MacTxCallback, context));
    context = Names::FindName(n2) + ":" + Names::FindName(net2);
    net2->TraceConnectWithoutContext("MacTx", MakeBoundCallback(&MacTxCallback, context));
    context = Names::FindName(n3) + ":" + Names::FindName(net3);
    net3->TraceConnectWithoutContext("MacTx", MakeBoundCallback(&MacTxCallback, context));

    //Callback to display the packet latency
    net0->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyCallback, context));
    context = Names::FindName(n1) + ":" + Names::FindName(net1);
    net1->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyCallback, context));
    context = Names::FindName(n2) + ":" + Names::FindName(net2);
    net2->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyCallback, context));
    context = Names::FindName(n3) + ":" + Names::FindName(net3);
    net3->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyCallback, context));


    net0->TraceConnectWithoutContext("PLCAState", MakeBoundCallback(&PLCAStateCallback, Names::FindName(net0)));
    net1->TraceConnectWithoutContext("PLCAState", MakeBoundCallback(&PLCAStateCallback, Names::FindName(net1)));
    net2->TraceConnectWithoutContext("PLCAState", MakeBoundCallback(&PLCAStateCallback, Names::FindName(net2)));
    net3->TraceConnectWithoutContext("PLCAState", MakeBoundCallback(&PLCAStateCallback, Names::FindName(net3)));

    //Execute the simulation
    Simulator::Stop(MilliSeconds(5));
    Simulator::Run();
    Simulator::Destroy();
    return 0;
}