#ifndef UDP_COMMAND_DEMUXER_H
#define UDP_COMMAND_DEMUXER_H

#include <boostAsioLinkageFix.h>
#include <atomic>
#include <memory>
#include <boost/asio/posix/stream_descriptor.hpp>
#include <componentThread.h>
#include <spinLock.h>

namespace livoxProto1 {

// Forward declarations
class DeviceManager;

namespace comms {

/**
 * UdpCommandDemuxer - Routes UDP command datagrams to appropriate devices
 *
 * This class listens on the command port (65000) for incoming UDP datagrams
 * from Livox devices and routes them to the appropriate Device based on
 * the source IP address.
 *
 * The reason we need a whole class for this is because we use the same port
 * numbers for all connected devices, so we have no way to distinguish between
 * devices except based on the devices' IP addrs. Since all commands are sent
 * over UDP, our sockets don't have built-in binding to a specific source IP.
 *
 * So we need to discriminate between source IPs manually, and demultiplex
 * the dgrams received from different devices manually.
 *
 * We'll prolly also have to do the same thing for point cloud and IMU data, so
 * we'll prolly end up renaming this class to UdpResponseDemuxer.
 */
class UdpCommandDemuxer
{
public:
	UdpCommandDemuxer(
		const std::shared_ptr<smo::ComponentThread>& componentThread,
		DeviceManager& deviceManager,
		uint16_t commandPort = 56001,
		uint16_t dataPort = 56000);

	~UdpCommandDemuxer();

	void start();
	void stop();
	bool isRunning() const { return isActive.load(); }

	// Get shared pointer to command endpoint for handshake use
	std::shared_ptr<boost::asio::posix::stream_descriptor>
	getCmdEndpointFdDesc() const
	{
		return cmdEndpointFdDesc;
	}

	// Get shared pointer to pcloud data fd for use in IoUringAssemblyEngine
	std::shared_ptr<boost::asio::posix::stream_descriptor>
	getPcloudDataFdDesc() const
	{
		return pcloudDataFdDesc;
	}

private:
	// Socket and async objects
	std::shared_ptr<boost::asio::posix::stream_descriptor> pcloudDataFdDesc;
	// Socket and async objects
	std::shared_ptr<boost::asio::posix::stream_descriptor> cmdEndpointFdDesc;

private:
	void setupSockets();
	void setupCommandSocket();
	void setupPcloudDataSocket();
	void startAsyncReceive();
	void onDataReady(const boost::system::error_code& error);
	void processIncomingData();

	std::shared_ptr<smo::ComponentThread> componentThread;
	DeviceManager& deviceManager;
	uint16_t commandPort;
	uint16_t dataPort;

	// State management
	smo::SpinLock isActiveAndShouldStopLock;
	std::atomic<bool> isActive{false};
	std::atomic<bool> shouldStop{false};

	// Receive buffer
	uint8_t receiveBuffer[1024];
	struct sockaddr_in senderAddr;
	socklen_t senderAddrLen;
	ssize_t bytesReceived;
};

} // namespace comms
} // namespace livoxProto1

#endif // UDP_COMMAND_DEMUXER_H