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

#pragma once

#include <zenhttp/websocket.h>

#include <zencore/thread.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <asio.hpp>
#if defined(ASIO_HAS_LOCAL_SOCKETS)
#	include <asio/local/stream_protocol.hpp>
#endif
#if ZEN_USE_OPENSSL
#	include <asio/ssl.hpp>
#endif
ZEN_THIRD_PARTY_INCLUDES_END

#include <deque>
#include <memory>
#include <vector>

namespace zen {
class HttpServer;
}  // namespace zen

namespace zen::asio_http {

/**
 * WebSocket connection over an ASIO stream socket
 *
 * Templated on SocketType to support both TCP and Unix domain sockets.
 * Owns the socket (moved from HttpServerConnection after the 101 handshake)
 * and runs an async read/write loop to exchange WebSocket frames.
 *
 * Lifetime is managed solely through intrusive reference counting (RefCounted).
 * The async read/write callbacks capture Ref<> to keep the connection alive
 * for the duration of the async operation.  The service layer also holds a
 * Ref<WebSocketConnection>.
 */
template<typename SocketType>
class WsAsioConnectionT : public WebSocketConnection
{
public:
	WsAsioConnectionT(std::unique_ptr<SocketType> Socket, IWebSocketHandler& Handler, HttpServer* Server);
	~WsAsioConnectionT() override;

	/**
	 * Start the async read loop. Must be called once after construction
	 * and the 101 response has been sent.
	 */
	void Start();

	// WebSocketConnection interface
	void SendText(std::string_view Text) override;
	void SendBinary(std::span<const uint8_t> Data) override;
	void Close(uint16_t Code, std::string_view Reason) override;
	bool IsOpen() const override;

private:
	void EnqueueRead();
	void OnDataReceived(const asio::error_code& Ec, std::size_t ByteCount);
	void ProcessReceivedData();

	void EnqueueWrite(std::vector<uint8_t> Frame);
	void FlushWriteQueue();
	void OnWriteComplete(const asio::error_code& Ec, std::size_t ByteCount);

	void DoClose(uint16_t Code, std::string_view Reason);

	std::unique_ptr<SocketType> m_Socket;
	IWebSocketHandler&			m_Handler;
	zen::HttpServer*			m_HttpServer;
	asio::streambuf				m_ReadBuffer;

	RwLock							 m_WriteLock;
	std::deque<std::vector<uint8_t>> m_WriteQueue;
	bool							 m_IsWriting = false;

	std::atomic<bool> m_IsOpen{true};
	std::atomic<bool> m_CloseSent{false};
};

using WsAsioConnection = WsAsioConnectionT<asio::ip::tcp::socket>;

#if defined(ASIO_HAS_LOCAL_SOCKETS)
using WsAsioUnixConnection = WsAsioConnectionT<asio::local::stream_protocol::socket>;
#endif

#if ZEN_USE_OPENSSL
using WsAsioSslConnection = WsAsioConnectionT<asio::ssl::stream<asio::ip::tcp::socket>>;
#endif

}  // namespace zen::asio_http