// Include a header file from your module to test.
#include "ns3/tsn-net-device.h"
#include "ns3/cbs.h"
#include "ns3/ethernet-channel.h"
#include "ns3/ethernet-header2.h"
#include "ns3/ethernet-generator.h"
#include "ns3/switch-net-device.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("CbsTestSuite");
// Add a doxygen group for tests.
// If you have more than one test, this should be in only one of them.
/**
 * \defgroup cbs-tests Tests for cbs
 * \ingroup tsn
 * \ingroup tests
 */




/**
 * \ingroup Cbs-tests
 * Check if message crossed a point to point tsn channel
 */
class CbsBasicTestCase : public TestCase
{
  public:
    CbsBasicTestCase();
    virtual ~CbsBasicTestCase();

  private:
    void DoRun() override;
    void SendTx(Ptr<const Packet> p);
    void ReceiveRx(Ptr<const Packet> p);
    uint64_t m_sent{0};          //!< number of bytes sent
    uint64_t m_received{0};      //!< number of bytes received
};

// Add some help text to this case to describe what it is intended to test
CbsBasicTestCase::CbsBasicTestCase()
    : TestCase("Check if paquets cross a point to point ethernet channel, tsn net device and a CBS")
{
}

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

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

void
CbsBasicTestCase::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
CbsBasicTestCase::DoRun()
{
  //Create two nodes
  Ptr<TsnNode> n0 = CreateObject<TsnNode>();
  Ptr<TsnNode> n1 = CreateObject<TsnNode>();

  //Create and add a netDevice to each node
  Ptr<TsnNetDevice> net0 = CreateObject<TsnNetDevice>();
  n0->AddDevice(net0);
  Ptr<TsnNetDevice> net1 = CreateObject<TsnNetDevice>();
  n1->AddDevice(net1);

  //Create a Tsn Channel and attach it two the two netDevices
  Ptr<EthernetChannel> channel = CreateObject<EthernetChannel>();
  net0->Attach(channel);
  net1->Attach(channel);

  //Allocate a Mac address for each netDevice.
  net0->SetAddress(Mac48Address::Allocate());
  net1->SetAddress(Mac48Address::Allocate());

  //Create CBS and add FIFOs
  Ptr<Cbs> cbs = CreateObject<Cbs>();
  cbs->SetTsnNetDevice(net0);
  cbs->SetAttribute("IdleSlope", DataRateValue(DataRate("5Mb/s")));
  cbs->SetAttribute("portTransmitRate", DataRateValue(DataRate("1Gb/s")));
  net0->SetQueue(CreateObject<DropTailQueue<Packet>>(), cbs);
  net1->SetQueue(CreateObject<DropTailQueue<Packet>>());

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


  //Callback to trace the message being send and received
  net0->TraceConnectWithoutContext("MacTx",
                                     MakeCallback(&CbsBasicTestCase::SendTx, this));
  net1->TraceConnectWithoutContext("MacRx",
                                   MakeCallback(&CbsBasicTestCase::ReceiveRx, this));

  //Execute the simulation
  Simulator::Stop(Seconds(12));
  Simulator::Run();
  Simulator::Destroy();

  NS_TEST_ASSERT_MSG_EQ(m_sent, 2 * 10 * (1400 + 22), "10 Packets have been sent two times");
  NS_TEST_ASSERT_MSG_EQ(m_sent, m_received, "All Packets sent have been received");
}




/**
 * \ingroup Cbs-tests
 * Check last message latency with a point to point network
 */
class CbsLatencyTestCase : public TestCase
{
  public:
    CbsLatencyTestCase(uint16_t burstSize, Time latency, DataRate datarate);
    virtual ~CbsLatencyTestCase();

  private:
    void DoRun() override;
    void Latency(Ptr<const Packet> p);
    uint16_t m_burstSize;
    Time m_true_latency;
    Time m_latency;
    DataRate m_datarate;
};

CbsLatencyTestCase::CbsLatencyTestCase(uint16_t burstSize, Time latency, DataRate datarate)
    : TestCase("Check last message latency with a point to point network with CBS")
{
  m_burstSize = burstSize;
  m_true_latency = latency;
  m_datarate = datarate;
}

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

void
CbsLatencyTestCase::Latency(Ptr<const Packet> p)
{
  TimestampTag tag;
  if (!p->FindFirstMatchingByteTag(tag))
  {
      return;
  }
  Time arrival = Simulator::Now();
  m_latency = arrival - tag.GetTimestamp();
}

