#pragma once

#include "Packet.h"
#include "IpAddress.h"
#include "PointerVector.h"
#include <unordered_map>
#include <chrono>
#include <map>
#include <list>
#include <time.h>
#include <array>
#include <functional>

/// @file
/// This is an implementation of TCP reassembly logic, which means reassembly of TCP messages spanning multiple TCP
/// segments (or packets).<BR> This logic can be useful in analyzing messages for a large number of protocols
/// implemented on top of TCP including HTTP, SSL/TLS, FTP and many many more.
///
/// __General Features:__
/// - Manage multiple TCP connections under one pcpp#TcpReassembly instance
/// - Support TCP retransmission
/// - Support out-of-order packets
/// - Support missing TCP data
/// - TCP connections can end "naturally" (by FIN/RST packets) or manually by the user
/// - Support callbacks for new TCP data, connection start and connection end
///
/// __Logic Description:__
/// - The user creates an instance of the pcpp#TcpReassembly class
/// - Then the user starts feeding it with TCP packets
/// - The pcpp#TcpReassembly instance manages all TCP connections from the packets it's being fed. For each connection
///   it manages its 2 sides (A->B and B->A).
/// - When a packet arrives, it is first classified to a certain TCP connection
/// - Then it is classified to a certain side of the TCP connection
/// - Then the pcpp#TcpReassembly logic tries to understand if the data in this packet is the expected data
///   (sequence-wise) and if it's new (e.g isn't a retransmission).
/// - If the packet data matches these criteria a callback is being invoked. This callback is supplied by the user in
///   the creation of the pcpp#TcpReassembly instance. This callback contains the new data (of course), but also
///   information about the connection (5-tuple, 4-byte hash key describing the connection, etc.) and also a pointer to
///   a "user cookie", meaning a pointer to a structure provided by the user during the creation of the
///   pcpp#TcpReassembly instance.
/// - If the data in this packet isn't new, it's being ignored
/// - If the data in this packet isn't expected (meaning this packet came out-of-order), then the data is being queued
///   internally and will be sent to the user when its turn arrives (meaning, after the data before arrives).
/// - If the missing data doesn't arrive until a new message from the other side of the connection arrives or until the
///   connection ends - this will be considered as missing data and the queued data will be sent to the user, but the
///   string "[X bytes missing]" will be added to the message sent in the callback.
/// - pcpp#TcpReassembly supports 2 more callbacks - one is invoked when a new TCP connection is first seen and the
///   other when it's ended (either by a FIN/RST packet or manually by the user). Both of these callbacks contain data
///   about the connection (5-tuple, 4-byte hash key describing the connection, etc.) and also a pointer to a "user
///   cookie", meaning a pointer to a structure provided by the user during the creation of the pcpp#TcpReassembly
///   instance. The end connection callback also provides the reason for closing it ("naturally" or manually).
///
/// __Basic Usage and APIs:__
/// - pcpp#TcpReassembly c'tor - Create an instance, provide the callbacks and the user cookie to the instance
/// - pcpp#TcpReassembly#reassemblePacket() - Feed pcpp#TcpReassembly instance with packets
/// - pcpp#TcpReassembly#closeConnection() - Manually close a connection by a flow key
/// - pcpp#TcpReassembly#closeAllConnections() - Manually close all currently opened connections
/// - pcpp#TcpReassembly#OnTcpMessageReady callback - Invoked when new data arrives on a certain connection. Contains
/// the new data as well as connection data (5-tuple, flow key).
/// - pcpp#TcpReassembly#OnTcpConnectionStart callback - Invoked when a new connection is identified
/// - pcpp#TcpReassembly#OnTcpConnectionEnd callback - Invoked when a connection ends (either by FIN/RST or manually by
///   the user).
///
/// __Additional information:__
/// When the connection is closed the information is not being deleted from memory immediately. There is a delay between
/// these moments. Existence of this delay is caused by two reasons:
/// - pcpp#TcpReassembly#reassemblePacket() should detect the packets that arrive after the FIN packet has been
/// received.
/// - The user can use the information about connections managed by pcpp#TcpReassembly instance. Following methods are
///   used for this purpose: pcpp#TcpReassembly#getConnectionInformation and pcpp#TcpReassembly#isConnectionOpen.
/// Cleaning of memory can be performed automatically (the default behavior) by pcpp#TcpReassembly#reassemblePacket() or
/// manually by calling pcpp#TcpReassembly#purgeClosedConnections in the user code. Automatic cleaning is performed once
/// per second.
///
/// The struct pcpp#TcpReassemblyConfiguration allows to setup the parameters of cleanup. Following parameters are
/// supported:
/// - pcpp#TcpReassemblyConfiguration#doNotRemoveConnInfo - if this member is set to false the automatic cleanup mode is
///   applied.
/// - pcpp#TcpReassemblyConfiguration#closedConnectionDelay - the value of delay expressed in seconds. The minimum value
///   is 1.
/// - pcpp#TcpReassemblyConfiguration#maxNumToClean - to avoid performance overhead when the cleanup is being performed,
///   this parameter is used. It defines the maximum number of items to be removed per one call of
///   pcpp#TcpReassembly#purgeClosedConnections.
/// - pcpp#TcpReassemblyConfiguration#maxOutOfOrderFragments - the maximum number of unmatched fragments to keep per
///   flow before missed fragments are considered lost. A value of 0 means unlimited.

