// Include a header file from your module to test.
#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"

// An essential include is test.h
#include "ns3/test.h"
#include "ns3/core-module.h"
#include "ns3/drop-tail-queue.h"
#include "ns3/timestamp-tag.h"

// Do not put your test classes in namespace ns3.  You may find it useful
// to use the using directive to access the ns3 namespace directly
using namespace ns3;
NS_LOG_COMPONENT_DEFINE("FRERTestSuite");
// Add a doxygen group for tests.
// If you have more than one test, this should be in only one of them.
/**
 * \defgroup FRER-tests Tests for FRER
 * \ingroup tsn
 * \ingroup tests
 */




/**
 * \ingroup FRER-tests
 * Check if paquets are correctly replicated and eliminated on a 4 SW / 2 ES network
 *
 *                  / ==== SW2 ==== \
 *  ESsource ==== SW1               SW4 ==== ESdest
 *                  \ ==== SW3 ==== /
 *
 * The FRER use in this test 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.
 *
 *
 */
class FRERBasicMulticastTestCase : public TestCase
{
  public:
    FRERBasicMulticastTestCase();
    virtual ~FRERBasicMulticastTestCase();

  private:
    void DoRun() override;
    void SendTx(Ptr<const Packet> p);
    void ReceiveRx(Ptr<const Packet> p);
    uint64_t m_replicated{0};          //!< number of bytes sent
    uint64_t m_received{0};      //!< number of bytes received
    uint64_t m_pktSize = 500;
    uint64_t m_expected_replicated = (1400 + 22 + 6) * 2 + (500+22);            //the 1400Bytes pkt is replicated two time and RTAGed (+6bits), the 500Bytes pkt is not replicated
    uint64_t m_expected_receive = (1400 + 22) + (500+22);
};

// Add some help text to this case to describe what it is intended to test
FRERBasicMulticastTestCase::FRERBasicMulticastTestCase()
    : TestCase("Check if paquets are correctly replicated and eliminated")
{}

// This destructor does nothing but we include it as a reminder that
// the test case should clean up after itself
FRERBasicMulticastTestCase::~FRERBasicMulticastTestCase()
{
}

void
FRERBasicMulticastTestCase::SendTx(Ptr<const Packet> p)
{
  m_replicated  += p->GetSize();
}

void
FRERBasicMulticastTestCase::ReceiveRx(Ptr<const Packet> p)
{
  m_received += p->GetSize();
}

//
// This method is the pure virtual method from class TestCase that every
// TestCase must implement
//
void
FRERBasicMulticastTestCase::DoRun()
{

  //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>();
  l0->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  net0->Attach(l0);
  net2_1->Attach(l0);
  Ptr<EthernetChannel> l1 = CreateObject<EthernetChannel>();
  l1->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  net2_2->Attach(l1);
  net3_1->Attach(l1);
  Ptr<EthernetChannel> l2 = CreateObject<EthernetChannel>();
  l2->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  net2_3->Attach(l2);
  net4_1->Attach(l2);
  Ptr<EthernetChannel> l3 = CreateObject<EthernetChannel>();
  l3->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  net3_2->Attach(l3);
  net5_1->Attach(l3);
  Ptr<EthernetChannel> l4 = CreateObject<EthernetChannel>();
  l4->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  net4_2->Attach(l4);
  net5_2->Attach(l4);
  Ptr<EthernetChannel> l5 = CreateObject<EthernetChannel>();
  l5->SetAttribute("Delay", TimeValue(MicroSeconds(0)));
  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(20)));
  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(500));
  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(20));


  //Callback to trace the message being replicated and eliminated
  net2_2->TraceConnectWithoutContext("MacTx",
                                     MakeCallback(&FRERBasicMulticastTestCase::SendTx, this));
  net2_3->TraceConnectWithoutContext("MacTx",
                                     MakeCallback(&FRERBasicMulticastTestCase::SendTx, this));
  net1->TraceConnectWithoutContext("MacRx",
                                   MakeCallback(&FRERBasicMulticastTestCase::ReceiveRx, this));

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

  NS_TEST_ASSERT_MSG_EQ(m_replicated, m_expected_replicated, "All Packets have been replicated");
  NS_TEST_ASSERT_MSG_EQ(m_received, m_expected_receive, "All Packets have been eliminated");
}


// The TestSuite class names the TestSuite, identifies what type of TestSuite,
// and enables the TestCases to be run.  Typically, only the constructor for
// this class must be defined

/**
 * \ingroup FRER-tests
 * TestSuite for module FRER
 */
class FRERTestSuite : public TestSuite
{
  public:
    FRERTestSuite();
};

FRERTestSuite::FRERTestSuite()
    : TestSuite("frer", UNIT)
{
    LogComponentEnable("FRERTestSuite", LOG_LEVEL_ALL);

    //Pkt in a switched network with multicast FRER
    AddTestCase(new FRERBasicMulticastTestCase(), TestCase::QUICK);


}

// Do not forget to allocate an instance of this TestSuite
/**
 * \ingroup FRER-tests
 * Static variable for test initialization
 */
static FRERTestSuite m_FRERTestSuite;