//
// This method is the pure virtual method from class TestCase that every
// TestCase must implement
//
void
CbsLatencyTestCase::DoRun()
{
  //Create two nodes
  Ptr<TsnNode> n0 = CreateObject<TsnNode>();
  Ptr<TsnNode> n1 = CreateObject<TsnNode>();

  //Create and add a netDevice to each node
  Ptr<TsnNetDevice> net0 = CreateObject<TsnNetDevice>();
  net0->SetAttribute("DataRate", DataRateValue(m_datarate));
  n0->AddDevice(net0);
  Ptr<TsnNetDevice> net1 = CreateObject<TsnNetDevice>();
  net1->SetAttribute("DataRate", DataRateValue(m_datarate));
  n1->AddDevice(net1);

  //Create a Tsn Channel and attach it two the two netDevices
  Ptr<EthernetChannel> channel = CreateObject<EthernetChannel>();
  net0->Attach(channel);
  net1->Attach(channel);

  //Allocate a Mac address for each netDevice.
  net0->SetAddress(Mac48Address::Allocate());
  net1->SetAddress(Mac48Address::Allocate());

  //Create CBS and add FIFOs
  Ptr<Cbs> cbs = CreateObject<Cbs>();
  cbs->SetTsnNetDevice(net0);
  cbs->SetAttribute("IdleSlope", DataRateValue(DataRate("5Mb/s")));
  cbs->SetAttribute("portTransmitRate", DataRateValue(m_datarate));
  net0->SetQueue(CreateObject<DropTailQueue<Packet>>(), cbs);
  net1->SetQueue(CreateObject<DropTailQueue<Packet>>());

  //Application description
  Ptr<EthernetGenerator> app0 = CreateObject<EthernetGenerator>();
  app0->Setup(net0);
  app0->SetAttribute("BurstSize", UintegerValue(m_burstSize));
  app0->SetAttribute("PayloadSize", UintegerValue(1400));
  app0->SetAttribute("Period", TimeValue(Seconds(5)));
  app0->SetAttribute("VlanID", UintegerValue(1));
  app0->SetAttribute("PCP", UintegerValue(0));
  n0->AddApplication(app0);
  app0->SetStartTime(Seconds(0));
  app0->SetStopTime(Seconds(1));


  //Callback to trace the latency
  net1->TraceConnectWithoutContext("Latency",
                                   MakeCallback(&CbsLatencyTestCase::Latency, this));

  //Execute the simulation
  Simulator::Stop(Seconds(12));
  Simulator::Run();
  Simulator::Destroy();

  NS_TEST_ASSERT_MSG_EQ(m_latency, m_true_latency, "Packet experience the correct latency");
}




// 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 cbs-tests
 * TestSuite for module cbs
 */
class CbsTestSuite : public TestSuite
{
  public:
    CbsTestSuite();
};

CbsTestSuite::CbsTestSuite()
    : TestSuite("cbs", UNIT)
{
    LogComponentEnable("CbsTestSuite", LOG_LEVEL_ALL);
    LogComponentEnable("Cbs", LOG_LEVEL_ALL);
    //Point to point network with cbs
    AddTestCase(new CbsBasicTestCase, TestCase::QUICK);
    AddTestCase(new CbsLatencyTestCase(1, Time(NanoSeconds((1400 + 22 + 8)*8 +25)), DataRate("1Gb/s")), TestCase::QUICK);
    AddTestCase(new CbsLatencyTestCase(2, Time(NanoSeconds(2307200 + (1400 + 22 + 8)*8 + 25)), DataRate("1Gb/s")), TestCase::QUICK); //2307200 is the time needed to restore the credit after an emission
    AddTestCase(new CbsLatencyTestCase(3, Time(NanoSeconds(2307200*2 + (1400 + 22 + 8)*8 + 25)), DataRate("1Gb/s")), TestCase::QUICK);
    AddTestCase(new CbsLatencyTestCase(10, Time(NanoSeconds(2307200*9 + (1400 + 22 + 8)*8 + 25)), DataRate("1Gb/s")), TestCase::QUICK);
    if(Time::GetResolution()==Time::PS){
      AddTestCase(new CbsLatencyTestCase(2, Time(NanoSeconds(2307200 + (1400 + 22 + 8)*3.2 + 25)), DataRate("2.5Gb/s")), TestCase::QUICK);
      AddTestCase(new CbsLatencyTestCase(2, Time(NanoSeconds(2307200 + (1400 + 22 + 8)*1.6 + 25)), DataRate("5Gb/s")), TestCase::QUICK);
      AddTestCase(new CbsLatencyTestCase(2, Time(NanoSeconds(2307200 + (1400 + 22 + 8)*0.8 + 25)), DataRate("10Gb/s")), TestCase::QUICK);
    }
}

// Do not forget to allocate an instance of this TestSuite
/**
 * \ingroup cbs-tests
 * Static variable for test initialization
 */
static CbsTestSuite m_cbsTestSuite;