/// @namespace pcpp
/// @brief The main namespace for the PcapPlusPlus lib
namespace pcpp
{
	/// @struct ConnectionData
	/// Represents basic TCP/UDP + IP connection data
	struct ConnectionData
	{
		/// Source IP address
		IPAddress srcIP;
		/// Destination IP address
		IPAddress dstIP;
		/// Source TCP/UDP port
		uint16_t srcPort;
		/// Destination TCP/UDP port
		uint16_t dstPort;
		/// A 4-byte hash key representing the connection
		uint32_t flowKey;
		/// Start timestamp of the connection with microsecond precision
		timeval startTime;
		/// End timestamp of the connection with microsecond precision
		timeval endTime;
		/// Start timestamp of the connection with nanosecond precision
		std::chrono::time_point<std::chrono::high_resolution_clock> startTimePrecise;
		/// End timestamp of the connection with nanosecond precision
		std::chrono::time_point<std::chrono::high_resolution_clock> endTimePrecise;

		/// A c'tor for this struct that basically zeros all members
		ConnectionData() : srcPort(0), dstPort(0), flowKey(0), startTime(), endTime()
		{}

		/// Set the start time of the connection
		/// @param[in] startTimeValue timestamp value
		void setStartTime(const std::chrono::time_point<std::chrono::high_resolution_clock>& startTimeValue);

		/// Set the end time of the connection
		/// @param[in] endTimeValue timestamp value
		void setEndTime(const std::chrono::time_point<std::chrono::high_resolution_clock>& endTimeValue);
	};

	class TcpReassembly;

	/// @class TcpStreamData
	/// When following a TCP connection each packet may contain a piece of the data transferred between the client and
	/// the server. This class represents these pieces: each instance of it contains a piece of data, usually extracted
	/// from a single packet, as well as information about the connection
	class TcpStreamData
	{
	public:
		/// A c'tor for this class that get data from outside and set the internal members
		/// @param[in] tcpData A pointer to buffer containing the TCP data piece
		/// @param[in] tcpDataLength The length of the buffer
		/// @param[in] missingBytes The number of missing bytes due to packet loss.
		/// @param[in] connData TCP connection information for this TCP data
		/// @param[in] timestamp when this packet was received
		TcpStreamData(const uint8_t* tcpData, size_t tcpDataLength, size_t missingBytes, const ConnectionData& connData,
		              std::chrono::time_point<std::chrono::high_resolution_clock> timestamp)
		    : m_Data(tcpData), m_DataLen(tcpDataLength), m_MissingBytes(missingBytes), m_Connection(connData),
		      m_Timestamp(timestamp)
		{}

