#include <cstdlib>
#include <cstdio>
#include <csignal>

#include <thread>
#include <atomic>

#include <unistd.h>
#include <errno.h>
#include <error.h>

#include <sys/socket.h>
#include <netinet/ip.h>
#include <arpa/inet.h>

/* UDP multicast example. Compile with threads and C++11 (g++ -pthread --std=c++11)
 * 
 * Sending to multicast group MCASTIP:MCASTPORT
 *   - create a UDP socket
 *   - send a datagram to MCASTIP:MCASTPORT
 * 
 * Receiving datagrams sent to a multicast group MCASTIP:MCASTPORT
 *   - create a UDP socket
 *   - bind to some address and MCASTPORT
 *   - prepare ip_mreqn structure – fill it with MCASTIP and optionally (but recommended) with local IP and interface id
 *   - call setsockopt to set option IP_ADD_MEMBERSHIP on level IPPROTO_IP to the aforementioned structure
 *   - receive datagrams
 * 
 * setsockopt of IP_ADD_MEMBERSHIP sends an IGMP message "join group" and tells OS that the program waits for messages for the group.
 * 
 * More info at: https://tldp.org/HOWTO/Multicast-HOWTO-6.html
 */


const sockaddr_in groupSockaddr {
	.sin_family = AF_INET,
	.sin_port = htons(6789),
	.sin_addr = {inet_addr("239.255.123.45")}
};

std::atomic<bool> quitting {false};

void ctrl_c(int){
	if(quitting)
		exit(-1);
	quitting = true;
}

void setupSignals();

int main(int argc, char **argv) {
	
	// create socket
	int sockfd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if(sockfd==-1) error(1,errno, "socket");
	
	// set SO_REUSEADDR (or only one app will be allowed to receive multicast messages from the groupIP:port pair)
	// this is one of very few cases when SO_REUSEADDR on UDP has any use
	const int one = 1;
	int res = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
	if(res) error(1,errno, "setsockopt SO_REUSEADDR");
	
	// bind 
	sockaddr_in myAddress {
		.sin_family = AF_INET,
		.sin_port = groupSockaddr.sin_port,
		.sin_addr = {INADDR_ANY}
	};
	res = bind(sockfd, (sockaddr*) &myAddress, sizeof(myAddress));
	if(res) error(1, errno, "bind");
	
	// send IGMP join group
	ip_mreqn groupDescription{
		.imr_multiaddr = groupSockaddr.sin_addr,
		.imr_address =  {INADDR_ANY},
		.imr_ifindex = 0
	};
	auto ret = setsockopt(sockfd, IPPROTO_IP, IP_ADD_MEMBERSHIP, &groupDescription, sizeof(groupDescription));
	if(ret) error(1, errno, "setsockopt IP_ADD_MEMBERSHIP");
	
	// handle ctrl+c for gracefull exit
	setupSignals();
	
	// receive and display ANY packet that arrives
	std::thread readerThread([&]{
		sockaddr_in peerAddr;
		socklen_t peerAddrSize;
		char buffer[255];
		while(true){
			peerAddrSize = sizeof(peerAddr);
			auto readSize = recvfrom(sockfd, buffer, 254, MSG_TRUNC, (sockaddr*) &peerAddr, &peerAddrSize);
			if(quitting) break;
			if(readSize==-1) error(1,errno, "recvfrom");
			if(readSize>254) { error(0,0,"UDP packet too long - truncated from %ld to %d", readSize, 254); readSize=254;}
			buffer[readSize] = 0;
			printf("<%s:%hu> %s", inet_ntoa(peerAddr.sin_addr), ntohs(peerAddr.sin_port), buffer);
		}
	});
	
	// read standard input and send it to the group
	char buffer[255];
	while(true){
		auto readSize = read(0, buffer, 254);
		if(quitting || !readSize) break;
		if(readSize<0) { error(0, errno, "reading stdin"); break; }
		
		auto res = sendto(sockfd, buffer, readSize, 0, (sockaddr*) &groupSockaddr, sizeof(groupSockaddr));
		if(res!=readSize) { error(0,  errno, "sendto"); break; }
	}
	
	quitting = true;
	// leave group
	setsockopt(sockfd, IPPROTO_IP, IP_DROP_MEMBERSHIP, &groupDescription, sizeof(groupDescription));
	// close socket AND interrupt all blocked recvfrom 
	shutdown(sockfd, SHUT_RDWR);
	close(sockfd);
	// wait for child
	readerThread.join();
	
	return 0;
}

void setupSignals(){
	struct sigaction action;
	
	// read sigint handler description
	auto res = sigaction(SIGINT, nullptr, &action);
	if(res) error(1, errno, "sigaction read");
	
	// set new interrupt handler
	action.sa_handler = ctrl_c;
	// disable SA_RESTART flag to exit blocking IO upon a signal
	action.sa_flags &= ~SA_RESTART;
	
	// update sigint handler description
	res = sigaction(SIGINT, &action, nullptr);
	if(res) error(1, errno, "sigaction modify");
}