#ifndef LIVOX_PROTO1_DEVICE_H
#define LIVOX_PROTO1_DEVICE_H

#include <string>
#include <cstdint>
#include <memory>
#include <atomic>
#include <optional>
#include <functional>
#include <unordered_map>
#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>

// 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 handshakeTimeoutMs, int retryDelayMs,
		const std::string& smoIp, uint8_t smoSubnetNbits,
		uint16_t dataPort, uint16_t cmdPort, uint16_t imuPort);
	~Device();

public:
	comms::DiscoveredDevice discoveredDevice;

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

	// Static collection for devices being constructed (not yet in DeviceManager)
	// Maps device IP to list of command-specific UDP handlers for that device
	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;

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;

public:
	// 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;

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

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

	// Point cloud data state
	std::unique_ptr<boost::asio::posix::stream_descriptor> pcloudDataSocketDesc;
	std::atomic<bool> pcloudDataActive;
	int pcloudDataFd;  // Socket file descriptor for point cloud data reception

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
	bool setupPcloudDataSocket();
	void cleanupPcloudDataSocket();

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

} // namespace livoxProto1

#endif // LIVOX_PROTO1_DEVICE_H