		/// A getter for the data buffer
		/// @return A pointer to the buffer
		const uint8_t* getData() const
		{
			return m_Data;
		}

		/// A getter for buffer length
		/// @return Buffer length
		size_t getDataLength() const
		{
			return m_DataLen;
		}

		/// A getter for missing byte count due to packet loss.
		/// @return Missing byte count
		size_t getMissingByteCount() const
		{
			return m_MissingBytes;
		}

		/// Determine if bytes are missing. getMissingByteCount can be called to determine the number of missing bytes.
		/// @return true if bytes are missing.
		bool isBytesMissing() const
		{
			return getMissingByteCount() > 0;
		}

		/// A getter for the connection data
		/// @return The const reference to connection data
		const ConnectionData& getConnectionData() const
		{
			return m_Connection;
		}

		/// @return A microsecond precision of the packet timestamp
		timeval getTimeStamp() const;

		/// @return A nanosecond precision of the packet timestamp
		std::chrono::time_point<std::chrono::high_resolution_clock> getTimeStampPrecise() const
		{
			return m_Timestamp;
		}

	private:
		const uint8_t* m_Data;
		size_t m_DataLen;
		size_t m_MissingBytes;
		const ConnectionData& m_Connection;
		std::chrono::time_point<std::chrono::high_resolution_clock> m_Timestamp;
	};

	/// @struct TcpReassemblyConfiguration
	/// A structure for configuring the TcpReassembly class
	struct TcpReassemblyConfiguration
	{
		/// The flag indicating whether to remove the connection data after a connection is closed
		bool removeConnInfo;

		/// How long the closed connections will not be cleaned up. The value is expressed in seconds. If the value is
		/// set to 0 then TcpReassembly should use the default value. This parameter is only relevant if removeConnInfo
		/// is equal to true.
		uint32_t closedConnectionDelay;

		/// The maximum number of items to be cleaned up per one call of purgeClosedConnections. If the value is set to
		/// 0 then TcpReassembly should use the default value. This parameter is only relevant if removeConnInfo is
		/// equal to true.
		uint32_t maxNumToClean;

		/// The maximum number of fragments with a non-matching sequence-number to store per connection flow before
		/// packets are assumed permanently missed. If the value is 0, TcpReassembly should keep out of order fragments
		/// indefinitely, or until a message from the paired side is seen.
		uint32_t maxOutOfOrderFragments;

		/// To enable to clear buffer once packet contains data from a different side than the side seen before
		bool enableBaseBufferClearCondition;

		/// A c'tor for this struct
		/// @param[in] removeConnInfo The flag indicating whether to remove the connection data after a connection is
		/// closed. The default is true
		/// @param[in] closedConnectionDelay How long the closed connections will not be cleaned up. The value is
		/// expressed in seconds. If it's set to 0 the default value will be used. The default is 5.
		/// @param[in] maxNumToClean The maximum number of items to be cleaned up per one call of
		/// purgeClosedConnections. If it's set to 0 the default value will be used. The default is 30.
		/// @param[in] maxOutOfOrderFragments The maximum number of unmatched fragments to keep per flow before missed
		/// fragments are considered lost. The default is unlimited.
		/// @param[in] enableBaseBufferClearCondition To enable to clear buffer once packet contains data from a
		/// different side than the side seen before
		explicit TcpReassemblyConfiguration(bool removeConnInfo = true, uint32_t closedConnectionDelay = 5,
		                                    uint32_t maxNumToClean = 30, uint32_t maxOutOfOrderFragments = 0,
		                                    bool enableBaseBufferClearCondition = true)
		    : removeConnInfo(removeConnInfo), closedConnectionDelay(closedConnectionDelay),
		      maxNumToClean(maxNumToClean), maxOutOfOrderFragments(maxOutOfOrderFragments),
		      enableBaseBufferClearCondition(enableBaseBufferClearCondition)
		{}
	};

