#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/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"

#include "ns3/animation-trace.h"

/**
 * \file
 *
 * Example with 2ES connected 4SW in a 1Gb/s full duplex link to demonstrate
 * FRER usage. To be more specific this example use the SW1 to do the replication
 * using mulitcast and SW4 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.
 *
 *                  / ==== SW2 ==== \
 *  ESsource ==== SW1               SW4 ==== ESdest
 *                  \ ==== SW3 ==== /
 *
 * This example generate a NetAnim trace for graphical purpose.
 *
 */


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

    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);
    Ptr<TsnNode> n4 = CreateObject<TsnNode>();
    Names::Add("SW3", n4);
    Ptr<TsnNode> n5 = CreateObject<TsnNode>();
    Names::Add("SW4", n5);

    //Create and add a netDevices to each node
    Ptr<TsnNetDevice> net0 = CreateObject<TsnNetDevice>();
    n0->AddDevice(net0);
    Names::Add("ESsource#01", net0);
    Ptr<TsnNetDevice> net1 = CreateObject<TsnNetDevice>();
    n1->AddDevice(net1);
    Names::Add("ESdest#01", net1);

    Ptr<TsnNetDevice> net2_1 = CreateObject<TsnNetDevice>();
    n2->AddDevice(net2_1);
    Ptr<TsnNetDevice> net2_2 = CreateObject<TsnNetDevice>();
    n2->AddDevice(net2_2);
    Ptr<TsnNetDevice> net2_3 = CreateObject<TsnNetDevice>();
    n2->AddDevice(net2_3);

    Ptr<TsnNetDevice> net3_1 = CreateObject<TsnNetDevice>();
    n3->AddDevice(net3_1);
    Ptr<TsnNetDevice> net3_2 = CreateObject<TsnNetDevice>();
    n3->AddDevice(net3_2);

    Ptr<TsnNetDevice> net4_1 = CreateObject<TsnNetDevice>();
    n4->AddDevice(net4_1);
    Ptr<TsnNetDevice> net4_2 = CreateObject<TsnNetDevice>();
    n4->AddDevice(net4_2);

    Ptr<TsnNetDevice> net5_1 = CreateObject<TsnNetDevice>();
    n5->AddDevice(net5_1);
    Ptr<TsnNetDevice> net5_2 = CreateObject<TsnNetDevice>();
    n5->AddDevice(net5_2);
    Ptr<TsnNetDevice> net5_3 = CreateObject<TsnNetDevice>();
    n5->AddDevice(net5_3);

    //Create Ethernet Channels and attach it to the netDevices
    Ptr<EthernetChannel> l0 = CreateObject<EthernetChannel>();
    net0->Attach(l0);
    net2_1->Attach(l0);
    Ptr<EthernetChannel> l1 = CreateObject<EthernetChannel>();
    net2_2->Attach(l1);
    net3_1->Attach(l1);
    Ptr<EthernetChannel> l2 = CreateObject<EthernetChannel>();
    net2_3->Attach(l2);
    net4_1->Attach(l2);
    Ptr<EthernetChannel> l3 = CreateObject<EthernetChannel>();
    net3_2->Attach(l3);
    net5_1->Attach(l3);
    Ptr<EthernetChannel> l4 = CreateObject<EthernetChannel>();
    net4_2->Attach(l4);
    net5_2->Attach(l4);
    Ptr<EthernetChannel> l5 = CreateObject<EthernetChannel>();
    net5_3->Attach(l5);
    net1->Attach(l5);

    //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(net2_1);
    sw1->AddSwitchPort(net2_2);
    sw1->AddSwitchPort(net2_3);
    Ptr<SwitchNetDevice> sw2 = CreateObject<SwitchNetDevice>();
    sw2->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
    sw2->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
    n3->AddDevice(sw2);
    sw2->AddSwitchPort(net3_1);
    sw2->AddSwitchPort(net3_2);
    Ptr<SwitchNetDevice> sw3 = CreateObject<SwitchNetDevice>();
    sw3->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
    sw3->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
    n4->AddDevice(sw3);
    sw3->AddSwitchPort(net4_1);
    sw3->AddSwitchPort(net4_2);
    Ptr<SwitchNetDevice> sw4 = CreateObject<SwitchNetDevice>();
    sw4->SetAttribute("MinForwardingLatency", TimeValue(MicroSeconds(10)));
    sw4->SetAttribute("MaxForwardingLatency", TimeValue(MicroSeconds(10)));
    n5->AddDevice(sw4);
    sw4->AddSwitchPort(net5_1);
    sw4->AddSwitchPort(net5_2);
    sw4->AddSwitchPort(net5_3);

    //Allocate a Mac address
    net0->SetAddress(Mac48Address::Allocate());
    net1->SetAddress(Mac48Address::Allocate());
    sw1->SetAddress(Mac48Address::Allocate());
    sw2->SetAddress(Mac48Address::Allocate());
    sw3->SetAddress(Mac48Address::Allocate());
    sw4->SetAddress(Mac48Address::Allocate());

    //Create and add eight FIFO on each net device
    for (int i=0; i<8; i++)
    {
      net0->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net1->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net2_1->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net2_2->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net2_3->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net3_1->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net3_2->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net4_1->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net4_2->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net5_1->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net5_2->SetQueue(CreateObject<DropTailQueue<Packet>>());
      net5_3->SetQueue(CreateObject<DropTailQueue<Packet>>());
    }

    //Add forwarding table
    sw1->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {net2_2, net2_3});
    sw2->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {net3_2});
    sw3->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {net4_2});
    sw4->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 100, {net5_3});

    sw1->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 101, {net2_2});
    sw2->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 101, {net3_2});
    sw4->AddForwardingTableEntry(Mac48Address("ff:ff:ff:ff:ff:ff"), 101, {net5_3});

    //Stream Indentification + FRER
    //First switch
    //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")));
    n2->AddStreamIdentificationFunction(StreamHandle, sif0, {net2_1}, {}, {}, {});
    //Sequencing : Sequence generation
    Ptr<SequenceGenerationFunction> seqf0 = CreateObject<SequenceGenerationFunction>();
    seqf0->SetAttribute("Direction", BooleanValue(false));   //in-facing
    seqf0->SetStreamHandle({StreamHandle});
    n2->AddSequenceGenerationFunction(seqf0);
    //Sequence encode
    Ptr<SequenceEncodeDecodeFunction> seqEnc0 = CreateObject<SequenceEncodeDecodeFunction>();
    seqEnc0->SetAttribute("Direction", BooleanValue(false)); //in-facing
    seqEnc0->SetAttribute("Active", BooleanValue(true));
    seqEnc0->SetStreamHandle({StreamHandle});
    seqEnc0->SetPort(net2_1);
    n2->AddSequenceEncodeDecodeFunction(seqEnc0);

    //Last switch
    //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")));
    n5->AddStreamIdentificationFunction(StreamHandle, sif1, {}, {}, {net5_3}, {});
    //Sequence Decode
    Ptr<SequenceEncodeDecodeFunction> seqEnc1 = CreateObject<SequenceEncodeDecodeFunction>();
    seqEnc1->SetAttribute("Direction", BooleanValue(false)); //in-facing
    seqEnc1->SetAttribute("Active", BooleanValue(false));
    seqEnc1->SetStreamHandle({StreamHandle});
    seqEnc1->SetPort(net5_3);
    n5->AddSequenceEncodeDecodeFunction(seqEnc1);
    //Sequencing : Sequence recovery
    Ptr<SequenceRecoveryFunction> seqfreco0 = CreateObject<SequenceRecoveryFunction>();
    seqfreco0->SetAttribute("Direction", BooleanValue(false)); //in-facing
    seqfreco0->SetAttribute("TakeNoSequence", BooleanValue(false));
    seqfreco0->SetAttribute("IndividualRecovery", BooleanValue(false));
    seqfreco0->SetStreamHandle({StreamHandle});
    seqfreco0->SetPorts({net5_3});
    n5->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(net0);
    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(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();
    net0->TraceConnectWithoutContext("PhyTxBegin", MakeBoundCallback(&PhyTxCallback, Names::FindName(n0) + ":" + Names::FindName(net0)));
    net1->TraceConnectWithoutContext("MacRx", MakeBoundCallback(&MacRxCallback, Names::FindName(n1) + ":" + Names::FindName(net1)));

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

    //NetAnim
    // Create the animation trace object and configure for specified output
    AnimationTrace anim = AnimationTrace("example-animation.xml");
    anim.SetNodePosition(n0, 0, 0);
    anim.SetNodePosition(n1, 40, 0);
    anim.SetNodePosition(n2, 10, 0);
    anim.SetNodePosition(n3, 20, -10);
    anim.SetNodePosition(n4, 20, 10);
    anim.SetNodePosition(n5, 30, 0);
    anim.StartAnimation();
    Packet::EnablePrinting(); //Need to enable packet printing for metadata display on NetAnim

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