/* -*-  Mode: C++; c-file-style: "gnu"; indent-tabs-mode:nil; -*- */
/*
 *
 * Author: Yu Cao <caoyu08@csnet1.cs.tsinghua.edu.cn>
 */
#ifndef NAMPT_SOCKET_H
#define NAMPT_SOCKET_H

#include <stdint.h>
#include <map>
#include <vector>
#include "ns3/callback.h"
#include "ns3/traced-value.h"
#include "ns3/tcp-socket.h"
#include "ns3/ptr.h"
#include "ns3/ipv4-address.h"
#include "ns3/ipv4-header.h"
#include "ns3/ipv4-interface.h"
#include "ns3/event-id.h"
#include "ns3/object.h"
#include "ns3/inet-socket-address.h"
#include "tcp-tx-buffer.h"
#include "tcp-rx-buffer.h"
#include "rtt-estimator.h"
#include "tcp-header.h"
#include "nampt-l4-protocol.h"


namespace ns3 {

class Ipv4EndPoint;
class Node;
class Packet;
class TcpL4Protocol;
class TcpSocketBase;
class Socket;
class NaMPTSubflow;
class NaMPTCC;


typedef std::map<Ptr<Socket>, Ptr<NaMPTSubflow> >	SubflowSet_t;
typedef std::map<uint32_t, Ptr<Ipv4Interface> >		AddressSet_t;
typedef std::map<SequenceNumber32, Ptr<Packet> >	PacketSet_t;

typedef enum {
	
	NAMPT_MPC	= 30,
	NAMPT_JOIN	= 31,
	
	NAMPT_ADDR	= 32,
	NAMPT_FIN	= 33,
	NAMPT_DSN	= 34,
	NAMPT_ACK	= 35,
	
	NAMPT_EXTHD	= 254,
	NAMPT_NONE	= 255
	
} TcpOptionType_t;

/**
 * \ingroup tcp
 * \brief Signal for the Name-based Multipath Transmission Protocol
 *
 * XXX
 */
class NaMPTSignal : public Object
{
public:
  static TypeId GetTypeId (void);
  /**
  * \brief Constructor
  */
  NaMPTSignal (void);
  NaMPTSignal (TcpOptionType_t t);
  virtual ~NaMPTSignal (void);
  
  TcpOptionType_t	sigType;
  uint8_t			sigLen; // bytes, including TLV
  union {
  	struct {
		uint32_t localToken;
  	} mpc;
	struct {
		uint32_t peerToken;
	} join;
	struct {
		uint16_t byteNum;
	} exthd;
	struct {
		uint32_t sigSeq;
		uint32_t ipv4Addr;
	} addr;
	struct {
		uint32_t sigSeq;
		uint64_t dataSeq;
		uint32_t dataLen;
		uint32_t sfSeq;
	} dsn;
	struct {
		uint32_t sigSeq;
	} fin;
	struct {
		uint32_t sigSeq;
	} ack;
  }sigData;
  
};


/**
 * \ingroup tcp
 * \brief Header for the Name-based Multipath Transmission Protocol
 *
 * XXX
 */
class NaMPTHeader : public TcpHeader
{
public:
  static TypeId GetTypeId (void);
  /**
  * \brief Constructor
  */
  NaMPTHeader ();
  NaMPTHeader (NaMPTHeader &res);
  NaMPTHeader (const TcpHeader &res);
  virtual ~NaMPTHeader ();
  virtual TypeId GetInstanceTypeId (void) const;
  virtual void Print (std::ostream &os) const;
  virtual uint32_t GetSerializedSize (void) const;
  virtual void Serialize (Buffer::Iterator start) const;
  virtual uint32_t Deserialize (Buffer::Iterator start);

  Ptr<NaMPTSignal> GetFirstSignal (std::vector<Ptr<NaMPTSignal> >::iterator &iter);
  Ptr<NaMPTSignal> GetNextSignal (std::vector<Ptr<NaMPTSignal> >::iterator &iter);
  Ptr<NaMPTSignal> AddSignal (TcpOptionType_t sigType);
  void AddSignal(Ptr<NaMPTSignal> signal);
  Ptr<NaMPTSignal> FindSignal (TcpOptionType_t sigType, bool bRemove);
  void RemoveAllSignals (void);
  
protected:
  void UpdateLength (void);
  
