#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/tsn-node.h"
#include "ns3/tsn-net-device.h"
#include "ns3/tsn-aggregated-net-device.h"
#include "ns3/switch-net-device.h"
#include "ns3/ethernet-channel.h"
#include "ns3/ethernet-generator.h"
#include "ns3/ethernet-header2.h"
#include "ns3/stream-identification-function.h"
#include "ns3/stream-identification-function-null.h"
#include "ns3/frer-sequence-generation-function.h"
#include "ns3/frer-sequence-recovery-function.h"
#include "ns3/frer-latent-error-detection-function.h"
#include "ns3/frer-match-recovery-function.h"


/**
 * \file
 *
 * Example with 2ES connected 2SW in a 1Gb/s full duplex link to demonstrate
 * FRER usage when replication and elimation are done on the ES. To be more
 * specific this example use the ESsource to do the replication using mulitcast
 * and ESdest to do the elimination. The stream identification is
 * done by the null stream identification function. The match recovery function,
 * latent error detection function are also used in this example.
 * Two flows go from ESsource to ESdest. Only the VLAN 100 flow use FRER mechanisms.
 *
 *          / ==== SW1 ==== \
 *  ESsource                 ESdest
 *          \ ==== SW2 ==== /
 *
 * Although mutliple ways of implementing elimination and replication on ES are
 * describe in the standard (IEEE802.1CB-2017 - Annex C), here we implement the
 * one describe in "IEEE802.1CB-2017 - Annex C - C.1 Example 1: End-to-end FRER"
 * . This ways is based on the use of IEEE802.1AX for port aggregation.
 *
 */


using namespace ns3;

NS_LOG_COMPONENT_DEFINE("Example");

//A callback to log the pkt reception
static void
MacRxCallback(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() << ") received !");
}

//A callback to log the pkt emission
static void
PhyTxCallback(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() << ") begin transmission !");
}

