#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/tsn-node.h"
#include "ns3/tsn-net-device.h"
#include "ns3/ethernet-generator.h"
#include "ns3/ethernet-header2.h"
#include "ns3/ethernet-channel.h"
#include "ns3/switch-net-device.h"


/**
 * \file
 *
 * Example of the use of tsn-net-device.cc tsn-multidrop-channel.cc on a network
 * switched network with three end-stations connected by a 10Base-T1S
 *  ES1 ==== SW ====== ES2
 *                ||
 *                ||== ES3
 *                ||
 *                ||== ES4
 *
 * There is two flows on this network :
 *        - vlanID = 10 from ES1 to ES2
 *        - vlanID = 20 from ES2 to ES1, 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("TsnMultidropNetChannel", 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);
    Ptr<TsnNode> n4 = CreateObject<TsnNode>();
    Names::Add("SW", n4);


    //Create and add a netDevice to each node
    Ptr<TsnNetDevice> net0 = CreateObject<TsnNetDevice>();
    n0->AddDevice(net0);
    Names::Add("ES1#01", net0);
    Ptr<TsnNetDevice> swnet0 = CreateObject<TsnNetDevice>();
    n4->AddDevice(swnet0);
    Names::Add("SW#01", swnet0);

    Ptr<TsnMultidropNetDevice> swnet1 = CreateObject<TsnMultidropNetDevice>();
    swnet1->SetAttribute("PLCALocalNodeId", UintegerValue(0));
    swnet1->SetAttribute("PLCANodeCount", UintegerValue(4));
    n4->AddDevice(swnet1);
    Names::Add("SW#02", swnet1);
    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 full-duplex channel
    Ptr<EthernetChannel> channel0 = CreateObject<EthernetChannel>();
    net0->Attach(channel0);
    swnet0->Attach(channel0);

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

    //Create and add a switch net device to the switch node
    Ptr<SwitchNetDevice> sw = CreateObject<SwitchNetDevice>();
    sw->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
    sw->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
    n4->AddDevice(sw);
    sw->AddSwitchPort(swnet0);
    sw->AddSwitchPort(swnet1);


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

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

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

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

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

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

    //Add forwarding table
    sw->AddForwardingTableEntry(Mac48Address::ConvertFrom(net2->GetAddress()), 10, {swnet1});
    sw->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 20, {swnet0});

    //Application description
    //From switched network to  10Base-T1S
    Ptr<EthernetGenerator> app0 = CreateObject<EthernetGenerator>();
    app0->Setup(net0);
    app0->SetAttribute("Address", AddressValue(net2->GetAddress()));
    app0->SetAttribute("BurstSize", UintegerValue(2));
    app0->SetAttribute("PayloadSize", UintegerValue(1400));
    app0->SetAttribute("Period", TimeValue(Seconds(5)));
    app0->SetAttribute("VlanID", UintegerValue(10));
    app0->SetAttribute("PCP", UintegerValue(0));

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

    //From 10Base-T1S to switched network
    Ptr<EthernetGenerator> app1 = CreateObject<EthernetGenerator>();
    app1->Setup(net1);
    app1->SetAttribute("BurstSize", UintegerValue(1));
    app1->SetAttribute("PayloadSize", UintegerValue(100));
    app1->SetAttribute("Period", TimeValue(Seconds(5)));
    app1->SetAttribute("VlanID", UintegerValue(20));
    app1->SetAttribute("PCP", UintegerValue(1));

    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));

    //Callback to display the packet latency
    context = Names::FindName(n0) + ":" + Names::FindName(net0);
    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));

    //Callback to display PLCA state evolution
    swnet1->TraceConnectWithoutContext("PLCAState", MakeBoundCallback(&PLCAStateCallback, Names::FindName(swnet1)));
    // 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(3));
    Simulator::Run();
    Simulator::Destroy();
    return 0;
}