/* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
/*
 * Copyright 2008 Sandia Corporation.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation;
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * Sandia is a multiprogram laboratory operated by Sandia
 * Corporation, a Lockheed Martin Company, for the United States
 * Department of Energy's National Nuclear Security Administration
 * under Contract DE-AC04-94AL85000.
 *
 * Author: David Evensky <evensky AT sandia dot gov>
 *
 */

#include <iostream>
#include <fstream>
#include <string>
#include <iomanip>

#include "ns3/core-module.h"
#include "ns3/helper-module.h"
#include "ns3/node-module.h"
#include "ns3/global-route-manager.h"
#include "ns3/simulator-module.h"
#include "ns3/log.h"

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("tcp-tester");


#define USE_TAGS
#define USE_HEADER
//#define VERBOSE_TESTING
//#define I_LIKE_CRASHES

#ifdef USE_TAGS
class Atag : public Tag {
public:
  static TypeId GetTypeId(void);
  virtual TypeId GetInstanceTypeId(void) const;
  virtual uint32_t GetSerializedSize(void) const;
  virtual void Serialize(TagBuffer) const;
  virtual void Deserialize(TagBuffer);

  void Print (std::ostream &os) const;
  void SetIntValue(uint32_t);
  uint32_t GetIntValue(void) const;
  void SetStrValue(const std::string &s);
  std::string GetStrValue(void) const;
private:
  uint32_t int_value;
  std::string str_value;
  const static bool debug = 0;
};


TypeId Atag::GetTypeId(void) {
  static TypeId tid = TypeId("ns3::Atag")
    .AddConstructor<Atag>() // NOTE: required Socket::PrintTags to work????
    ;
    return tid;
}

TypeId Atag::GetInstanceTypeId(void) const {
  return GetTypeId();
}

uint32_t Atag::GetSerializedSize (void) const {
  return 2*sizeof(uint32_t) + str_value.size();
}

void Atag::Serialize (TagBuffer i) const {
  if (debug)
    std::cout << "Atag::Serialize() called" << std::endl;
  if (debug)
    std::cout << " int_value = " << int_value << std::endl;
  i.WriteU32 (int_value);
  uint32_t str_len = str_value.size();
  if (debug)
    std::cout << " str_len = " << str_len << std::endl;
  i.WriteU32 (str_len);
  for (uint32_t j = 0; j < str_value.size(); j++) {
    if (debug)
      std::cout << " str_value[" << j << "] = " << str_value[j] << std::endl;
    i.WriteU8(str_value[j]);
  }
}
  
void Atag::Deserialize (TagBuffer i) {
  if (debug)
    std::cout << "Atag::Deserialize() called" << std::endl;
  int_value = i.ReadU32 ();
  if (debug)
    std::cout << " int_value = " << int_value << std::endl;
  uint32_t str_len = i.ReadU32();
  if (debug)
    std::cout << " str_len = " << str_len << std::endl;
  str_value.reserve(str_len);
  for (uint32_t j = 0; j < str_len; j++) {
    char v = i.ReadU8();
    str_value.push_back(v);
    if (debug)
      std::cout << " str_value[" << j << "] = " << str_value[j] << std::endl;
  }
  if (debug)
    std::cout << "string is " << str_value << std::endl;
}


void Atag::SetIntValue(uint32_t v_) {
  int_value = v_;
}

uint32_t Atag::GetIntValue(void) const {
  return int_value;
}

void Atag::SetStrValue(const std::string &s) {
  str_value = s;
}

std::string Atag::GetStrValue(void) const {
  return str_value;
}
   
void Atag::Print (std::ostream &os) const {
  os << "(Atag=" << str_value << ":" << int_value << ")" << std::endl;
}
#endif


#ifdef USE_HEADER
class ArchHeader : public Header {
public:
  virtual uint32_t  GetSerializedSize(void) const;
  virtual void      Serialize(Buffer::Iterator) const;
  virtual uint32_t  Deserialize(Buffer::Iterator);
  virtual void      Print(std::ostream &) const;
  static  TypeId    GetTypeId(void);
  virtual TypeId    GetInstanceTypeId(void) const;

  void        SetIntValue(uint32_t);
  uint32_t    GetIntValue(void) const;
  void        SetStrValue(const std::string &s);
  std::string GetStrValue(void) const;
private:
  uint32_t          int_value;
  std::string       str_value;
  const static bool debug = 0;
};