  std::vector<Ptr<NaMPTSignal> >	m_signals;
};


/**
 * \ingroup socket
 * \ingroup tcp
 *
 * \brief An implementation of a stream socket using Name-based Multipath Transmission
 *
 * XXX
 */
typedef struct _TcpAttributes
{
  bool		m_shutdownSend;
  bool		m_shutdownRecv;
  uint32_t	m_segmentSize;
  Time		m_cnTimeout;
  uint32_t	m_cnRetries;
  Time		m_delAckTimeout;
  uint32_t	m_delAckMaxCount;
  Time		m_persistTimeout;
} TcpAttributes;

class NaMPTSocket : public TcpSocket
{
public:
  static TypeId GetTypeId (void);
  /**
   * \brief Constructor
   */
  NaMPTSocket (void);
  NaMPTSocket (const NaMPTSocket& sock);
  virtual ~NaMPTSocket (void);

  // Set associated Node, NaMPTL4Protocol, typeids to this socket
  virtual void SetNode (Ptr<Node> node);
  virtual void SetTcp (Ptr<TcpL4Protocol> tcp);
  
  // Hook application callback functions: coming from ns3::socket
  virtual void SetConnectCallback (Callback<void, Ptr<Socket> > connectionSucceeded,
                                             Callback<void,  Ptr<Socket> > connectionFailed);
  virtual void SetCloseCallbacks (Callback<void, Ptr<Socket> > normalClose,
                                           Callback<void, Ptr<Socket> > errorClose);
  virtual void SetAcceptCallback (Callback<bool, Ptr<Socket>, const Address &> connectionRequest,
                                           Callback<void, Ptr<Socket>, const Address&> newConnectionCreated);
  virtual void SetDataSentCallback (Callback<void, Ptr<Socket>, uint32_t> dataSent);
  virtual void SetSendCallback (Callback<void, Ptr<Socket>, uint32_t> sendCb);
  virtual void SetRecvCallback (Callback<void, Ptr<Socket> > receivedData);

  // Necessary implementations of null functions from ns3::Socket
  virtual enum SocketErrno GetErrno (void) const;	// returns m_errno
  virtual enum SocketType GetSocketType (void) const;	// returns socket type 
  virtual Ptr<Node> GetNode (void) const;	// returns m_node
  virtual int Bind (void);	// Bind a socket by setting up endpoint in TcpL4Protocol
  virtual int Bind (const Address &address);	// ... endpoint of specific addr or port
  virtual int Bind6 (void);
  virtual int Connect (const Address &address);	// Setup endpoint and call ProcessAction() to connect
  virtual int Listen (void);	// Verify the socket is in a correct state and call ProcessAction() to listen  
  virtual int Close (void);	// Close by app: Kill socket upon tx buffer emptied
  virtual int ShutdownSend (void);	// Assert the m_shutdownSend flag to prevent send to network
  virtual int ShutdownRecv (void);	// Assert the m_shutdownRecv flag to prevent forward to app
  virtual int Send (Ptr<Packet> p, uint32_t flags);	// Call by app to send data to network
  virtual int SendTo (Ptr<Packet> p, uint32_t flags, const Address &toAddress);	// Same as Send(), toAddress is insignificant
  virtual Ptr<Packet> Recv (uint32_t maxSize, uint32_t flags);	// Return a packet to be forwarded to app
  virtual Ptr<Packet> RecvFrom (uint32_t maxSize, uint32_t flags, Address &fromAddress);	// ... and write the remote address at fromAddress
  virtual uint32_t GetTxAvailable (void) const;	// Available Tx buffer size
  virtual uint32_t GetRxAvailable (void) const;	// Available-to-read data size, i.e. value of m_rxAvailable
  virtual int GetSockName (Address &address) const;	// Return local addr:port in address
  virtual bool SetAllowBroadcast (bool allowBroadcast);
  virtual bool GetAllowBroadcast () const;
  
