#include "ns3/core-module.h"
#include "ns3/node-module.h"
#include "ns3/simulator-module.h"
#include "ns3/contrib-module.h"
#include "ns3/point-to-point-net-device.h"
#include "ns3/point-to-point-channel.h"
#include "ns3/packet-socket-helper.h"

#include <string>
#include <ostream>
#include <iostream>
#include <list>
#include <iomanip>
#include <sstream>


using namespace  ns3;

#define USE_HANDLESEND
//------------------------------------------------------------------------------

/*
 * returns a string containing the length of the buffer,
 * followed by an run length encoded version of the buffer
 * un-printables are represented by '.' (i.e. special_marker
 * const).
 */
std::string rle(const uint8_t *buf, int buflen) 
{
  int justify = 5;
  std::ostringstream o;
  const uint8_t min_rep = 3;
  uint8_t cur_char = 0;
  const uint8_t special_marker = '.';
  int cur_cnt = 0;
  if (justify) 
    {
      o << std::setw(justify) << std::right << buflen << ":";
    } 
  else 
    {
      o << buflen << ":";
    }

  if (buf) {
    for (int i = 0; i < buflen; i++) 
      {
        if ((!isprint(buf[i]) && cur_char == special_marker) ||
            (buf[i] == cur_char))
          {
            cur_cnt++;
          }
        else 
          {
            if (cur_cnt >= min_rep)
              {
                o << (isprint(cur_char) ? static_cast<char>(cur_char) : '.') << "{"
                  << cur_cnt << "}";
              }
            else
              {
                for (int irep=0; irep < cur_cnt; irep++)
                  o << (isprint(cur_char) ? static_cast<char>(cur_char) : '.');
              }
            if (isprint(buf[i]))
              {
                cur_char = buf[i];
              }
            else
              {
                cur_char = special_marker;
              }
            cur_cnt = 1;
          }
      }
    if (cur_cnt >= min_rep)
      {
        o << (isprint(cur_char) ? static_cast<char>(cur_char) : '.') << "{"
          << cur_cnt << "}";
      }
    else
      {
        for (int irep=0; irep < cur_cnt; irep++)
          {
            o << (isprint(cur_char) ? static_cast<char>(cur_char) : '.');
          }
      }
  }
  return o.str();
}


//------------------------------------------------------------------------------

void pretty_packet(std::ostream &o,
		   const std::string &node_label,
		   const std::string &direction_label,
		   const Ptr<Packet> &packet) {
  packet->Print(o 
		<< Simulator::Now() << ": " 
		<< node_label
		<< ":" << direction_label << ":");
  o << rle(packet->PeekData(),packet->GetSize())
    << std::endl;
}
void pretty_packet(std::ostream &o,
		   const std::string &node_label,
		   const std::string &direction_label,
		   const Ptr<Packet> &packet,
		   const uint32_t bytes_left) {
  packet->Print(o 
		<< Simulator::Now() << ": " 
		<< node_label
		<< ":" << direction_label << ":");
  o << " left{" << bytes_left << "} ";
  o << rle(packet->PeekData(),packet->GetSize())
    << std::endl;
}
//------------------------------------------------------------------------------
class SrcRtSwitch : public Application
{
public:
  explicit SrcRtSwitch(const char *lab_);
  uint32_t AddChannel(Ptr<PointToPointChannel> &channel);
  void InjectPacket (const Time & dt,
		     const char fill, const uint32_t size) const;
  mutable uint32_t m_bytes_left;
protected:
  void DoDispose (void);
  void StartApplication();
  void StopApplication();

private:
  void HandleRead(Ptr<Socket>);
  void HandleSend (Ptr<Socket> socket, uint32_t avail) const;
  void SendPacket (Ptr<Packet> packet) const;
  std::string m_label;
  Ptr<Node> m_node;
  Ptr<Socket>  m_socket;
  mutable std::list<Ptr<Packet> >         m_txQueue;
  const static uint16_t FragSize = 50000;

};

void
SrcRtSwitch::DoDispose(void)
{
  Application::DoDispose ();
}

void
SrcRtSwitch::StartApplication()
{
  std::cout << "SrcRtSwitch::StartApplication called"
	    << std::endl;
}

void
SrcRtSwitch::StopApplication()
{
  std::cout << "SrcRtSwitch::StopApplication called"
	    << std::endl;
}

SrcRtSwitch::SrcRtSwitch(const char * lab_) :  m_bytes_left(0), m_label(lab_)
{
  m_node = CreateObject<Node> ();
  this->Start(Seconds(0.0));
  PacketSocketHelper sockhelper;
  sockhelper.Install(m_node);
}

uint32_t
SrcRtSwitch::AddChannel(Ptr<PointToPointChannel> &channel)
{
  Ptr<PointToPointNetDevice> netdev =
    CreateObject<PointToPointNetDevice> ();
  netdev->Attach (channel);
  netdev->SetAddress (Mac48Address::Allocate ());
  netdev->SetQueue (CreateObject<DropTailQueue> ());
  netdev->SetDataRate(DataRate (100000000)); // to make output prettier

  uint16_t max_mtu = FragSize;
  std::cout << "setting max_mtu = " << max_mtu << std::endl;
  if (!netdev->SetMtu(max_mtu)) {
    std::cout << "Failed to set MTU to << " << max_mtu
	      << "I'm outta here." << std::endl;
    exit(1);
  }
  uint32_t rtn = m_node->AddDevice(netdev);
  m_socket = Socket::CreateSocket (m_node,
				   PacketSocketFactory::GetTypeId ());

  PacketSocketAddress psa;
  psa.SetSingleDevice(m_node->GetDevice(0)->GetIfIndex ());
  psa.SetPhysicalAddress (m_node->GetDevice(0)->GetAddress());
  psa.SetProtocol (0x800);
  m_socket->Bind(psa);

  psa.SetPhysicalAddress (m_node->GetDevice(0)->GetBroadcast());
  m_socket->Connect(psa);

  m_socket->
    SetRecvCallback (MakeCallback (&SrcRtSwitch::HandleRead, this));
#ifdef USE_HANDLESEND
  m_socket->
    SetSendCallback (MakeCallback (&SrcRtSwitch::HandleSend, this));
#endif
  return rtn;
}

