#ifndef MS_RTC_SCTP_STATE_COOKIE_HPP
#define MS_RTC_SCTP_STATE_COOKIE_HPP

#include "common.hpp"
#include "RTC/SCTP/association/Capabilities.hpp"
#include "RTC/SCTP/public/SctpTypes.hpp"
#include "RTC/Serializable.hpp"
#include "Utils.hpp"

namespace RTC
{
	namespace SCTP
	{
		/**
		 * This is the State Cookie we generate and put into a State Cookie
		 * parameter when we send INIT-ACK chunk to the remote peer.
		 *
		 * The syntax we use is as follows. Note that we use a fixed length of
		 * StateCookieLength bytes.
		 *
		 *  0                   1                   2                   3
		 *  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                            Magic 1                            |
		 * |                                                               |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                    Local Verification Tag                     |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                    Remote Verification Tag                    |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                      Local Initial TSN                        |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                      Remote Initial TSN                       |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |       Remote Advertised Receiver Window Credit (a_rwnd)       |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                            Tie-Tag                            |
		 * |                                                               |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * \                                                               \
		 * /                      Remote Capabilities                      /
		 * \                                                               \
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 *
		 * The Remote Capabilities are the raw capabilities announced by the remote
		 * endpoint in its INIT or INIT-ACK chunk (before being negotiated against
		 * our local options). Storing the raw remote capabilities (instead of the
		 * already negotiated ones) limits the effect of a tampered/forged cookie:
		 * when the COOKIE-ECHO is received they are re-negotiated against our local
		 * options, so an attacker can never enable an extension we didn't signal
		 * nor raise the stream limits above what we announced.
		 *
		 * Remote Capabilities are serialized as follows:
		 *
		 *  0                   1                   2                   3
		 *  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |   (Reserved)  |         |C|B|A|            Magic 2            |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |         Zero Checksum Alternate Error Detection Method        |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |      Max Outbound Streams     |       Max Inbound Streams     |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 *
		 * - Flag A (partialReliability): Partial Reliability Extension.
		 * - Flag B (messageInterleaving): Stream Schedulers and User Message
		 *     Interleaving (I-DATA).
		 * - Flag C (reConfig): Stream Reconfiguration.
		 * - Zero Checksum Alternate Error Detection Method: the method announced
		 *     by the remote endpoint (0 means none).
		 *
		 * When State Cookie authentication is enabled (see
		 * `SctpOptions::requireAuthenticatedCookie`), two extra fields are appended
		 * after the Remote Capabilities so the receiver can verify that the cookie
		 * was generated by itself and is not stale:
		 *
		 *  0                   1                   2                   3
		 *  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * |                       Creation Timestamp                      |
		 * |                                                               |
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 * \                                                               \
		 * /            MAC (HMAC-SHA1 over all preceding bytes)           /
		 * \                                                               \
		 * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
		 *
		 * The MAC is keyed with a per-association secret that never leaves the
		 * worker, making the cookie unforgeable. The Creation Timestamp (in ms) is
		 * used to detect stale cookies.
		 *
		 * @see RFC 9260 sections 5.1.3 and 5.1.4.
		 */
		class StateCookie : public Serializable
		{
		private:
			struct RemoteCapabilitiesField
			{
				uint8_t reserved;
#if defined(MS_LITTLE_ENDIAN)
				uint8_t bitA : 1;
				uint8_t bitB : 1;
				uint8_t bitC : 1;
				uint8_t unusedBits : 5;
#elif defined(MS_BIG_ENDIAN)
				uint8_t unusedBits : 5;
				uint8_t bitC : 1;
				uint8_t bitB : 1;
				uint8_t bitA : 1;
#endif
				uint16_t magic2;
				uint32_t zeroChecksumAlternateErrorDetectionMethod;
				uint16_t maxOutboundStreams;
				uint16_t maxInboundStreams;
			};