	/// @class TcpReassembly
	/// A class containing the TCP reassembly logic. Please refer to the documentation at the top of TcpReassembly.h for
	/// understanding how to use this class
	class TcpReassembly
	{
	public:
		/// An enum for connection end reasons
		enum ConnectionEndReason
		{
			/// Connection ended because of FIN or RST packet
			TcpReassemblyConnectionClosedByFIN_RST,
			/// Connection ended manually by the user
			TcpReassemblyConnectionClosedManually
		};

		/// An enum for providing reassembly status for each processed packet
		enum ReassemblyStatus
		{
			/// The processed packet contains valid TCP payload, and its payload is processed by
			/// `OnMessageReadyCallback` callback function. The packet may be:
			/// 1. An in-order TCP packet, meaning `packet_sequence == sequence_expected`.
			///    Note if there's any buffered out-of-order packet waiting for this packet, their associated callbacks
			///    are called in this `reassemblePacket` call.
			/// 2. An out-of-order TCP packet which satisfy `packet_sequence < sequence_expected && packet_sequence +
			///    packet_payload_length > sequence_expected`. Note only the new data (the `[sequence_expected,
			///    packet_sequence + packet_payload_length]` part ) is processed by `OnMessageReadyCallback` callback
			///    function.
			TcpMessageHandled,
			/// The processed packet is an out-of-order TCP packet, meaning `packet_sequence > sequence_expected`. It's
			/// buffered so no `OnMessageReadyCallback` callback function is called. The callback function for this
			/// packet maybe called LATER, under different circumstances:
			/// 1. When an in-order packet which is right before this packet arrives(case 1 and case 2 described in
			///    `TcpMessageHandled` section above).
			/// 2. When a FIN or RST packet arrives, which will clear the buffered out-of-order packets of this side.
			///    If this packet contains "new data", meaning `(packet_sequence <= sequence_expected) &&
			///    (packet_sequence + packet_payload_length > sequence_expected)`, the new data is processed by
			///    `OnMessageReadyCallback` callback.
			OutOfOrderTcpMessageBuffered,
			/// The processed packet is a FIN or RST packet with no payload.
			/// Buffered out-of-order packets will be cleared.
			/// If they contain "new data", the new data is processed by `OnMessageReadyCallback` callback.
			FIN_RSTWithNoData,
			/// The processed packet is not a SYN/SYNACK/FIN/RST packet and has no payload.
			/// Normally it's just a bare ACK packet.
			/// It's ignored and no callback function is called.
			Ignore_PacketWithNoData,

			/// The processed packet comes from a closed flow(an in-order FIN or RST is seen).
			/// It's ignored and no callback function is called.

			Ignore_PacketOfClosedFlow,
			/// The processed packet is a restransmission packet with no new data, meaning the `packet_sequence +
			/// packet_payload_length < sequence_expected`. It's ignored and no callback function is called.
			Ignore_Retransimission,
			/// The processed packet is not an IP packet.
			/// It's ignored and no callback function is called.
			NonIpPacket,
			/// The processed packet is not a TCP packet.
			/// It's ignored and no callback function is called.
			NonTcpPacket,
			/// The processed packet does not belong to any known TCP connection.
			/// It's ignored and no callback function is called.
			/// Normally this will be happen.
			Error_PacketDoesNotMatchFlow,
		};

		/// The type for storing the connection information
		using ConnectionInfoList = std::unordered_map<uint32_t, ConnectionData>;

