#ifndef LIVOX_PROTO1_DEVICE_H
#define LIVOX_PROTO1_DEVICE_H

#include <boostAsioLinkageFix.h>

#include <string>
#include <cstdint>
#include <cstddef>
#include <memory>
#include <atomic>
#include <optional>
#include <functional>
#include <unordered_map>
#include <stdexcept>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <boost/asio/deadline_timer.hpp>
#include <boost/asio/posix/stream_descriptor.hpp>
#include "protocol.h"
#include <spinscale/co/dynamicPostingInvoker.h>
#include <spinscale/co/invokers.h>
#include <spinscale/co/nonViralTaskNursery.h>
#include <spinscale/syncCancelerForAsyncWork.h>

// Custom hash function for std::pair<uint8_t, uint8_t>
namespace std {
	template<>
	struct hash<std::pair<uint8_t, uint8_t>> {
		size_t operator()(const std::pair<uint8_t, uint8_t>& p) const noexcept {
			return (static_cast<size_t>(p.first) << 8) | static_cast<size_t>(p.second);
		}
	};
}

// Forward declaration
namespace smo {
	class ComponentThread;
}

namespace livoxProto1 {
namespace comms {

/**		EXPLANATION:
 * This class represents a discovered device. It is used to store the
 * device identifier and IP address of a discovered device.
 */
class DiscoveredDevice
{
public:
	DiscoveredDevice(
		const std::string &deviceIdentifier,
		DeviceType deviceType,
		const std::string &ipAddr);

	// "Conversion" constructor from BroadcastMessage
	DiscoveredDevice(const BroadcastMessage &msg, const std::string &ipAddr);

	~DiscoveredDevice() = default;

	bool operator==(const DiscoveredDevice &other) const
	{
		return comms::deviceIdentifiersEqual(
			deviceIdentifier, other.deviceIdentifier);
	}

	std::string stringify(void) const;
	std::string getDeviceTypeName(void) const;

public:
	std::string deviceIdentifier;
	DeviceType deviceType;
	std::string ipAddr;
};

} // namespace comms

class Device
{
public:
	Device(const std::string &deviceIdentifier,
		const std::shared_ptr<sscl::ComponentThread>& componentThread,
		int commandTimeoutMs, int retryDelayMs,
		const std::string& smoIp, uint8_t smoSubnetNbits,
		uint16_t dataPort, uint16_t cmdPort, uint16_t imuPort);
	~Device();

private:
	// Heartbeat mechanism
	void startHeartbeat();
	void stopHeartbeat();
	void sendHeartbeatOnce();

	/**	EXPLANATION:
	 * deviceCDaemon is a dynamic posting non-viral coroutine: startHeartbeat()
	 * passes ExplicitPostTarget{componentThread->getIoContext()} so the daemon
	 * body always runs on componentThread. Extensible for future per-device
	 * background work beyond heartbeats.
	 */
	sscl::co::DynamicNonViralPostingInvoker deviceCDaemon(
		sscl::co::ExplicitPostTarget postTarget,
		std::exception_ptr &exceptionPtr,
		std::function<void()> callback,
		sscl::SyncCancelerForAsyncWork &canceler);

	std::string generateClientDeviceIpFromSerialNumber(
		const std::string& broadcastCode);

	// IP detection methods
	std::optional<std::string> detectSmoIp(const std::string& deviceIP);
	uint32_t getSubnetMaskFor(uint8_t nbits);

public:
	enum class ReturnMode : uint8_t
	{
		SingleFirst		= 0x00,
		SingleStrongest	= 0x01,
		Dual			= 0x02,
		Triple			= 0x03
	};

	/**
	 * Get the number of points per datagram based on return mode
	 * @param returnMode The return mode (0=SingleFirst, 1=SingleStrongest, 2=Dual, 3=Triple)
	 * @return Number of points per datagram
	 */
	static inline size_t getNPointsPerDgram(int returnMode)
	{
		/*
		 * Map modes to points per datagram based on Livox docs
		 * 1: first, 2: strongest -> 96 samples => 96 points
		 * 3: dual -> 48 samples * 2 points = 96
		 * 4: triple -> 30 samples * 3 points = 90
		 */
		switch (returnMode)
		{
		case static_cast<int>(ReturnMode::SingleFirst):
		case static_cast<int>(ReturnMode::SingleStrongest):
		case static_cast<int>(ReturnMode::Dual):
			return 96u;
		case static_cast<int>(ReturnMode::Triple):
			return 90u;
		default:
			throw std::runtime_error(
				std::string(__func__) + ": Unknown returnMode "
				+ std::to_string(returnMode));
		}
	}