		public:
			// Fixed length of our generated State Cookies when authentication is
			// disabled.
			static constexpr size_t StateCookieLength{ 48 };
			// Offset in the State Cookie where the Remote Capabilities are located.
			static constexpr size_t RemoteCapabilitiesOffset{ 36 };
			// Magic value we prefix the State Cookie with. Note that it is
			// "msworker" in ASCII bytes.
			static constexpr uint64_t Magic1{ 0x6D73776F726B6572 };
			static constexpr size_t Magic1Length{ 8 };
			// Magic value used within the Remote Capabilities block.
			static constexpr uint16_t Magic2{ 0xAD81 };
			// Offset of the creation timestamp present in authenticated State
			// Cookies.
			static constexpr size_t TimestampOffset{ StateCookie::StateCookieLength };
			// Length of the creation timestamp field (a uint64_t with the time in
			// milliseconds).
			static constexpr size_t TimestampLength{ 8 };
			// Offset of the MAC present in authenticated State Cookies.
			static constexpr size_t MacOffset{ StateCookie::TimestampOffset + StateCookie::TimestampLength };
			// Length of the MAC. We use HMAC-SHA1 so it's 20 bytes.
			static constexpr size_t MacLength{ 20 };
			// Fixed length of our generated State Cookies when authentication is
			// enabled.
			static constexpr size_t AuthenticatedStateCookieLength{ StateCookie::MacOffset +
			                                                        StateCookie::MacLength };
			// State Cookie lifespan (Valid.Cookie.Life) in microseconds. Used to
			// reject stale authenticated cookies.
			//
			// @see RFC 9260 section 16.
			static constexpr int64_t ValidCookieLifeUs{ 60000 * 1000 };

		public:
			/**
			 * Parse a StateCookie supposely generated by mediasoup.
			 *
			 * @remarks
			 * `bufferLength` must be the exact length of the State Cookie.
			 */
			static StateCookie* Parse(const uint8_t* buffer, size_t bufferLength);

			/**
			 * Create a StateCookie.
			 *
			 * @remarks
			 * - `bufferLength` could be greater than the real length of the State
			 *   cookie.
			 * - If `macKey` is not nullptr, an authenticated cookie (with creation
			 *   timestamp and MAC) is generated. Otherwise a plain cookie is
			 *   generated and `creationTimestampUs`/`macKey` are ignored.
			 */
			static StateCookie* Factory(
			  uint8_t* buffer,
			  size_t bufferLength,
			  uint32_t localVerificationTag,
			  uint32_t remoteVerificationTag,
			  uint32_t localInitialTsn,
			  uint32_t remoteInitialTsn,
			  uint32_t remoteAdvertisedReceiverWindowCredit,
			  uint64_t tieTag,
			  const Capabilities& remoteCapabilities,
			  int64_t creationTimestampUs = 0,
			  const uint8_t* macKey       = nullptr,
			  size_t macKeyLength         = 0);

			/**
			 * Serialize a StateCookie (based on given arguments) in the given buffer.
			 *
			 * @remarks
			 * - If `macKey` is not nullptr, an authenticated cookie (with creation
			 *   timestamp and MAC) is generated. Otherwise a plain cookie is
			 *   generated and `creationTimestampUs`/`macKey` are ignored.
			 */
			static void Write(
			  uint8_t* buffer,
			  size_t bufferLength,
			  uint32_t localVerificationTag,
			  uint32_t remoteVerificationTag,
			  uint32_t localInitialTsn,
			  uint32_t remoteInitialTsn,
			  uint32_t remoteAdvertisedReceiverWindowCredit,
			  uint64_t tieTag,
			  const Capabilities& remoteCapabilities,
			  int64_t creationTimestampUs = 0,
			  const uint8_t* macKey       = nullptr,
			  size_t macKeyLength         = 0);

			/**
			 * Whether the given buffer is a StateCookie generated by mediasoup.
			 *
			 * @remarks
			 * - This only checks magic values and length. It does NOT verify the
			 *   MAC of authenticated cookies (which requires the per-association
			 *   secret). Use `VerifyMac()` for that.
			 */
			static bool IsMediasoupStateCookie(const uint8_t* buffer, size_t bufferLength);