TypeId ArchHeader::GetTypeId(void) {
  static TypeId tid = TypeId("ns3::ArchHeader")
    .AddConstructor<ArchHeader>() // NOTE: Required for ????
    .SetParent<Header>()          // ditto
    ;
  return tid;
}

TypeId ArchHeader::GetInstanceTypeId(void) const {
  return GetTypeId();
}

uint32_t ArchHeader::GetSerializedSize (void) const {
  return 2*sizeof(uint32_t) + str_value.size();
}

void ArchHeader::Serialize (Buffer::Iterator i) const {
  if (debug)
    std::cout << "ArchHeader::Serialize() called" << std::endl;
  if (debug)
    std::cout << " int_value = " << int_value << std::endl;
  i.WriteU32 (int_value);
  uint32_t str_len = str_value.size();
  if (debug)
    std::cout << " str_len = " << str_len << std::endl;
  i.WriteU32 (str_len);
  for (uint32_t j = 0; j < str_value.size(); j++) {
    if (debug)
      std::cout << " str_value[" << j << "] = " << str_value[j] << std::endl;
    i.WriteU8(str_value[j]);
  }
}
  
uint32_t ArchHeader::Deserialize (Buffer::Iterator i) {
  if (debug)
    std::cout << "ArchHeader::Deserialize() called" << std::endl;
  int_value = i.ReadU32 ();
  if (debug)
    std::cout << " int_value = " << int_value << std::endl;
  uint32_t str_len = i.ReadU32();
  if (debug)
    std::cout << " str_len = " << str_len << std::endl;
  str_value.reserve(str_len);
  for (uint32_t j = 0; j < str_len; j++) {
    char v = i.ReadU8();
    str_value.push_back(v);
    if (debug)
      std::cout << " str_value[" << j << "] = " << str_value[j] << std::endl;
  }
  if (debug)
    std::cout << "string is " << str_value << std::endl;
  return 2*sizeof(uint32_t) + str_value.size();
}


void ArchHeader::SetIntValue(uint32_t v_) {
  int_value = v_;
}

uint32_t ArchHeader::GetIntValue(void) const {
  return int_value;
}

void ArchHeader::SetStrValue(const std::string &s) {
  str_value = s;
}

std::string ArchHeader::GetStrValue(void) const {
  return str_value;
}
   
void ArchHeader::Print (std::ostream &os) const {
  os << "(ArchHeader=" << str_value << ":" << int_value << ")";
}
#endif

//-------------  Application ------------------------
class TalkerApp : public Application {
public:
  TalkerApp();
  virtual ~TalkerApp ();
  void ConfRecv(Ptr<Node>,const Address,int);
  void ConfSend(Ptr<Node>,const Address,int);
protected:
  virtual void DoDispose(void);
private:
  virtual void StartApplication(void);
  virtual void StopApplication(void);

  void StartApplicationRecv(void);
  void StartApplicationSend(void);

  void HandleRead(Ptr<Socket>);
  void SendPacket(std::string,char,int);
  void MarkPacket(Ptr<Packet>,std::string,char);

  void CloseConnection(Ptr<Socket> socket);
  void ConnectionSucceeded(Ptr<Socket>);
  void ConnectionFailed(Ptr<Socket>);
  void ConnectionHalfClosed(Ptr<Socket>);
  bool ConnectionRequested(Ptr<Socket>, const Address &);
  void ConnectionCreated(Ptr<Socket>, const Address &);
  void ConnectionCloseRequested(Ptr<Socket>);


  bool            verbose;
  Ptr<Socket>     m_socket;
  Ptr<Socket>     m_servsocket;
  Address         m_local;
  Address         m_remote;
  int             i_am_listener;
  int             recv_cnt;
  int             m_send_size;
  int             m_recv_size;
};

TalkerApp::TalkerApp () :
  verbose(0),
  m_socket(0),
  m_servsocket(0),
  i_am_listener(0),
  recv_cnt(0),
  m_send_size(0),
  m_recv_size(0) {}

TalkerApp::~TalkerApp() {}

void TalkerApp::DoDispose (void) {
  m_socket = 0;
  m_servsocket = 0;
  Application::DoDispose ();
}


