#include <boostAsioLinkageFix.h>

#include <array>
#include <algorithm>
#include <cstring>
#include <iostream>
#include <functional>
#include <opts.h>
#include <componentThread.h>
#include <adapters/boostAsio/udpReceiveFromAReq.h>
#include <spinscale/co/nonViralCompletion.h>
#include "broadcastListener.h"
#include "core.h"
#include "protocol.h"

namespace livoxProto1 {
namespace comms {

namespace {

bool isShutdownReceiveError(const boost::system::error_code &ec)
{
	return ec == boost::asio::error::operation_aborted
		|| ec == boost::asio::error::bad_descriptor;
}

void logEventHandlerException(std::exception_ptr &exceptionPtr)
{
	sscl::co::NonViralCompletion nvc(exceptionPtr);
	if (!nvc.hasException()) {
		return;
	}

	try {
		nvc.checkAndRethrowException();
	} catch (const std::exception &e) {
		std::cerr << "BroadcastListener: event handler: "
			<< e.what() << std::endl;
	}
}

void logReceiveDaemonException(std::exception_ptr &exceptionPtr)
{
	sscl::co::NonViralCompletion nvc(exceptionPtr);
	if (!nvc.hasException()) {
		return;
	}

	try {
		nvc.checkAndRethrowException();
	} catch (const std::exception &e) {
		std::cerr << "BroadcastListener: receive daemon: "
			<< e.what() << std::endl;
	}
}

bool parseAndValidateBroadcastMessage(
	const std::array<uint8_t, UDP_BCAST_MSG_BUFFER_NBYTES> &bytes,
	std::size_t nbytes,
	BroadcastMessage &msgOut)
{
	if (nbytes < sizeof(BroadcastMessage))
	{
		std::cerr << "broadcastMsgIndCInd"
			<< ": Received packet too small: " << nbytes
			<< " bytes (expected at least "
			<< sizeof(BroadcastMessage) << ")" << std::endl;
		return false;
	}

	std::memcpy(&msgOut, bytes.data(), sizeof(BroadcastMessage));

	// Following the clean receiving flow:
	// 1. Swap CRC32 to host endianness first
	msgOut.footer.swapCrc32ToHostEndianness();
	// 2. Validate CRC32 (on whole message excluding footer CRC32 field)
	if (!msgOut.validateCrc32())
	{
		std::cerr << "broadcastMsgIndCInd"
			<< ": Broadcast message failed CRC32 validation"
			<< std::endl;
		return false;
	}

	// 3. Swap CRC16 to host endianness
	msgOut.header.swapCrc16ToHostEndianness();
	// 4. Validate CRC16 (on header only)
	if (!msgOut.header.validateCrc16())
	{
		std::cerr << "broadcastMsgIndCInd"
			<< ": Broadcast message failed CRC16 validation"
			<< std::endl;
		return false;
	}

	// 5. Swap content to host endianness
	msgOut.swapContentsToHostEndianness();
	// 6. Validate message sanity
	if (!msgOut.sanityCheck())
	{
		std::cerr << "broadcastMsgIndCInd"
			<< ": Broadcast message failed sanity check"
			<< std::endl;
		return false;
	}

	return true;
}

/** RAII: open eventHandlerNursery at CDaemon entry; drain at exit. */
class EventHandlerNurseryScopeGuard
{
public:
	explicit EventHandlerNurseryScopeGuard(
		sscl::co::NonViralTaskNursery &nursery)
	: nursery(nursery)
	{
		nursery.openAdmission();
	}

	~EventHandlerNurseryScopeGuard()
	{
		nursery.closeAdmission();
		nursery.syncAwaitAllSettlements(
			sscl::ComponentThread::getSelf()->getIoContext());
	}

