#ifndef packetsort_h
#define packetsort_h

// Timestamps can be imprecise and even inconsistent among packets
// from different sources. This class tries to guess a "correct"
// order by looking at TCP sequence numbers.
// 
// In particular, it tries to eliminate "false" content gaps. 

#include "Dict.h"
#include "Conn.h"

enum {
	CONN_PQ,
	GLOBAL_PQ,
	NUM_OF_PQ_LEVEL,
};

class PacketSortElement {
public:
	PacketSortElement(double timestamp, struct pcap_pkthdr* hdr,
				const u_char* pkt, int hdr_size);
	~PacketSortElement();

	double TimeStamp() const		{ return timestamp; }
	const struct pcap_pkthdr* Hdr() const	{ return &hdr; }
	const u_char* Pkt() const		{ return pkt; }
	int HdrSize() const			{ return hdr_size; }
	const IP_Hdr* IPHdr() const		{ return ip_hdr; }

protected:
	double timestamp;
	struct pcap_pkthdr hdr;
	u_char* pkt;
	int hdr_size;

	IP_Hdr* ip_hdr;
	int valid_id;
	ConnID id;
	uint32 seq[2];	// indexed by endpoint
	int tcp_flags;
	int endp;
	int payload_length;

	HashKey* key;

	int pq_index[NUM_OF_PQ_LEVEL];

	friend class PacketSortPQ;
	friend class PacketSortConnPQ;
	friend class PacketSortGlobalPQ;
};

class PacketSortPQ {
public:
	PacketSortPQ()
		{ pq_level = -1; }
	virtual ~PacketSortPQ();

	PacketSortElement* Min() const	{ return (pq.size() > 0) ? pq[0] : 0; }

protected:
	virtual int Cmp(PacketSortElement* a, PacketSortElement* b) = 0;
	int Timestamp_Cmp(PacketSortElement* a, PacketSortElement* b);

	int UpdatePQ(PacketSortElement* prev_e, PacketSortElement* new_e);
	int AddToPQ(PacketSortElement* e);
	int RemoveFromPQ(PacketSortElement* e);

	void Assign(int k, PacketSortElement* e);
	int FixUp(PacketSortElement* e, int k);
	void FixDown(PacketSortElement* e, int k);

	vector<PacketSortElement*> pq;
	int pq_level;
};

// Sort by sequence numbers within a connection
class PacketSortConnPQ : public PacketSortPQ {
public:
	PacketSortConnPQ() 
		{ pq_level = CONN_PQ; }

	int Add(PacketSortElement* e);

	int Remove(PacketSortElement* e) 
		{ return RemoveFromPQ(e); }

protected:
	int Cmp(PacketSortElement* a, PacketSortElement* b);
};

declare(PDict, PacketSortConnPQ);

// Sort by timestamps.
class PacketSortGlobalPQ : public PacketSortPQ {
public:
	PacketSortGlobalPQ();

	int Add(PacketSortElement* e);

	PacketSortElement* RemoveMin(double timestamp);
	int Empty() const { return conn_pq_table.Length() == 0; }

protected:
	int Cmp(PacketSortElement* a, PacketSortElement* b) 
		{ return Timestamp_Cmp(a, b); }
	PacketSortConnPQ* FindConnPQ(PacketSortElement* e);

	PDict(PacketSortConnPQ) conn_pq_table;
};

#define PACKET_SORT	1
extern PacketSortGlobalPQ* packet_sorter;

#endif