		/// @typedef OnTcpMessageReady
		/// A callback invoked when new data arrives on a connection
		/// @param[in] side The side this data belongs to (MachineA->MachineB or vice versa). The value is 0 or 1 where
		/// 0 is the first side seen in the connection and 1 is the second side seen
		/// @param[in] tcpData The TCP data itself + connection information
		/// @param[in] userCookie A pointer to the cookie provided by the user in TcpReassembly c'tor (or nullptr if no
		/// cookie provided)
		using OnTcpMessageReady = std::function<void(int8_t side, const TcpStreamData& tcpData, void* userCookie)>;

		/// @typedef OnTcpConnectionStart
		/// A callback invoked when a new TCP connection is identified (whether it begins with a SYN packet or not)
		/// @param[in] connectionData Connection information
		/// @param[in] userCookie A pointer to the cookie provided by the user in TcpReassembly c'tor (or nullptr if no
		/// cookie provided)
		using OnTcpConnectionStart = std::function<void(const ConnectionData& connectionData, void* userCookie)>;

		/// @typedef OnTcpConnectionEnd
		/// A callback invoked when a TCP connection is terminated, either by a FIN or RST packet or manually by the
		/// user
		/// @param[in] connectionData Connection information
		/// @param[in] reason The reason for connection termination: FIN/RST packet or manually by the user
		/// @param[in] userCookie A pointer to the cookie provided by the user in TcpReassembly c'tor (or nullptr if no
		/// cookie provided)
		using OnTcpConnectionEnd =
		    std::function<void(const ConnectionData& connectionData, ConnectionEndReason reason, void* userCookie)>;

		/// A c'tor for this class
		/// @param[in] onMessageReadyCallback The callback to be invoked when new data arrives
		/// @param[in] userCookie A pointer to an object provided by the user. This pointer will be returned when
		/// invoking the various callbacks. This parameter is optional, default cookie is nullptr
		/// @param[in] onConnectionStartCallback The callback to be invoked when a new connection is identified. This
		/// parameter is optional
		/// @param[in] onConnectionEndCallback The callback to be invoked when a new connection is terminated (either by
		/// a FIN/RST packet or manually by the user). This parameter is optional
		/// @param[in] config Optional parameter for defining special configuration parameters. If not set the default
		/// parameters will be set
		explicit TcpReassembly(OnTcpMessageReady onMessageReadyCallback, void* userCookie = nullptr,
		                       OnTcpConnectionStart onConnectionStartCallback = nullptr,
		                       OnTcpConnectionEnd onConnectionEndCallback = nullptr,
		                       const TcpReassemblyConfiguration& config = TcpReassemblyConfiguration());

		/// The most important method of this class which gets a packet from the user and processes it. If this packet
		/// opens a new connection, ends a connection or contains new data on an existing connection, the relevant
		/// callback will be called (TcpReassembly#OnTcpMessageReady, TcpReassembly#OnTcpConnectionStart,
		/// TcpReassembly#OnTcpConnectionEnd)
		/// @param[in] tcpData A reference to the packet to process
		/// @return A enum of `TcpReassembly::ReassemblyStatus`, indicating status of TCP reassembly
		ReassemblyStatus reassemblePacket(Packet& tcpData);

		/// The most important method of this class which gets a raw packet from the user and processes it. If this
		/// packet opens a new connection, ends a connection or contains new data on an existing connection, the
		/// relevant callback will be invoked (TcpReassembly#OnTcpMessageReady, TcpReassembly#OnTcpConnectionStart,
		/// TcpReassembly#OnTcpConnectionEnd)
		/// @param[in] tcpRawData A reference to the raw packet to process
		/// @return A enum of `TcpReassembly::ReassemblyStatus`, indicating status of TCP reassembly
		ReassemblyStatus reassemblePacket(RawPacket* tcpRawData);

		/// Close a connection manually. If the connection doesn't exist or already closed an error log is printed. This
		/// method will cause the TcpReassembly#OnTcpConnectionEnd to be invoked with a reason of
		/// TcpReassembly#TcpReassemblyConnectionClosedManually
		/// @param[in] flowKey A 4-byte hash key representing the connection. Can be taken from a ConnectionData
		/// instance
		void closeConnection(uint32_t flowKey);

