#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/trace-helper.h"
#include <fstream>

#include "ns3/ethernet-net-device.h"
#include "ns3/ethernet-channel.h"
#include "ns3/ethernet-generator.h"
#include "ns3/ethernet-header2.h"
#include "ns3/switch-net-device.h"

/**
 * \file
 *
 * Example of the use of switch-net-device.cc switch-channel.cc
 * ethernet-net-device.cc ethernet-channel.cc on a network composed of three
 * end-stations connected with a 1Gb/s full duplex link
 *  ES1 ==== SW ==== ES2
 *            ==== ES3
 */

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("Example");


//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 write in a file the pkt latency hop by hop
static void
LatencyHopByHopCallback(Ptr<OutputStreamWrapper> stream, std::string context, Ptr<const Packet> p)
{
	Ptr<Packet> originalPacket = p->Copy();
	EthernetHeader2 ethHeader;
	originalPacket->RemoveHeader(ethHeader);

	*stream->GetStream() << context << " : Packet #"<< p->GetUid() << " from "<< ethHeader.GetSrc() << "(VID:" << ethHeader.GetVid() << ") : ";
	TimestampTag tag;
	TypeId tid = tag.GetInstanceTypeId();
	ByteTagIterator i = p->GetByteTagIterator();
	while (i.HasNext())
	{
		ByteTagIterator::Item item = i.Next();
		if (tid == item.GetTypeId())
		{
			item.GetTag(tag);
			*stream->GetStream() << tag.GetTimestamp().GetNanoSeconds() << ",";
		}
	}
	*stream->GetStream() << "\n";
}


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 four nodes
    Ptr<Node> n0 = CreateObject<Node>();
    Names::Add("ES1", n0);
    Ptr<Node> n1 = CreateObject<Node>();
    Names::Add("ES2", n1);
    Ptr<Node> n2 = CreateObject<Node>();
    Names::Add("ES3", n2);
    Ptr<Node> n3 = CreateObject<Node>();
    Names::Add("SW", n3);

    //Create and add a netDevice to each end station node
    Ptr<EthernetNetDevice> net0 = CreateObject<EthernetNetDevice>();
    n0->AddDevice(net0);
    Names::Add("ES1#01", net0);
    Ptr<EthernetNetDevice> net1 = CreateObject<EthernetNetDevice>();
    n1->AddDevice(net1);
    Names::Add("ES2#01", net1);
    Ptr<EthernetNetDevice> net2 = CreateObject<EthernetNetDevice>();
    n2->AddDevice(net2);
    Names::Add("ES3#01", net2);

    //Create and add a netDevice to each switch port
    Ptr<EthernetNetDevice> swnet0 = CreateObject<EthernetNetDevice>();
    n3->AddDevice(swnet0);
    Names::Add("SW#01", swnet0);
    Ptr<EthernetNetDevice> swnet1 = CreateObject<EthernetNetDevice>();
    n3->AddDevice(swnet1);
    Names::Add("SW#02", swnet1);
    Ptr<EthernetNetDevice> swnet2 = CreateObject<EthernetNetDevice>();
    n3->AddDevice(swnet2);
    Names::Add("SW#03", swnet2);

    //Create Ethernet Channels and connect switch to the end-stations
    Ptr<EthernetChannel> channel0 = CreateObject<EthernetChannel>();
    net0->Attach(channel0);
    swnet0->Attach(channel0);
    Ptr<EthernetChannel> channel1 = CreateObject<EthernetChannel>();
    net1->Attach(channel1);
    swnet1->Attach(channel1);
    Ptr<EthernetChannel> channel2 = CreateObject<EthernetChannel>();
    net2->Attach(channel2);
    swnet2->Attach(channel2);

    //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)));
    n3->AddDevice(sw);
    sw->AddSwitchPort(swnet0);
    sw->AddSwitchPort(swnet1);
    sw->AddSwitchPort(swnet2);

    //Allocate a Mac address and create 2 FIFOs (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>>());
    sw->SetAddress(Mac48Address::Allocate());
    swnet0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    swnet1->SetQueue(CreateObject<DropTailQueue<Packet>>());
    swnet2->SetQueue(CreateObject<DropTailQueue<Packet>>());
    swnet0->SetQueue(CreateObject<DropTailQueue<Packet>>());
    swnet1->SetQueue(CreateObject<DropTailQueue<Packet>>());
    swnet2->SetQueue(CreateObject<DropTailQueue<Packet>>());

    //Add forwarding table
    sw->AddForwardingTableEntry(Mac48Address::ConvertFrom(net2->GetAddress()), 1, {swnet2});

    //Application description
    //ES1 -> ES3 with priority 1
    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(1));
    app0->SetAttribute("PCP", UintegerValue(1));
    n0->AddApplication(app0);
    app0->SetStartTime(Seconds(0));
    app0->SetStopTime(Seconds(10));

    //Callback to display the packet latency log
    std::string context = Names::FindName(n2) + ":" + Names::FindName(net2);
    net2->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyCallback, context));

    //Callback to write the packet latency hop by hop in a file
    AsciiTraceHelper asciiTraceHelper;
    std::string filename = "example.hopByHopLat";
  	Ptr<OutputStreamWrapper> stream = asciiTraceHelper.CreateFileStream(filename);
    net2->TraceConnectWithoutContext("Latency", MakeBoundCallback(&LatencyHopByHopCallback, stream, context));

    //Execute the simulation
    Simulator::Stop(Seconds(10));
    Simulator::Run();
    Simulator::Destroy();
    return 0;
}