			/**
			 * Verify the MAC of an authenticated StateCookie generated by mediasoup.
			 *
			 * @remarks
			 * - Returns false if the buffer is not an authenticated cookie (i.e. its
			 *   length is not `AuthenticatedStateCookieLength`) or if the MAC doesn't
			 *   match the one computed with the given secret key.
			 *
			 * @see RFC 9260 section 5.1.4.
			 */
			static bool VerifyMac(
			  const uint8_t* buffer, size_t bufferLength, const uint8_t* macKey, size_t macKeyLength);

			/**
			 * Determine the SCTP implementation of the generator of State Cookie
			 * given in the buffer.
			 */
			static Types::SctpImplementation DetermineSctpImplementation(
			  const uint8_t* buffer, size_t bufferLength);

		public:
			StateCookie(uint8_t* buffer, size_t bufferLength);

			~StateCookie() override;

			void Dump(int indentation = 0) const final;

			StateCookie* Clone(uint8_t* buffer, size_t bufferLength) const final;

			/**
			 * The value of the Initiate Tag field we put in our INIT or INIT-ACK
			 * chunk. Packets sent by the remote peer must include this value in
			 * their Verification Tag field.
			 */
			uint32_t GetLocalVerificationTag() const
			{
				return Utils::Byte::Get4Bytes(GetBuffer(), 8);
			}

			/**
			 * The value of the Initiate Tag field the peer put in its INIT or
			 * INIT-ACK chunk. Packets sent by us to the peer must include this value
			 * in their Verification Tag field.
			 */
			uint32_t GetRemoteVerificationTag() const
			{
				return Utils::Byte::Get4Bytes(GetBuffer(), 12);
			}

			/**
			 * The value of the Initial TSN field we put in our INIT or INIT-ACK
			 * chunk.
			 */
			uint32_t GetLocalInitialTsn() const
			{
				return Utils::Byte::Get4Bytes(GetBuffer(), 16);
			}

			/**
			 * The value of the Initial TSN field the peer put in its INIT or
			 * INIT-ACK chunk.
			 */
			uint32_t GetRemoteInitialTsn() const
			{
				return Utils::Byte::Get4Bytes(GetBuffer(), 20);
			}

			/**
			 * The value of the Advertised Receiver Window Credit field we put in our
			 * INIT or INIT-ACK chunk.
			 */
			uint32_t GetRemoteAdvertisedReceiverWindowCredit() const
			{
				return Utils::Byte::Get4Bytes(GetBuffer(), 24);
			}

			/**
			 * Tie-Tag used as a nonce when connecting.
			 */
			uint64_t GetTieTag() const
			{
				return Utils::Byte::Get8Bytes(GetBuffer(), 28);
			}

			/**
			 * Raw capabilities announced by the remote endpoint (before being
			 * negotiated against our local options).
			 */
			Capabilities GetRemoteCapabilities() const;

			/**
			 * Whether this is an authenticated StateCookie (i.e. it carries a
			 * creation timestamp and a MAC).
			 */
			bool IsAuthenticated() const
			{
				return GetLength() == StateCookie::AuthenticatedStateCookieLength;
			}

			/**
			 * The time (in microseconds) at which this StateCookie was created.
			 *
			 * @remarks
			 * - Only meaningful in authenticated cookies (see `IsAuthenticated()`).
			 */
			int64_t GetCreationTimestampUs() const
			{
				return static_cast<int64_t>(Utils::Byte::Get8Bytes(GetBuffer(), StateCookie::TimestampOffset));
			}

		private:
			RemoteCapabilitiesField* GetRemoteCapabilitiesField() const
			{
				return reinterpret_cast<RemoteCapabilitiesField*>(
				  const_cast<uint8_t*>(GetBuffer()) + StateCookie::RemoteCapabilitiesOffset);
			}
		};
	} // namespace SCTP
} // namespace RTC

#endif