#include "IP.h"
#include "PacketSort.h"

PacketSortGlobalPQ* packet_sorter = 0;

PacketSortElement::PacketSortElement(double arg_timestamp,
			struct pcap_pkthdr* arg_hdr, const u_char* arg_pkt,
			int arg_hdr_size)
	{
	timestamp = arg_timestamp;
	hdr = *arg_hdr;
	hdr_size = arg_hdr_size;

	pkt = new u_char[hdr.caplen];
	memcpy(pkt, arg_pkt, hdr.caplen);

	valid_id = 0;
	ip_hdr = 0;
	key = 0;

	// Now check if it is a "parsable" TCP packet.
	uint32 caplen = hdr.caplen;
	uint32 tcp_offset;

	if ( caplen >= sizeof(struct ip) + hdr_size )
		{
		const struct ip* ip = (const struct ip*) (pkt + hdr_size);
		if ( ip->ip_v == 4 )
			ip_hdr = new IP_Hdr(ip);
		else
			ip_hdr = new IP_Hdr((const struct ip6_hdr*) ip);

		if ( ip_hdr->NextProto() == IPPROTO_TCP &&
		      // Note: can't sort fragmented packets
		     (ip_hdr->FragField() & 0x3fff) == 0 )
			{
			tcp_offset = hdr_size + ip_hdr->HdrLen();
			if ( caplen >= tcp_offset + sizeof(struct tcphdr) )
				{
				const struct tcphdr* tp = (const struct tcphdr*)
							(pkt + tcp_offset);

				id.src_addr = ip_hdr->SrcAddr();
				id.dst_addr = ip_hdr->DstAddr();
				id.src_port = tp->th_sport;
				id.dst_port = tp->th_dport;

				endp = addr_port_canon_lt(id.src_addr,
							  id.src_port, 
							  id.dst_addr,
							  id.dst_port) ? 0 : 1;

				seq[endp] = ntohl(tp->th_seq);

				if ( tp->th_flags & TH_ACK )
					seq[1-endp] = ntohl(tp->th_ack);
				else
					seq[1-endp] = 0;

				tcp_flags = tp->th_flags;

				// DEBUG_MSG("%.6f: %u, %u\n", timestamp, seq[0], seq[1]);

				payload_length = ip_hdr->PayloadLen() - tp->th_off * 4;

				key = id.BuildConnKey();

				valid_id = 1;
				}
			}
		}

	if ( ! valid_id )
		DEBUG_MSG("%.6f packet without valid id\n", timestamp);
	}

PacketSortElement::~PacketSortElement()
	{
	delete [] pkt;
	delete ip_hdr;
	delete key;
	}

PacketSortPQ::~PacketSortPQ()
	{
	for ( int i = 0; i < (int) pq.size(); ++i )
		{
		delete pq[i];
		pq[i] = 0;
		}
	}

int PacketSortPQ::Timestamp_Cmp(PacketSortElement* a, PacketSortElement* b) 
	{ 
	double d = a->timestamp - b->timestamp;

	if ( d > 0 ) return 1;
	else if ( d < 0 ) return -1;
	else return 0;
	}

int PacketSortPQ::UpdatePQ(PacketSortElement* prev_e, PacketSortElement* new_e)
	{
	int index = prev_e->pq_index[pq_level];

	new_e->pq_index[pq_level] = index;
	pq[index] = new_e;

	if ( Cmp(prev_e, new_e) > 0 )
		return FixUp(new_e, index);
	else
		{
		FixDown(new_e, index);
		return index == 0;
		}
	}

int PacketSortPQ::AddToPQ(PacketSortElement* new_e)
	{
	int index = pq.size();

	new_e->pq_index[pq_level] = index;
	pq.push_back(new_e);

	return FixUp(new_e, index);
	}

int PacketSortPQ::RemoveFromPQ(PacketSortElement* prev_e)
	{
	if ( pq.size() > 1 )
		{
		PacketSortElement* new_e = pq[pq.size() - 1];
		pq.pop_back();
		return UpdatePQ(prev_e, new_e);		
		}
	else
		{
		pq.pop_back();
		return 1;
		}
	}

void PacketSortPQ::Assign(int k, PacketSortElement* e)
	{
	pq[k] = e;
	e->pq_index[pq_level] = k;
	}

