// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "hordetransport.h"

#include <cstdint>
#include <memory>
#include <mutex>
#include <vector>

namespace zen::horde {

/** AES-256-GCM encrypted transport wrapper.
 *
 *  Wraps an inner ComputeTransport, encrypting all outgoing data and decrypting
 *  all incoming data using AES-256-GCM. The nonce is mutated per message using
 *  the Horde nonce mangling scheme: n32[0]++; n32[1]--; n32[2] = n32[0] ^ n32[1].
 *
 *  Wire format per encrypted message:
 *    [plaintext length (4B little-endian)][nonce (12B)][ciphertext][GCM tag (16B)]
 *
 *  Uses BCrypt on Windows and OpenSSL EVP on Linux/macOS (selected at compile time).
 */
class AesComputeTransport final : public ComputeTransport
{
public:
	AesComputeTransport(const uint8_t (&Key)[KeySize], std::unique_ptr<ComputeTransport> InnerTransport);
	~AesComputeTransport() override;

	bool   IsValid() const override;
	size_t Send(const void* Data, size_t Size) override;
	size_t Recv(void* Data, size_t Size) override;
	void   MarkComplete() override;
	void   Close() override;

private:
	static constexpr size_t NonceBytes = 12;  ///< AES-GCM nonce size
	static constexpr size_t TagBytes   = 16;  ///< AES-GCM authentication tag size

	struct CryptoContext;

	std::unique_ptr<CryptoContext>	  m_Crypto;
	std::unique_ptr<ComputeTransport> m_Inner;
	std::vector<uint8_t>			  m_EncryptBuffer;
	std::vector<uint8_t>			  m_RemainingData;	///< Buffered decrypted data from a partially consumed Recv
	size_t							  m_RemainingOffset = 0;
	std::mutex						  m_Lock;
	bool							  m_IsClosed = false;
};

}  // namespace zen::horde