void TalkerApp::CloseConnection (Ptr<Socket> sock) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  );
  sock->Close ();
}
void TalkerApp::ConnectionSucceeded(Ptr<Socket> sock) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this 
		   << sock
	
		   );
}
void TalkerApp::ConnectionFailed(Ptr<Socket> sock) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  );
}
void TalkerApp::ConnectionHalfClosed(Ptr<Socket> sock) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  );
}
bool TalkerApp::ConnectionRequested(Ptr<Socket> sock, const Address &addr) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  << addr
                  );
  return(1);
}
void TalkerApp::ConnectionCreated(Ptr<Socket> sock, const Address &addr) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  << addr
                  );
  m_socket = sock;
  m_socket->SetRecvCallback(MakeCallback(&TalkerApp::HandleRead, this));
}
void TalkerApp::ConnectionCloseRequested(Ptr<Socket> sock) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << sock
                  );
}

void TalkerApp::StartApplication() {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  );
  if (i_am_listener)
    StartApplicationRecv();
  else
    StartApplicationSend();
}

void TalkerApp::StartApplicationSend() {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  );

  if (!m_socket) {
    m_socket = Socket::CreateSocket(GetNode(),
				    TcpSocketFactory::GetTypeId());
    m_socket->Bind ();
  }

  m_socket->Connect (m_remote);
  m_socket->SetConnectCallback
    (
     MakeCallback(&TalkerApp::ConnectionSucceeded,this),
     MakeCallback(&TalkerApp::ConnectionFailed,this),
     MakeCallback(&TalkerApp::ConnectionHalfClosed,this)
     );
  m_socket->SetRecvCallback(MakeCallback(&TalkerApp::HandleRead, this));
  if (m_send_size)
    SendPacket(std::string("Hello"),'S',m_send_size);
}

void TalkerApp::StartApplicationRecv() {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  );
  if (!m_servsocket) {
    m_servsocket = Socket::CreateSocket(GetNode(),
					TcpSocketFactory::GetTypeId());
    m_servsocket->Bind (m_local);
    m_servsocket->Listen (0);
    if (verbose)
      std::cout << "  bound and listening on serversocket "
		<< m_servsocket
		<< std::endl;
  }

  m_servsocket->SetAcceptCallback
    (
     MakeCallback(&TalkerApp::ConnectionRequested,this),
     MakeCallback(&TalkerApp::ConnectionCreated,this),
     MakeCallback(&TalkerApp::ConnectionCloseRequested,this)
     );
}
 
void TalkerApp::ConfRecv(Ptr<Node> node, const Address addr,int size) {
  i_am_listener = 1;
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << node
                  << addr
                  << size
                  );
  m_local = addr;
  m_recv_size = size;
  node->AddApplication(this);
}

void TalkerApp::ConfSend(Ptr<Node> node, const Address addr,int size) {
  i_am_listener = 0;
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << node
                  << addr
                  << size
                  );
  m_remote = addr;
  m_send_size = size;
  node->AddApplication(this);
}

void TalkerApp::MarkPacket(Ptr<Packet> p,std::string label,char fill) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << p
                  << label
                  << fill
                  );
  
  if (verbose) {
    p->Print(std::cout << "Before Marking|");
    std::cout << std::endl;
  }

  uint32_t size = p->GetSize();

#ifdef USE_TAGS
  Atag tag;
  tag.SetStrValue(label);
  tag.SetIntValue(size);
  p->AddTag(tag);
  if (verbose) {
    p->Print(std::cout << "After Tag Marking|");
    std::cout << std::endl;
  }
#endif

#ifdef USE_HEADER
  ArchHeader hdr;
  hdr.SetStrValue(label);
  hdr.SetIntValue(size);
  p->AddHeader(hdr);
  if (verbose) {
    p->Print(std::cout << "After Header Marking|");
    std::cout << std::endl;
  }
#endif

}

void TalkerApp::SendPacket(std::string label,char fill,int payloadsize) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << label
                  << fill
                  << payloadsize
                  );

  Ptr<Packet> p = Create<Packet>(payloadsize);
  MarkPacket(p,label,fill);
  if (verbose)
    std::cout << "  sending on socket " << m_socket
	      << std::endl;

#if defined(USE_TAGS) && defined(VERBOSE_TESTING)
  std::cout << std::endl;
  std::cout << "Test of tags on Tx side." << std::endl;
  Atag tag_copy;
  p->FindFirstMatchingTag(tag_copy);
  tag_copy.Print(std::cout << "VRFY Tag::Print| ");
  p->PrintTags(std::cout << "VRFY Packet::PrintTags| ");
  std::cout << std::endl;
#endif

#if defined(USE_HEADER) && defined(VERBOSE_TESTING)
  std::cout << std::endl;
  std::cout << "Test of header on Tx side." << std::endl;
  Packet pcopy;
  pcopy.Deserialize(p->Serialize());
  pcopy.Print(std::cout << "VRFY pcopy: ");
  std::cout << std::endl;
