#import "UDPListener.h"
#import "Categories.h"

@implementation UDPListener

- (nullable instancetype)init
{
	return [self initWithQueue:nil];
}

- (nullable instancetype)initWithQueue:(dispatch_queue_t __nullable)queue
{
	self = [super init];

	_recvPacketHandler = nil;
	_childConnections = [NSMutableDictionary dictionary];

	self.queue = (queue) ? queue : dispatch_queue_create("UDPListener", 
			dispatch_queue_attr_make_with_autorelease_frequency(
			DISPATCH_QUEUE_SERIAL, DISPATCH_AUTORELEASE_FREQUENCY_WORK_ITEM));

	return self;
}

- (void)dealloc
{
	[self stopListening];
}

- (void)startListening:(uint16_t)port packetHandler:(nonnull RecvPacketHandler)packetHandler
{
	nw_parameters_t parameters = nw_parameters_create_secure_udp(
			NW_PARAMETERS_DISABLE_PROTOCOL, NW_PARAMETERS_DEFAULT_CONFIGURATION);
	const char *portString = [[NSString stringWithFormat:@"%u", port] UTF8String];
	nw_endpoint_t localEndpoint = nw_endpoint_create_host("::", portString);
	nw_parameters_set_local_endpoint(parameters, localEndpoint);
	
	_recvPacketHandler = packetHandler;
	
	self.listener = nw_listener_create(parameters);

	nw_listener_set_state_changed_handler(self.listener, 
			^(nw_listener_state_t state, nw_error_t error)
	{
		[self listenerDidChangeState:state error:error];
	});
	nw_listener_set_new_connection_handler(self.listener, ^(nw_connection_t connection)
	{
		[self listenerDidAcceptConnection:connection];
	});
	nw_listener_set_queue(self.listener, self.queue);
	nw_listener_start(self.listener);
}

- (void)stopListening
{
	if (self.listener)
	{
		nw_listener_set_new_connection_handler(self.listener, NULL);
		nw_listener_set_state_changed_handler(self.listener, NULL);
		nw_listener_cancel(self.listener);
		self.listener = nil;
	}

	for(nw_connection_t connection in [_childConnections allValues])
		nw_connection_cancel(connection);
}

- (void)sendPacket:(nonnull NSData *)packetData toEndpoint:(nonnull NSString *)remoteEndpoint
		completionHandler:(nullable SendCompletionHandler)completionHandler
{
	nw_connection_t		connection = [self childConnectionForHostString:remoteEndpoint];
	dispatch_data_t		dispatchData = [packetData createDispatchData];
	
	nw_connection_send(connection, dispatchData, NW_CONNECTION_DEFAULT_MESSAGE_CONTEXT, 
			true, ^(nw_error_t _Nullable nwError)
	{
		if (completionHandler)
			completionHandler((nwError) ? (__bridge NSError *)nw_error_copy_cf_error(nwError) : nil);
	});
}

#pragma mark -

- (nullable nw_connection_t)childConnectionForHostString:(nonnull NSString *)hostString
{
	nw_connection_t	connection = nil;
	
	connection = _childConnections[hostString];
	if (nil == connection)
	{
		nw_endpoint_t	endpoint = [hostString nwEndpointFromStringRepresentation];
		nw_parameters_t parameters = nw_parameters_create_secure_udp(
								NW_PARAMETERS_DISABLE_PROTOCOL, 
								NW_PARAMETERS_DEFAULT_CONFIGURATION);
		
		connection = nw_connection_create(endpoint, parameters);
		[self listenerDidCreateConnection:connection];
	}
	
	return connection;
}