	EventHandlerNurseryScopeGuard(
		const EventHandlerNurseryScopeGuard &) = delete;
	EventHandlerNurseryScopeGuard &operator=(
		const EventHandlerNurseryScopeGuard &) = delete;

private:
	sscl::co::NonViralTaskNursery &nursery;
};

} // namespace

BroadcastListener::BroadcastIndPayload::BroadcastIndPayload(
	const uint8_t *receiveBytes, std::size_t receiveNbytes,
	const boost::asio::ip::udp::endpoint &endpoint)
:	nbytes(receiveNbytes),
senderEndpoint(endpoint)
{
	if (receiveNbytes > 0) {
		std::memcpy(bytes.data(), receiveBytes, receiveNbytes);
	}
}

BroadcastListener::BroadcastIndPayload::BroadcastIndPayload(
	BroadcastIndPayload &&other) noexcept
:	nbytes(other.nbytes),
senderEndpoint(std::move(other.senderEndpoint))
{
	if (nbytes > 0) {
		std::memcpy(bytes.data(), other.bytes.data(), nbytes);
	}
	other.nbytes = 0;
}

BroadcastListener::BroadcastListener(
	const std::shared_ptr<sscl::ComponentThread>& componentThread,
	uint16_t listeningPort, uint16_t connectPort
)
: componentThread(componentThread),
listeningPort(listeningPort),
connectPort(connectPort),
deviceGoneAwayCb(nullptr),
socket(componentThread->getIoContext()),
listeningEndpoint(boost::asio::ip::udp::v4(), listeningPort)
{
}

std::shared_ptr<DiscoveredDevice>
BroadcastListener::getDevice(const std::string &deviceIdentifier) const
{
	sscl::SpinLock::Guard lock(discoveredDevices.lock);

	auto it = std::find_if(
		discoveredDevices.rsrc.devices.begin(),
		discoveredDevices.rsrc.devices.end(),
		[&deviceIdentifier](const std::shared_ptr<DiscoveredDevice>& device) {
			return comms::deviceIdentifiersEqual(
				device->deviceIdentifier, deviceIdentifier);
		}
	);

	return it != discoveredDevices.rsrc.devices.end() ? *it : nullptr;
}

void BroadcastListener::registerDiscoveredDevice(
	const BroadcastMessage &msg, const std::string &senderIp)
{
	std::string broadcastCode(
		reinterpret_cast<const char*>(msg.broadcast_code));

	sscl::SpinLock::Guard lock(discoveredDevices.lock);

	// Early return if device already exists
	const auto existingIt = std::find_if(
		discoveredDevices.rsrc.devices.begin(),
		discoveredDevices.rsrc.devices.end(),
		[&broadcastCode](const std::shared_ptr<DiscoveredDevice> &device) {
			return comms::deviceIdentifiersEqual(
				device->deviceIdentifier, broadcastCode);
		});

	if (existingIt != discoveredDevices.rsrc.devices.end())
	{
		if (getProtoState().smoCallbacks.OptionParser_getOptions().verbose)
		{
			std::cout << "broadcastMsgIndCInd"
				<< ": Received broadcast from known device: "
				<< broadcastCode << " at " << senderIp << "\n";
		}
		return;
	}

	// Create new DiscoveredDevice using conversion constructor
	auto device = std::make_shared<DiscoveredDevice>(msg, senderIp);
	discoveredDevices.rsrc.devices.push_back(device);
	std::cout << "broadcastMsgIndCInd: Discovered new Livox device: "
		<< device->stringify() << "\n";
}

sscl::co::NonViralNonPostingInvoker
BroadcastListener::broadcastMsgIndCInd(
	[[maybe_unused]] std::exception_ptr &exceptionPtr,
	[[maybe_unused]] std::function<void()> callback,
	[[maybe_unused]] sscl::SyncCancelerForAsyncWork &canceler,
	BroadcastIndPayload payload)
{
	BroadcastMessage msg;
	if (!parseAndValidateBroadcastMessage(
		payload.bytes, payload.nbytes, msg)) {
		co_return;
	}

	const std::string senderIp = payload.senderEndpoint.address().to_string();
	registerDiscoveredDevice(msg, senderIp);

	co_return;
}

sscl::co::DynamicNonViralPostingInvoker
BroadcastListener::broadcastReceiveCDaemon(
	[[maybe_unused]] sscl::co::ExplicitPostTarget postTarget,
	[[maybe_unused]] std::exception_ptr &exceptionPtr,
	[[maybe_unused]] std::function<void()> callback,
	sscl::SyncCancelerForAsyncWork &canceler)
{
	EventHandlerNurseryScopeGuard eventHandlerScope(eventHandlerNursery);

	boost::asio::io_context &ioContext =
		sscl::ComponentThread::getSelf()->getIoContext();

	std::array<uint8_t, UDP_BCAST_MSG_BUFFER_NBYTES> receiveBuffer{};
	boost::asio::ip::udp::endpoint senderEndpoint;

	while (!canceler.isCancellationRequested())
	{
		const adapters::boostAsio::UdpReceiveFromAReq::Result receiveResult =
			co_await adapters::boostAsio::getUdpReceiveFromAReqAwaiter(
				ioContext, socket,
				boost::asio::buffer(receiveBuffer), senderEndpoint);

		if (isShutdownReceiveError(receiveResult.ec)) { break; }

		if (receiveResult.ec)
		{
			std::cerr << __func__
				<< ": Error receiving broadcast message: "
				<< receiveResult.ec.message() << std::endl;
			continue;
		}

		if (receiveResult.nbytes == 0) { continue; }

		if (!canceler.execUncancelableSegmentOrAbort([&]()
		{
			eventHandlerNursery.launch(
				[this,
				payload = BroadcastIndPayload(
					receiveBuffer.data(),
					receiveResult.nbytes,
					senderEndpoint)](
					sscl::co::NonViralTaskNursery::Slot::Lease &lease) mutable
				{
					return broadcastMsgIndCInd(
						lease.getExceptionStorage(),
						lease.getCallerLambda(),
						lease.getSyncCanceler(),
						std::move(payload));
				},
				logEventHandlerException);
		})) {
			break;
		}
	}

	co_return;
}

void BroadcastListener::start(void)
{
	if (daemonNursery.admissionIsOpen()) { return; }

	try
	{
		/**		EXPLANATION:
		 * Set up a boost::asio udp listening socket on the broadcast listening
		 * port.
		 *
		 *		FIXME:
		 * We should also set up a timer to check for devices that have gone
		 * away.
		 */
		socket.open(boost::asio::ip::udp::v4());
		socket.bind(listeningEndpoint);

		daemonNursery.openAdmission();
		// Launch the receive daemon coroutine on componentThread
		daemonNursery.launch(
			[this](sscl::co::NonViralTaskNursery::Slot::Lease &lease)
			{
				return broadcastReceiveCDaemon(
					sscl::co::ExplicitPostTarget{
						componentThread->getIoContext()},
					lease.getExceptionStorage(),
					lease.getCallerLambda(),
					lease.getSyncCanceler());
			},
			logReceiveDaemonException);

		std::cout << __func__ << ": BroadcastListener started on port "
			<< listeningPort << std::endl;
	}
	catch (const boost::system::system_error& e)
	{
		std::cerr << __func__ << ": Failed to start BroadcastListener: "
			<< e.what() << std::endl;
		throw;
	}
}

void BroadcastListener::stop(void)
{
	if (!daemonNursery.admissionIsOpen()) { return; }

	daemonNursery.requestCancelOnAll();

	try
	{
		socket.cancel();
		socket.close();
	}
	catch (const boost::system::system_error& e)
	{
		std::cerr << __func__ << ": Error stopping BroadcastListener: "
			<< e.what() << std::endl;
		throw;
	}

	daemonNursery.closeAdmission();
	daemonNursery.syncAwaitAllSettlements(
		sscl::ComponentThread::getSelf()->getIoContext());

	std::cout << __func__ << ": BroadcastListener stopped" << std::endl;
}

} // namespace comms
} // namespace livoxProto1