		/// Close all open connections manually. This method will cause the TcpReassembly#OnTcpConnectionEnd to be
		/// invoked for each connection with a reason of TcpReassembly#TcpReassemblyConnectionClosedManually
		void closeAllConnections();

		/// Get a map of all connections managed by this TcpReassembly instance (both connections that are open and
		/// those that are already closed)
		/// @return A map of all connections managed. Notice this map is constant and cannot be changed by the user
		const ConnectionInfoList& getConnectionInformation() const
		{
			return m_ConnectionInfo;
		}

		/// Check if a certain connection managed by this TcpReassembly instance is currently opened or closed
		/// @param[in] connection The connection to check
		/// @return A positive number (> 0) if connection is opened, zero (0) if connection is closed, and a negative
		/// number (< 0) if this connection isn't managed by this TcpReassembly instance
		int isConnectionOpen(const ConnectionData& connection) const;

		/// Clean up the closed connections from the memory
		/// @param[in] maxNumToClean The maximum number of items to be cleaned up per one call. This parameter, when its
		/// value is not zero, overrides the value that was set by the constructor.
		/// @return The number of cleared items
		uint32_t purgeClosedConnections(uint32_t maxNumToClean = 0);

	private:
		struct TcpFragment
		{
			uint32_t sequence = 0;
			size_t dataLength = 0;
			uint8_t* data = nullptr;
			std::chrono::time_point<std::chrono::high_resolution_clock> timestamp;

			~TcpFragment()
			{
				delete[] data;
			}
		};

		struct TcpOneSideData
		{
			PointerVector<TcpFragment> tcpFragmentList;
			uint32_t sequence = 0;
			uint16_t srcPort = 0;
			IPAddress srcIP;
			bool gotFinOrRst = false;
		};

		struct TcpReassemblyData
		{
			bool closed = false;
			int8_t numOfSides = 0;
			int8_t prevSide = -1;
			std::array<TcpOneSideData, 2> twoSides;
			ConnectionData connData;
		};

		class OutOfOrderProcessingGuard
		{
		private:
			bool& m_Flag;

		public:
			explicit OutOfOrderProcessingGuard(bool& flag) : m_Flag(flag)
			{
				m_Flag = true;
			}

			~OutOfOrderProcessingGuard()
			{
				m_Flag = false;
			}

			// Disable copy and move operations
			OutOfOrderProcessingGuard(const OutOfOrderProcessingGuard&) = delete;
			OutOfOrderProcessingGuard& operator=(const OutOfOrderProcessingGuard&) = delete;
		};

		using ConnectionList = std::unordered_map<uint32_t, TcpReassemblyData>;
		using CleanupMultiMap = std::multimap<time_t, uint32_t>;

		OnTcpMessageReady m_OnMessageReadyCallback;
		OnTcpConnectionStart m_OnConnStart;
		OnTcpConnectionEnd m_OnConnEnd;
		void* m_UserCookie;
		ConnectionList m_ConnectionList;
		ConnectionInfoList m_ConnectionInfo;
		CleanupMultiMap m_CleanupMultimap;
		bool m_RemoveConnInfo;
		uint32_t m_ClosedConnectionDelay;
		uint32_t m_MaxNumToClean;
		size_t m_MaxOutOfOrderFragments;
		time_t m_PurgeTimepoint;
		bool m_EnableBaseBufferClearCondition;
		bool m_ProcessingOutOfOrder = false;

		void checkOutOfOrderFragments(TcpReassemblyData* tcpReassemblyData, int8_t sideIndex, bool cleanWholeFragList);

		void handleFinOrRst(TcpReassemblyData* tcpReassemblyData, int8_t sideIndex, uint32_t flowKey, bool isRst);

		void closeConnectionInternal(uint32_t flowKey, ConnectionEndReason reason);

		void scheduleCleanup(uint32_t flowKey);
	};

}  // namespace pcpp