- (void)listenerDidChangeState:(nw_listener_state_t)state error:(nullable nw_error_t)error
{
	switch(state)
	{
		case nw_listener_state_invalid:
			break;
		case nw_listener_state_waiting:
			NSLog(@"Listener waiting");
			break;
		case nw_listener_state_ready:
			NSLog(@"Listener ready");
			break;
		case nw_listener_state_cancelled:
			// We don't normally get this, as we detach our callback before cancelling
			NSLog(@"Listener cancelled");
			break;
		case nw_listener_state_failed:
			{
			NSError *nsError = (error) ? (__bridge NSError *)nw_error_copy_cf_error(error) : nil;
			NSLog(@"Listener failed: %@", (nsError) ? nsError : @"Unknown");
			
			[self stopListening];
			}
			break;
		default:
			break;
	}
}

- (void)childConnection:(nonnull nw_connection_t)connection didChangeState:
		(nw_connection_state_t)state error:(nullable nw_error_t)error
{
	switch(state)
	{
		case nw_connection_state_invalid:
			break;
		case nw_connection_state_waiting:
			NSLog(@"Child connection waiting");
			break;
		case nw_connection_state_preparing:
			NSLog(@"Child connection preparing");
			break;
		case nw_connection_state_ready:
			NSLog(@"Child connection ready");
			break;
		case nw_connection_state_cancelled:
			NSLog(@"Child connection cancelled");
			for(NSString *key in [_childConnections allKeysForObject:connection])
				_childConnections[key] = nil;
			break;
		case nw_connection_state_failed:
			{
			NSError *nsError = (error) ? (__bridge NSError *)nw_error_copy_cf_error(error) : nil;
			NSLog(@"Child connection failed: %@", (nsError) ? nsError : @"Unknown");
			
			nw_connection_cancel(connection);
			}
			break;
		default:
			break;
	}
}

- (void)listenerDidCreateConnection:(nonnull nw_connection_t)connection
{
	nw_endpoint_t	endpoint = nw_connection_copy_endpoint(connection);
	NSString *		hostString = [NSString stringRepresentationForNWEndpoint:endpoint];

	nw_connection_set_state_changed_handler(connection, 
			^(nw_connection_state_t state, nw_error_t error)
	{
		[self childConnection:connection didChangeState:state error:error];
	});
	nw_connection_set_queue(connection, self.queue);
	nw_connection_start(connection);

	_childConnections[hostString] = connection;
	[self waitForNextRead:connection fromHost:hostString];
}

- (void)listenerDidAcceptConnection:(nonnull nw_connection_t)connection
{
	nw_endpoint_t	endpoint = nw_connection_copy_endpoint(connection);
	NSString *		hostString = [NSString stringRepresentationForNWEndpoint:endpoint];

	nw_connection_set_state_changed_handler(connection, 
			^(nw_connection_state_t state, nw_error_t error)
	{
		[self childConnection:connection didChangeState:state error:error];
	});
	nw_connection_set_queue(connection, self.queue);
	nw_connection_start(connection);

	_childConnections[hostString] = connection;
	[self waitForNextRead:connection fromHost:hostString];
}

- (void)waitForNextRead:(nonnull nw_connection_t)connection fromHost:(nonnull NSString *)hostString
{
	nw_connection_receive(connection, 0, 64 * 1024, ^(dispatch_data_t data,
				nw_content_context_t context, bool isComplete, nw_error_t error)
	{
		bool	wasCancelled = error && (ECANCELED == nw_error_get_error_code(error)) &&
						(nw_error_domain_posix == nw_error_get_error_domain(error));

		if (data)
		{
			self->_recvPacketHandler([NSData dataFromDispatchData:data], hostString);
			[self waitForNextRead:connection fromHost:hostString];
		}
		else if (isComplete)
			nw_connection_cancel(connection);
		else if (error && ! wasCancelled)
		{
			NSError *nsError = (error) ? (__bridge NSError *)nw_error_copy_cf_error(error) : nil;
			NSLog(@"Connection read failed: %@", (nsError) ? nsError : @"Unknown");
			
			nw_connection_cancel(connection);
		}
	});
}

@end // UDPListener