/** really sends the packet when enabled.
 */
#ifdef USE_HANDLESEND
void 
SrcRtSwitch::HandleSend (Ptr<Socket> socket, uint32_t avail) const
{
  std::cout 
    << Simulator::Now() << ": " << m_label << ": "
    << "HandleSend(" << this << ","
    << socket << "," << avail << ")"
    << std::endl;

  NS_ASSERT (avail > 0);
  if (m_txQueue.empty ()) {
    std::cout << "  queue is empty"
	      << std::endl;
    return;
  }
  std::cout << "HandleSend: queue size " << m_txQueue.size()
	    << std::endl;
  Ptr<Packet> p = m_txQueue.front ();
  m_txQueue.pop_front ();
  if (p->GetSize () <= avail) {
    int r = socket->Send (p);
    std::cout << "Send returned " << r << std::endl;
  } else
    std::cout << "HandleSend() needs to fragment, but I can't."
	      << std::endl;
}
#endif

/** On read we print what we see
 */
void
SrcRtSwitch::HandleRead( Ptr<Socket> socket)
{
  std::cout 
    << Simulator::Now() << ": "
    << m_label << ": "
    << "HandleRead("
    << this << ","
    << socket
    << ")"
    << std::endl;
  Ptr<Packet> packet = socket->Recv();
  m_bytes_left -= packet->GetSize();
  pretty_packet(std::cout ,m_label,"recv",packet,m_bytes_left);
  if (m_bytes_left == 0)
    {
      std::cout << Simulator::Now() << ": "
		<< m_label << ": " << "Packet Delivered"
		<< std::endl;
    }
}

/** Sends packets, it only queues a packet for sending if USE_HANDLESEND
    is defined.
 */
void
SrcRtSwitch::SendPacket (Ptr<Packet> packet) const
{
  std::cout
    << Simulator::Now() << ": "
    << m_label << ": "
    << "SendPacket("
    << this << ","
    << packet << ")"
    << std::endl;
  pretty_packet(std::cout,m_label,"forwarding",packet);
#ifdef USE_HANDLESEND
  m_txQueue.push_back (packet);
  if (m_socket->GetTxAvailable () >= packet->GetSize())
    {
      std::cout << "SendPacket(): space avail...." << std::endl;
      HandleSend (m_socket, m_socket->GetTxAvailable ());
    }
  else
    std::cout << "SendPacket(): No space to send." << std::endl;
#else
  std::cout << "m_socket{" << m_socket << "}->Send(packet): sock= "
	    << m_socket
	    << std::endl;
  int r = m_socket->Send(packet);
  std::cout << "Send returned " << r << std::endl;
#endif
}

/** User injects packets into the system with this call. If necessary this
    method will fragment the packet into FragSize chunks. Fragmentation is
    weird because the real code needs headers which are removed here.
 */
void
SrcRtSwitch::InjectPacket (const Time &dt, // time to inject
			   const char fill, // char to fill buffer with
			   const uint32_t size // size of buffer to send
			   ) const
{
  std::cout
    << Simulator::Now() << ": "
    << m_label << ": "
    << "InjectPacket("
    << this << ","
    << dt << ","
    << fill << ","
    << size    
    << ")"
    << std::endl;

  m_bytes_left = size;
  uint32_t data_left = size;
  uint32_t frag_counter = 0;
  while(data_left > 0) {
    uint32_t frag_data_size = data_left;
    if (frag_data_size > FragSize) {
      frag_data_size = FragSize;
    }
    data_left -= frag_data_size;
    uint8_t *buf = (uint8_t *) malloc(frag_data_size*sizeof(uint8_t));
    memset(buf,fill,frag_data_size);
    Ptr<Packet> packet = Create<Packet> (buf,frag_data_size);
    free(buf);
    pretty_packet(std::cout,m_label,"inject frag",packet);
    Simulator::Schedule(
			dt,
			&SrcRtSwitch::SendPacket,this,
			packet);
    frag_counter++;
  }
}
//------------------------------------------------------------------------------
int
main(int argc, char** argv)
{
  Packet::EnablePrinting();

  uint32_t size = 10000000;
  CommandLine cmd;
  cmd.AddValue("size",
               "size of packet to send",
               size);
  cmd.Parse(argc,argv);
  SrcRtSwitch node1("node1");
  SrcRtSwitch node2("node2");


  Ptr<PointToPointChannel> chan0 = CreateObject<PointToPointChannel> ();
  node1.AddChannel(chan0);
  node2.AddChannel(chan0);



    ;
  node2.m_bytes_left = size;
  node1.InjectPacket(Seconds(1.0),'A',size);

  std::cout << "\n" << "Simulation Start." << std::endl;
  Simulator::Run ();
  Simulator::Destroy ();
}