#pragma once

#include "zeek/File.h"
#include "zeek/Reassem.h"
#include "zeek/analyzer/protocol/tcp/TCP_Endpoint.h"
#include "zeek/analyzer/protocol/tcp/TCP_Flags.h"

namespace zeek
	{
namespace packet_analysis::TCP
	{
class TCPSessionAdapter;
	}

class Connection;

namespace analyzer
	{

class Analyzer;

namespace tcp
	{

class TCP_Reassembler final : public Reassembler
	{
public:
	enum Type
		{
		Direct, // deliver to destination analyzer itself
		Forward, // forward to destination analyzer's children
		};

	TCP_Reassembler(analyzer::Analyzer* arg_dst_analyzer,
	                packet_analysis::TCP::TCPSessionAdapter* arg_tcp_analyzer, Type arg_type,
	                TCP_Endpoint* arg_endp);

	void Done();

	void SetDstAnalyzer(analyzer::Analyzer* analyzer) { dst_analyzer = analyzer; }
	void SetType(Type arg_type) { type = arg_type; }

	packet_analysis::TCP::TCPSessionAdapter* GetTCPAnalyzer() { return tcp_analyzer; }

	// Returns the volume of data buffered in the reassembler.
	// First parameter returns data that is above a hole, and thus is
	// waiting on the hole being filled.  Second parameter returns
	// data that has been processed but is awaiting an ACK to free
	// it up.
	//
	// If we're not processing contents, then naturally each of
	// these is empty.
	//
	// WARNING: this is an O(n) operation and potentially very slow.
	void SizeBufferedData(uint64_t& waiting_on_hole, uint64_t& waiting_on_ack) const;

	// How much data is pending delivery since it's not yet reassembled.
	// Includes the data due to holes (so this value is a bit different
	// from waiting_on_hole above; and is computed in a different fashion).
	uint64_t NumUndeliveredBytes() const;

	void SetContentsFile(FilePtr f);
	const FilePtr& GetContentsFile() const { return record_contents_file; }

	void MatchUndelivered(uint64_t up_to_seq, bool use_last_upper);

	// Skip up to seq, as if there's a content gap.
	// Can be used to skip HTTP data for performance considerations.
	void SkipToSeq(uint64_t seq);

	bool DataSent(double t, uint64_t seq, int len, const u_char* data,
	              analyzer::tcp::TCP_Flags flags, bool replaying = true);
	void AckReceived(uint64_t seq);

	// Checks if we have delivered all contents that we can possibly
	// deliver for this endpoint.  Calls TCPSessionAdapter::EndpointEOF()
	// when so.
	void CheckEOF();

	bool HasUndeliveredData() const { return HasBlocks(); }
	bool HadGap() const { return had_gap; }
	bool DataPending() const;
	uint64_t DataSeq() const { return LastReassemSeq(); }

	void DeliverBlock(uint64_t seq, int len, const u_char* data);
	void Deliver(uint64_t seq, int len, const u_char* data);

	TCP_Endpoint* Endpoint() { return endp; }
	const TCP_Endpoint* Endpoint() const { return endp; }

	bool IsOrig() const { return endp->IsOrig(); }

	bool IsSkippedContents(uint64_t seq, int length) const { return seq + length <= seq_to_skip; }

private:
	void Undelivered(uint64_t up_to_seq) override;
	void Gap(uint64_t seq, uint64_t len);

	void RecordToSeq(uint64_t start_seq, uint64_t stop_seq, const FilePtr& f);
	void RecordBlock(const DataBlock& b, const FilePtr& f);
	void RecordGap(uint64_t start_seq, uint64_t upper_seq, const FilePtr& f);

	void BlockInserted(DataBlockMap::const_iterator it) override;
	void Overlap(const u_char* b1, const u_char* b2, uint64_t n) override;

	TCP_Endpoint* endp;

	bool deliver_tcp_contents;
	bool had_gap;
	bool did_EOF;
	bool skip_deliveries;

	uint64_t seq_to_skip;

	bool in_delivery;
	analyzer::tcp::TCP_Flags flags;

	FilePtr record_contents_file; // file on which to reassemble contents

	analyzer::Analyzer* dst_analyzer;
	packet_analysis::TCP::TCPSessionAdapter* tcp_analyzer;

	Type type;
	};

	} // namespace tcp
	} // namespace analyzer
	} // namespace zeek
