// 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>
ZEN_THIRD_PARTY_INCLUDES_END

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

namespace zen::asio_http {

/**
 * WebSocket connection over an ASIO TCP socket
 *
 * Owns the TCP 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<WsAsioConnection> to keep the
 * connection alive for the duration of the async operation.  The service layer
 * also holds a Ref<WebSocketConnection>.
 */
class WsAsioConnection : public WebSocketConnection
{
public:
	WsAsioConnection(std::unique_ptr<asio::ip::tcp::socket> Socket, IWebSocketHandler& Handler);
	~WsAsioConnection() 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<asio::ip::tcp::socket> m_Socket;
	IWebSocketHandler&					   m_Handler;
	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};
	bool			  m_CloseSent = false;
};

}  // namespace zen::asio_http