  virtual void BindToNetDevice (Ptr<NetDevice> netdevice);	// NetDevice with my m_endPoint

protected:
  // Necessary Implementation of Attribute get/set from ns3::TcpSocket
  virtual void     SetSndBufSize (uint32_t size);
  virtual uint32_t GetSndBufSize (void) const;
  virtual void     SetRcvBufSize (uint32_t size);
  virtual uint32_t GetRcvBufSize (void) const;
  virtual void     SetSegSize (uint32_t size);
  virtual uint32_t GetSegSize (void) const;
  virtual void     SetConnTimeout (Time timeout);
  virtual Time     GetConnTimeout (void) const;
  virtual void     SetConnCount (uint32_t count);
  virtual uint32_t GetConnCount (void) const;
  virtual void     SetDelAckTimeout (Time timeout);
  virtual Time     GetDelAckTimeout (void) const;
  virtual void     SetDelAckMaxCount (uint32_t count);
  virtual uint32_t GetDelAckMaxCount (void) const;
  virtual void     SetPersistTimeout (Time timeout);
  virtual Time     GetPersistTimeout (void) const;

  virtual void     SetSSThresh (uint32_t threshold);
  virtual uint32_t GetSSThresh (void) const;
  virtual void     SetInitialCwnd (uint32_t cwnd);
  virtual uint32_t GetInitialCwnd (void) const;

  /* 
     *************************
     *   NaMPT-related functions  *
     *************************
   */
public:
  void SetRttTypeId (TypeId rttTypeId);
  void SetSocketTypeId (TypeId socketTypeId);
  void SetFirstSocket (Ptr<TcpSocketBase> baseSock);
  static Ptr<NaMPTSocket> EncapsulateSocket (Ptr<TcpSocketBase> baseSock);
  static void HookPacketDelivery (Ptr<TcpSocketBase> baseSocket, Ptr<NaMPTSocket> shellSocket);
  static bool HaveSignals (Ptr<Packet> packet);
  Ptr<NaMPTSubflow> SwitchToMultipathMode(void);
  Ptr<NaMPTSubflow> SetupSubflow(Ptr<TcpSocketBase> baseSocket, bool bFirst);
  void ReceiveSubflowJoin(Ptr<Packet> p, Ipv4Header const &header, Ptr<Ipv4Interface> incomingInterface);
  void TryPeerAddress(uint32_t addr);
  void TryMultiPoints(void);
  void AdvertizeAddresses(Ptr<NaMPTSubflow> subflow);
  void CleanupSubflow(Ptr<NaMPTSubflow> subflow);
  
  uint16_t GetLocalPort (void);
  uint16_t GetPeerPort (void);
  bool GetLocalAddress (Ipv4Address &addr);
  bool GetPeerAddress (Ipv4Address &addr);
  uint32_t GetLocalToken (void);
  uint32_t GetPeerToken(void);

protected:
  void ForwardUp (Ptr<Packet> packet, Ipv4Header header, 
                       uint16_t port, Ptr<Ipv4Interface> incomingInterface);
  Ptr<TcpSocketBase> CreateBaseSocket(void);
  static Ptr<TcpSocketBase> ConvertCongestionControl(Ptr<TcpSocketBase> baseSock);
  Ptr<NaMPTSubflow> GetSubflow(Ptr<Socket> sock);
  void InitAddressSet(void);
  bool AddSubflow(uint32_t srcAddr, uint16_t srcPort, uint32_t desAddr, uint16_t desPort, bool bCheck = true);
  void ReschedulingPackets(Ptr<NaMPTSubflow> subflow); 
  uint64_t CalculateLinkedAlpha(Ptr<NaMPTCC> sock, uint32_t segmentSize);