//A callback to log pkt elimination
static void
FrerDrop(std::string context, Ptr<const Packet> p)
{
  NS_LOG_INFO(context << " : An instance of Packet #"<< p->GetUid() <<" was dropped by FRER recovery function");
}


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

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

  //Create six nodes
  Ptr<TsnNode> n0 = CreateObject<TsnNode>();
  Names::Add("ESsource", n0);

  Ptr<TsnNode> n1 = CreateObject<TsnNode>();
  Names::Add("ESdest", n1);

  Ptr<TsnNode> n2 = CreateObject<TsnNode>();
  Names::Add("SW1", n2);
  Ptr<TsnNode> n3 = CreateObject<TsnNode>();
  Names::Add("SW2", n3);

  //Create and add a netDevices to each node
  Ptr<TsnNetDevice> n0_net0 = CreateObject<TsnNetDevice>();
  n0->AddDevice(n0_net0);
  Names::Add("ESsource#01", n0_net0);
  Ptr<TsnNetDevice> n0_net1 = CreateObject<TsnNetDevice>();
  n0->AddDevice(n0_net1);
  Names::Add("ESsource#02", n0_net1);
  Ptr<TsnAggregatedNetDevice> n0_netagg = CreateObject<TsnAggregatedNetDevice>();
  n0_netagg->AddNetDevice(n0_net0);
  n0_netagg->AddNetDevice(n0_net1);
  n0->AddDevice(n0_netagg);

  Ptr<TsnNetDevice> n1_net0 = CreateObject<TsnNetDevice>();
  n1->AddDevice(n1_net0);
  Names::Add("ESdest#01", n1_net0);
  Ptr<TsnNetDevice> n1_net1 = CreateObject<TsnNetDevice>();
  n1->AddDevice(n1_net1);
  Names::Add("ESdest#02", n1_net1);
  Ptr<TsnAggregatedNetDevice> n1_netagg = CreateObject<TsnAggregatedNetDevice>();
  n1_netagg->AddNetDevice(n1_net0);
  n1_netagg->AddNetDevice(n1_net1);
  n1->AddDevice(n1_netagg);

  Ptr<TsnNetDevice> n2_net0 = CreateObject<TsnNetDevice>();
  n2->AddDevice(n2_net0);
  Ptr<TsnNetDevice> n2_net1 = CreateObject<TsnNetDevice>();
  n2->AddDevice(n2_net1);

  Ptr<TsnNetDevice> n3_net0 = CreateObject<TsnNetDevice>();
  n3->AddDevice(n3_net0);
  Ptr<TsnNetDevice> n3_net1 = CreateObject<TsnNetDevice>();
  n3->AddDevice(n3_net1);

  //Create Ethernet Channels and attach it to the netDevices
  Ptr<EthernetChannel> l0 = CreateObject<EthernetChannel>();
  l0->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  n0_net0->Attach(l0);
  n2_net0->Attach(l0);
  Ptr<EthernetChannel> l1 = CreateObject<EthernetChannel>();
  l1->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  n2_net1->Attach(l1);
  n1_net0->Attach(l1);
  Ptr<EthernetChannel> l2 = CreateObject<EthernetChannel>();
  l2->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  n0_net1->Attach(l2);
  n3_net0->Attach(l2);
  Ptr<EthernetChannel> l3 = CreateObject<EthernetChannel>();
  l3->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  n3_net1->Attach(l3);
  n1_net1->Attach(l3);

  //Create and add switche net devices to the switch nodes
  Ptr<SwitchNetDevice> sw1 = CreateObject<SwitchNetDevice>();
  sw1->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
  sw1->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
  n2->AddDevice(sw1);
  sw1->AddSwitchPort(n2_net0);
  sw1->AddSwitchPort(n2_net1);
  Ptr<SwitchNetDevice> sw2 = CreateObject<SwitchNetDevice>();
  sw2->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
  sw2->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
  n3->AddDevice(sw2);
  sw2->AddSwitchPort(n3_net0);
  sw2->AddSwitchPort(n3_net1);

  //Allocate a Mac address
  n0_net0->SetAddress(Mac48Address::Allocate());
  n0_net1->SetAddress(Mac48Address::Allocate());
  n1_net0->SetAddress(Mac48Address::Allocate());
  n1_net1->SetAddress(Mac48Address::Allocate());
  sw1->SetAddress(Mac48Address::Allocate());
  sw2->SetAddress(Mac48Address::Allocate());

  //Create and add eight FIFO on each net device
  for (int i=0; i<8; i++)
  {
    n0_net0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n0_net1->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n1_net0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n1_net1->SetQueue(CreateObject<DropTailQueue<Packet>>());

    n2_net0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n2_net1->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n3_net0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    n3_net1->SetQueue(CreateObject<DropTailQueue<Packet>>());
  }

  //Add forwarding table
  sw1->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {n2_net1});
  sw2->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {n3_net1});

  sw1->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 101, {n2_net1});

  //Stream Indentification + FRER
  //On aggregated interface of ESsource
  //Stream identification
  Ptr<NullStreamIdentificationFunction> sif0 = CreateObject<NullStreamIdentificationFunction>();
  uint16_t StreamHandle = 1;
  sif0->SetAttribute("VlanID", UintegerValue(100));
  sif0->SetAttribute("Address", AddressValue(Mac48Address("ff:ff:ff:ff:ff:ff")));
  n0->AddStreamIdentificationFunction(StreamHandle, sif0, {}, {}, {n0_netagg}, {});
  //Sequencing : Sequence generation
  Ptr<SequenceGenerationFunction> seqf0 = CreateObject<SequenceGenerationFunction>();
  seqf0->SetAttribute("Direction", BooleanValue(true));   //out-facing
  seqf0->SetStreamHandle({StreamHandle});
  n0->AddSequenceGenerationFunction(seqf0);
  //Sequence encode
  Ptr<SequenceEncodeDecodeFunction> seqEnc0 = CreateObject<SequenceEncodeDecodeFunction>();
  seqEnc0->SetAttribute("Direction", BooleanValue(true)); //out-facing
  seqEnc0->SetAttribute("Active", BooleanValue(true));
  seqEnc0->SetStreamHandle({StreamHandle});
  seqEnc0->SetPort(n0_netagg);
  n0->AddSequenceEncodeDecodeFunction(seqEnc0);

  // On aggregated interface of ESdest
  //Stream identification
  Ptr<NullStreamIdentificationFunction> sif1 = CreateObject<NullStreamIdentificationFunction>();
  StreamHandle = 1;
  sif1->SetAttribute("VlanID", UintegerValue(100));
  sif1->SetAttribute("Address", AddressValue(Mac48Address("ff:ff:ff:ff:ff:ff")));
  n1->AddStreamIdentificationFunction(StreamHandle, sif1, {n1_netagg}, {}, {}, {});
  //Sequence Decode
  Ptr<SequenceEncodeDecodeFunction> seqEnc1 = CreateObject<SequenceEncodeDecodeFunction>();
  seqEnc1->SetAttribute("Direction", BooleanValue(true)); //out-facing
  seqEnc1->SetAttribute("Active", BooleanValue(false));
  seqEnc1->SetStreamHandle({StreamHandle});
  seqEnc1->SetPort(n1_netagg);
  n1->AddSequenceEncodeDecodeFunction(seqEnc1);
  //Sequencing : Sequence recovery
  Ptr<SequenceRecoveryFunction> seqfreco0 = CreateObject<SequenceRecoveryFunction>();
  seqfreco0->SetAttribute("Direction", BooleanValue(true)); //out-facing
  seqfreco0->SetAttribute("TakeNoSequence", BooleanValue(false));
  seqfreco0->SetAttribute("IndividualRecovery", BooleanValue(false));
  seqfreco0->SetStreamHandle({StreamHandle});
  seqfreco0->SetPorts({n1_netagg});
  n1->AddSequenceRecoveryFunction(seqfreco0);
  //Sequencing : Sequence recovery : recovery function
  Ptr<MatchRecoveryFunction> recf0 = CreateObject<MatchRecoveryFunction>();
  seqfreco0->SetRecoveryFunction(recf0);
  //Sequencing : Sequence recovery : latent error detection function
  Ptr<LatentErrorDetectionFunction> latf0 = CreateObject<LatentErrorDetectionFunction>();
  latf0->SetAttribute("LatentErrorPaths", UintegerValue(2));
  latf0->SetRecoveryFunction(recf0);
  seqfreco0->SetLatentErrorDetectionFunction(latf0);

  //Application description
  Ptr<EthernetGenerator> app0 = CreateObject<EthernetGenerator>();
  app0->Setup(n0_netagg);
  app0->SetAttribute("BurstSize", UintegerValue(1));
  app0->SetAttribute("PayloadSize", UintegerValue(1400));
  app0->SetAttribute("Period", TimeValue(MilliSeconds(10)));
  app0->SetAttribute("PCP", UintegerValue(1));
  app0->SetAttribute("VlanID", UintegerValue(100));
  n0->AddApplication(app0);
  app0->SetStartTime(MilliSeconds(0));
  app0->SetStopTime(MilliSeconds(20));

  Ptr<EthernetGenerator> app1 = CreateObject<EthernetGenerator>();
  app1->Setup(n0_net0);
  app1->SetAttribute("BurstSize", UintegerValue(1));
  app1->SetAttribute("PayloadSize", UintegerValue(1400));
  app1->SetAttribute("Period", TimeValue(MilliSeconds(20)));
  app1->SetAttribute("PCP", UintegerValue(0));
  app1->SetAttribute("VlanID", UintegerValue(101));
  n0->AddApplication(app1);
  app1->SetStartTime(Seconds(0));
  app1->SetStopTime(MilliSeconds(41));

  //Callback to display the packet transmitted and received log
  Packet::EnablePrinting();
  n0_net0->TraceConnectWithoutContext("PhyTxBegin", MakeBoundCallback(&PhyTxCallback, Names::FindName(n0) + ":" + Names::FindName(n0_net0)));
  n0_net1->TraceConnectWithoutContext("PhyTxBegin", MakeBoundCallback(&PhyTxCallback, Names::FindName(n0) + ":" + Names::FindName(n0_net1)));
  n1_netagg->TraceConnectWithoutContext("MacRx", MakeBoundCallback(&MacRxCallback, Names::FindName(n1) + ":" + Names::FindName(n1_netagg)));

  //Callback to display the packet elimination on last output port
  n1_netagg->TraceConnectWithoutContext("FrerDrop", MakeBoundCallback(&FrerDrop, Names::FindName(n1_netagg)));

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