int PacketSortConnPQ::Cmp(PacketSortElement* a, PacketSortElement* b) 
	{ 
	// Note: here we do not distinguish between packets without
	// an ACK and packets with seq/ack of 0. The later will sorted
	// only by their timestamps.

	if ( a->seq[0] && b->seq[0] && a->seq[0] != b->seq[0] )
		return (a->seq[0] > b->seq[0]) ? 1 : -1;

	else if ( a->seq[1] && b->seq[1] && a->seq[1] != b->seq[1] )
		return (a->seq[1] > b->seq[1]) ? 1 : -1;

	else
		return Timestamp_Cmp(a, b);
	}

int PacketSortPQ::FixUp(PacketSortElement* e, int k)
	{
	if ( k == 0 )
		{
		Assign(0, e);
		return 1;
		}

	int parent = (k-1) / 2;
	if ( Cmp(pq[parent], e) > 0 )
		{
		Assign(k, pq[parent]);
		return FixUp(e, parent);
		}
	else
		{
		Assign(k, e);
		return 0;
		}
	}

void PacketSortPQ::FixDown(PacketSortElement* e, int k)
	{
	uint32 kid = k * 2 + 1;

	if ( kid >= pq.size() )
		{
		Assign(k, e);
		return;
		}

	if ( kid + 1 < pq.size() && Cmp(pq[kid], pq[kid+1]) > 0 )
		++kid;

	if ( Cmp(e, pq[kid]) > 0 )
		{
		Assign(k, pq[kid]);
		FixDown(e, kid);
		}
	else
		Assign(k, e);
	}


int PacketSortConnPQ::Add(PacketSortElement* e) 
	{ 
#if 0
	int endp = e->endp;
	uint32 end_seq = e->seq[endp] + e->payload_length;

	int p = 1 - endp;
	if ( (e->tcp_flags & TH_RST) && ! (e->tcp_flags & TH_ACK) )
		{
		DEBUG_MSG("%.6f %c: %u -> %u\n",
			  e->TimeStamp(), (p == endp) ? 'S' : 'A', 
			  e->seq[p], next_seq[p]);
		e->seq[p] = next_seq[p];
		}

	if ( end_seq > next_seq[endp] )
		next_seq[endp] = end_seq;
#endif

	return AddToPQ(e); 
	}

static void DeleteConnPQ(void* p)
	{ 
	delete (PacketSortConnPQ*) p;
	}

PacketSortGlobalPQ::PacketSortGlobalPQ() 
	{ 
	pq_level = GLOBAL_PQ; 
	conn_pq_table.SetDeleteFunc(DeleteConnPQ);
	}

int PacketSortGlobalPQ::Add(PacketSortElement* e)
	{
	if ( e->valid_id )
		{
		PacketSortConnPQ* conn_pq = FindConnPQ(e);
		PacketSortElement* prev_min = conn_pq->Min();

		if ( conn_pq->Add(e) )
			{
			ASSERT(conn_pq->Min() != prev_min);

			if ( prev_min )
				return UpdatePQ(prev_min, e);
			else
				return AddToPQ(e);
			}

		else
			{
			ASSERT(conn_pq->Min() == prev_min);
			return 0;
			}
		}
	else
		return AddToPQ(e);
	}

PacketSortElement* PacketSortGlobalPQ::RemoveMin(double timestamp)
	{
	PacketSortElement* e = Min();

	if ( ! e || e->timestamp > timestamp )
		return 0;

	if ( e->valid_id )
		{
		PacketSortConnPQ* conn_pq = FindConnPQ(e);

		conn_pq->Remove(e);
		PacketSortElement* new_e = conn_pq->Min();

		if ( new_e )
			UpdatePQ(e, new_e);
		else
			{
			RemoveFromPQ(e);
			conn_pq_table.Remove(e->key);
			}
		}
	else
		RemoveFromPQ(e);

	return e;
	}

PacketSortConnPQ* PacketSortGlobalPQ::FindConnPQ(PacketSortElement* e)
	{
	if ( ! e->valid_id )
		internal_error("cannot find a connection for an invalid id");
	
	PacketSortConnPQ* pq = (PacketSortConnPQ*) conn_pq_table.Lookup(e->key);
	if ( ! pq )
		{
		pq = new PacketSortConnPQ();
		conn_pq_table.Insert(e->key, pq);
		}

	return pq;
	}
