#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 <callback.h>
#include <spinLock.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<smo::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 sendHeartbeat();
	void onHeartbeatTimer(const boost::system::error_code& error);
	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);

	class ConnectReq;
	class ConnectToKnownDeviceReq;
	class ConnectByDeviceIdentifierReq;
	class ExecuteHandshakeReq;
	class DisconnectReq;
	class EnablePcloudDataReq;
	class DisablePcloudDataReq;
	class SetReturnModeReq;
	class GetReturnModeReq;

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);

	// Callback function type definitions for async methods
	typedef std::function<void(bool success)> connectReqCbFn;
	typedef std::function<
		void(bool success, const std::string& ipAddr)>
		connectToKnownDeviceReqCbFn;
	typedef std::function<
		void(bool success, const std::string& ipAddr)>
		connectByDeviceIdentifierReqCbFn;
	typedef std::function<void(bool success)> executeHandshakeReqCbFn;
	typedef std::function<void(bool success)> disconnectReqCbFn;
	typedef std::function<void(bool success)> enablePcloudDataReqCbFn;
	typedef std::function<void(bool success)> disablePcloudDataReqCbFn;
	typedef std::function<void(bool success)> setReturnModeReqCbFn;
	typedef std::function<void(bool success, uint8_t returnMode)>
		getReturnModeReqCbFn;

	// Async connection methods
	void connectReq(smo::Callback<connectReqCbFn> callback);
	void connectToKnownDeviceReq(
		smo::Callback<connectToKnownDeviceReqCbFn> callback);
	void connectByDeviceIdentifierReq(
		smo::Callback<connectByDeviceIdentifierReqCbFn> callback);
	void executeHandshakeReq(
		const std::string& deviceIP,
		smo::Callback<executeHandshakeReqCbFn> callback);
	void disconnectReq(smo::Callback<disconnectReqCbFn> callback);
	void enablePcloudDataReq(smo::Callback<enablePcloudDataReqCbFn> callback);
	void disablePcloudDataReq(smo::Callback<disablePcloudDataReqCbFn> callback);
	void setReturnModeReq(
		uint8_t returnMode, smo::Callback<setReturnModeReqCbFn> callback);
	void getReturnModeReq(smo::Callback<getReturnModeReqCbFn> callback);

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

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

	// Heartbeat state
	std::unique_ptr<boost::asio::deadline_timer> heartbeatTimer;
	std::atomic<bool> heartbeatActive;
	smo::SpinLock heartbeatActiveLock;

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