	// Utility methods
	std::optional<std::string> getSmoIp(const std::string& deviceIP);

	struct ConnectIpResult
	{
		bool success = false;
		std::string ipAddr;
	};

	// Async connection methods
	sscl::co::ViralNonPostingInvoker<bool> connectCReq();
	sscl::co::ViralNonPostingInvoker<ConnectIpResult> connectToKnownDeviceCReq();
	sscl::co::ViralNonPostingInvoker<ConnectIpResult> connectByDeviceIdentifierCReq();
	sscl::co::ViralNonPostingInvoker<bool> executeHandshakeCReq(
		const std::string& deviceIP);
	sscl::co::ViralNonPostingInvoker<bool> disconnectCReq();
	sscl::co::ViralNonPostingInvoker<bool> enablePcloudDataCReq();
	sscl::co::ViralNonPostingInvoker<bool> disablePcloudDataCReq();
	
	struct GetReturnModeResult
	{
		bool success = false;
		uint8_t returnMode = 0;
	};

	sscl::co::ViralNonPostingInvoker<bool> setReturnModeCReq(uint8_t returnMode);
	sscl::co::ViralNonPostingInvoker<GetReturnModeResult> getReturnModeCReq();

public:
	comms::DiscoveredDevice discoveredDevice;
	std::atomic<size_t> nAttachedStimulusProducers;

	// Configuration
	std::shared_ptr<sscl::ComponentThread> componentThread;
	int commandTimeoutMs, retryDelayMs;
	std::string smoIp;
	std::string detectedSmoListeningIp;
	uint8_t smoSubnetNbits;
	uint16_t dataPort, cmdPort, imuPort;

	// Heartbeat state (timer lifetime tied to Device ctor/dtor)
	boost::asio::deadline_timer heartbeatTimer;
	sscl::co::NonViralTaskNursery daemonNursery;

	// Point cloud data state
	std::atomic<bool> pcloudDataActive;

	// Cached last-known return mode for this device
	ReturnMode currentReturnMode = ReturnMode::SingleFirst;

public:
	// UDP datagram handling
	void handleUdpDgram(
		const uint8_t* data, ssize_t bytesReceived,
		const struct sockaddr_in& senderAddr);

	// Command handler registration
	void registerUdpCommandHandler(
		uint8_t cmd_set, uint8_t cmd_id,
		std::function<void(
			const uint8_t* data, ssize_t bytesReceived,
			const struct sockaddr_in& senderAddr)> handler,
		const std::string& deviceIP = "");

	void unregisterUdpCommandHandler(
		uint8_t cmd_set, uint8_t cmd_id, const std::string& deviceIP = "");

private:
	// Point cloud data setup
	void cleanupPcloudDataSocket();

	/**		EXPLANATION:
	 * This is the "straightforward" map of command set and command id to
	 * handlers. This is useful for any commands which are guaranteed to be
	 * issued to the device *AFTER* the device has successfully been added
	 * to the DeviceManager's list of devices.
	 *
	 * I.e: it cannot be used for commands which are issued to the device before
	 * getOrCreateDevice() has added the device to the DeviceManager's list of
	 * devices.
	 */
	// Command handler map
	std::unordered_map<
		std::pair<uint8_t, uint8_t>,
		std::function<void(
			const uint8_t* data, ssize_t bytesReceived,
			const struct sockaddr_in& senderAddr)>> udpCommandHandlers;

public:
	/**		EXPLANATION:
	 * This is the "temporary" map of command set and command id to
	 * handlers. This is useful for any commands which are issued to the device
	 * while it is being constructed.
	 *
	 * I.e: it shouldn't be used for cmds which are issued to the device after
	 * getOrCreateDevice() has added the device to the DeviceManager's list of
	 * devices. It will work for such commands, but we'd kind of prefer to use
	 * the "straightforward" map above for such commands.
	 *
	 *	NOTE:
	 * There's a strong argument to be made for just getting rid of the
	 * "straightforward" map above and just using this one, tho.
	 */
	struct CommandHandler {
		uint8_t cmd_set;
		uint8_t cmd_id;
		std::function<void(
			const uint8_t* data, ssize_t bytesReceived,
			const struct sockaddr_in& senderAddr)> handler;
	};
	static std::unordered_map<std::string, std::vector<CommandHandler>>
		devicesUnderConstruction;
};

} // namespace livoxProto1

#endif // LIVOX_PROTO1_DEVICE_H