  // Hook app's callback functions
  void DoConnectionSucceeds (Ptr<Socket> sock);
  void DoConnectionFails (Ptr<Socket> sock);
  void DoNormalClose (Ptr<Socket> sock);
  void DoErrorClose (Ptr<Socket> sock);
  bool DoConnectionRequests (Ptr<Socket> sock, const Address& from);
  void DoConnectionAccepted (Ptr<Socket> sock, const Address& from);
  void DoDataSent (Ptr<Socket> sock, uint32_t txSize);
  void DoSendMore (Ptr<Socket> sock, uint32_t txAvailable);
  void DoRead (Ptr<Socket> sock);

public:
  // For subflows
  void DoSubflowSucceeds(Ptr<Socket> sock);
  void DoSubflowFails(Ptr<Socket> sock);
  bool DoSubflowRequests(Ptr<Socket> sock, const Address& from);
  void DoSubflowAccepted(Ptr<Socket> sock, const Address& from);
  void DoSubflowDataSent(Ptr<Socket> sock, uint32_t bytes);
  void DoSubflowSend(Ptr<Socket> sock, uint32_t txAvailable);
  void DoSubflowRecv(Ptr<Socket> sock);
  void DoSubflowPeerClose(Ptr<Socket> sock);
  void DoSubflowNormalClose(Ptr<Socket> sock);
  void DoSubflowErrorClose(Ptr<Socket> sock);

  void DoDCTCPEnterCWR(Ptr<NaMPTSubflow> subflow, Ptr<NaMPTCC> sock);
  void DoXMPEnterCWR(Ptr<NaMPTSubflow> subflow, Ptr<NaMPTCC> sock);

protected:
  bool SchedulingQuota(Ptr<NaMPTSubflow> subflow, uint32_t& minQuota, uint32_t& maxQuota);
  void SchedulingPackets(Ptr<NaMPTSubflow> subflow);
  void NaiveAssembling(Ptr<NaMPTSubflow> subflow);
  //void OrderAssembling(Ptr<NaMPTSubflow> subflow);
  bool AssignPackets(uint32_t maxSize, Ptr<NaMPTSubflow> subflow);

  bool DoSubflowNewAck(Ptr<NaMPTCC> sock, const SequenceNumber32& ackSeq);
  bool DoSubflowDupAck(Ptr<NaMPTCC> sock, const SequenceNumber32& ackSeq, uint32_t count);
  bool DoSubflowTimeout(Ptr<NaMPTCC> sock);
  
protected:
  Ptr<Node>				m_node;
  Ptr<TcpL4Protocol>		m_tcp;
  TypeId					m_rttTypeId;
  TypeId					m_socketTypeId;

private:
friend class TcpSocketBase;

  Ptr<TcpSocketBase>		m_firstSocket;
  SubflowSet_t			m_subflows;
  AddressSet_t			m_addresses; // init by bind() and EncapsulateSocket()
  uint32_t				m_nextAddr; // for round robin path management
  std::vector<InetSocketAddress>		m_accessPoints;

  uint32_t				m_localToken;
  uint32_t				m_peerToken;
  uint16_t				m_localPort;
  uint16_t				m_peerPort;
  std::string			m_sessionName; // init by SwitchToMultipathMode()

  bool						m_bClosed;
  bool						m_bScheduling;
  std::vector<Ptr<Packet> >	m_pendingSynPacket;
  uint64_t					m_alpha; // coupled parameter

  // Attributes
  NaMPTAlgo_t			m_ccAlgo;
  NaMPTAlgo_t			m_schedulingAlgo;
  NaMPTAlgo_t			m_assemblingAlgo;
  uint32_t				m_initGamma;
  uint32_t				m_backoffBeta;
  uint32_t				m_pathManagement;
  TcpAttributes			m_attributes;

  // Statistics
  uint32_t				m_peakOfOutputQueue;
  double					m_instRate;
  double					m_timeTag; // init by Connect and SwitchToMultipathMode
  uint64_t				m_flowSize; // how many bytes are transmitted.

  // Queue Manager
  SequenceNumber32		m_nextTx;
  TcpTxBuffer			m_inputQueue;
  TcpRxBuffer			m_outputQueue;
  PacketSet_t			m_reschedulingQueue;
  uint32_t				m_maxIngoingSize; // serve for subflows
  uint32_t				m_maxOutgoingSize;
  
};

} // namespace ns3

#endif /* NAMPT_SOCKET_H */