#endif

  m_socket->Send(p);
}

void TalkerApp::StopApplication() {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  );
  if (m_socket)
    m_socket->SetRecvCallback (MakeNullCallback<void, Ptr<Socket> > ());
}

void TalkerApp::HandleRead(Ptr<Socket> socket) {
  NS_LOG_FUNCTION((i_am_listener ? "Server|" : "Client|")
                  << this
                  << socket
                  );

  if (i_am_listener && recv_cnt++ == 0 && m_recv_size) {
    SendPacket(std::string("BackAtYou"),'R', m_recv_size);
  }

  Ptr<Packet> packet;
  while (packet = socket->Recv()){
    SocketRxAddressTag rx_addr_tag;
    bool found;
    found = packet->FindFirstMatchingTag (rx_addr_tag);
    NS_ASSERT (found);
    Address from = rx_addr_tag.GetAddress ();
    // XXX packet->RemoveTag (tag);
    if (InetSocketAddress::IsMatchingType (from)) {
      InetSocketAddress address = InetSocketAddress::ConvertFrom (from);
      std::cout
	<< std::setw(4) << packet->GetSize() << " bytes from "
	<< address.GetIpv4() << " [" << address << "] ";
      packet->Print(std::cout << " ");
      packet->PrintTags(std::cout << "|");
      std::cout << std::endl;


#if defined(USE_TAGS) && defined(VERBOSE_TESTING)
      Atag tg;
      if (packet->FindFirstMatchingTag (tg))
	tg.Print(std::cout << "yanked tag ");
      else 
	std::cout << "Atag not on Rx side" << std::endl;
#endif
#if defined(USE_HEADER) && defined(VERBOSE_TESTING) && defined(I_LIKE_CRASHES)
      /* NOTE: There is a bug here, because the hdr is only on the first
       * Packet, but IF we fragment we'll read the Packet and die. Currently
       * must use Print routines to see if there is a header or not. If we
       * have in order packet delivery then there isn't a problem, if we
       * have out of order delivery then we need a better way....
       */
      ArchHeader hdr; 
      if (packet->RemoveHeader(hdr)) {
	hdr.Print(std::cout << "yanked out ");
      } else {
	std::cout << " no header ";
      }
      std::cout << std::endl;
#endif
    } else {
      std::cout << "not a match in HandleRead" << std::endl;
    }
  }
}


int main (int argc, char *argv[]) {

  int send_size = 400; // lets keep it in one fragment...
  int recv_size = 500;
  if (argc == 3) {
    send_size = atoi(argv[1]);
    recv_size = atoi(argv[2]);
  }

  LogComponentEnable("tcp-2way",LOG_LEVEL_ALL);

  RandomVariable::UseGlobalSeed (1, 1, 2, 3, 5, 8);
  Packet::EnableMetadata();
  NodeContainer c0;
  c0.Create (2);

  // We create the channels first without any IP addressing information
  PointToPointHelper p2p;
  p2p.SetDeviceParameter("DataRate",DataRateValue(DataRate(10000000)));
  p2p.SetChannelParameter ("Delay",TimeValue(MilliSeconds(10)));
  NetDeviceContainer dev0 = p2p.Install(c0);

  // add ip/tcp stack to nodes.
  InternetStackHelper internet;
  internet.Install (c0);

  // Later, we add IP addresses.  
  Ipv4AddressHelper ipv4;
  ipv4.SetBase ("10.1.0.0", "255.255.255.0");
  Ipv4InterfaceContainer ifaces = ipv4.Assign (dev0);

  // Configure the apps.
  uint16_t servPort = 32;
  TalkerApp source;
  source.ConfSend(c0.Get(0),
                  InetSocketAddress (ifaces.GetAddress (1),
                                     servPort),send_size);
  source.Start(Seconds(0.0));

  TalkerApp sink;
  sink.ConfRecv(c0.Get(1),
		InetSocketAddress (Ipv4Address::GetAny(), servPort),
		recv_size);
  sink.Start (Seconds (0.0));

  Simulator::Stop (Seconds(10));

  std::ofstream ascii;
  ascii.open("tcp-2way.tr");
  PointToPointHelper::EnableAsciiAll(ascii);
  PointToPointHelper::EnablePcapAll("tcp-2way");

  Simulator::Run ();
  Simulator::Destroy ();
}