//
//  MIUdpDirect.m
//  NeTunnel
//
//  Created by Paul Ling on 9/25/18.
//  Copyright © 2018 MobileIron. All rights reserved.
//

#import "DnsManager.h"
#import "MIAppProxyProvider.h"
#import "MIConnection.h"
#import "MIUdpHandler.h"
#import "NELog.h"

#pragma mark * Interface
@interface MIUdpHandler ()

@property (nonatomic, weak) MIConnectionUDP *connection;
+ (NSString *)endpointName:(NWHostEndpoint *)endpoint;
- (NSUInteger)enqueuePacket:(NSData *)packet toEndpoint:(NWHostEndpoint *)endpoint;
- (NSMutableArray<NSData *> *)emptyQueueForEndpoint:(NWHostEndpoint *)endpoint;

@end

#import "MIUdpDirect.h"

#pragma mark * Interface
@interface MIUdpDirect ()

@property (nonatomic, strong) GCDAsyncUdpSocket *udpSocket;
@property (nonatomic) NSString * connName;
@property (nonatomic) NSString * localEpName;
@property (atomic)  long tag;

@end


#import "GCDAsyncUdpSocket.h"

#pragma mark * Implementation
@implementation MIUdpDirect

//static NSMutableDictionary<NSString *, GCDAsyncUdpSocket *> * gs_sockets;

- (instancetype)initFromConnection:(MIConnectionUDP *)connection {
    self = [super initFromConnection:connection];
    //if (gs_sockets == nil) {
    //    gs_sockets = [[NSMutableDictionary alloc] init];
    //}
    
    _tag = 0;
    _connName = [connection name];
    return self;
}

- (BOOL)open:(NWHostEndpoint *)localEndpoint {
    //NSString *localEpName = [MIUdpHandler endpointName:localEndpoint];
    GCDAsyncUdpSocket *socket = nil; //[gs_sockets objectForKey:localEpName];
    if (socket == nil) {
        NWHostEndpoint *localEP = (NWHostEndpoint *)[[self.connection flow] localEndpoint];
        NELogS(@"Allocate new UDP socket for %@ from EP %@", _connName, localEP);
        NSError *error = nil;
        _udpSocket = [[GCDAsyncUdpSocket alloc] initWithDelegate:self delegateQueue:dispatch_get_main_queue()];
                    // delegateQueue:dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_BACKGROUND, 0)];
        if (![_udpSocket bindToPort:[localEndpoint.port intValue] error:&error]) {
            NELogError(error, @"%@ Error when bind UDP socket to port %@", _connName, localEndpoint.port);
            _udpSocket = nil;
            return NO;
        }
        if (![_udpSocket beginReceiving:&error]) {
            NELogError(error, @"%@ Error when begin receiving from UDP socket", self.connName);
            [_udpSocket close];
            _udpSocket = nil;
            return NO;
        }
        _localEndpoint = [NWHostEndpoint endpointWithHostname:localEndpoint.hostname port:@(_udpSocket.localPort).stringValue];
        _localEpName = [MIUdpHandler endpointName:_localEndpoint];
        //gs_sockets[_localEpName] = socket;

        NELogD(@"Added UDP socket from %@ for %@", _localEpName, _connName);
    }
    else {
        _localEndpoint = [NWHostEndpoint endpointWithHostname:localEndpoint.hostname port:@(socket.localPort).stringValue];
        _localEpName = [MIUdpHandler endpointName:_localEndpoint];
        NELogD(@"Opened existing UDP socket from %@ for %@", _localEpName, _connName);
    }

    return YES;
}

- (BOOL)isReady:(NWHostEndpoint *)endpoint {
    //GCDAsyncUdpSocket *socket = [gs_sockets objectForKey:_localEpName];
    return _udpSocket != nil && !_udpSocket.isClosed;
}

- (void)close {
    NELogD(@"Close UDP %@ with %@", _localEpName, _connName);
    //GCDAsyncUdpSocket *socket = [gs_sockets objectForKey:_localEpName];
    //[gs_sockets removeObjectForKey:_localEpName];
    [_udpSocket close];
}

- (void)sendPacket:(NSData *)packet toEndpoint:(NWHostEndpoint *)endpoint {
    //NELogD(@"- send 1 packet at a time");
    NELogS(@"- send packet length = %lu, endpoint %@:%@", [packet length], [endpoint hostname], [endpoint port]);
    if ([self isReady:endpoint]) {
        //GCDAsyncUdpSocket *socket = [gs_sockets objectForKey:_localEpName];
        if (!_udpSocket) {
            NELogE(@"Unable to get a UDP socket of %@", _localEpName);
            return;
        }
        if ([_udpSocket isClosed]) {
            NELogE(@"Cannot send packet to %@ due to socket closed", endpoint);
            return;
        }
        
        self.connection.idleCounter = 0;
        NSMutableArray<NSData *> *queuedPkts = [self emptyQueueForEndpoint:endpoint];
        if (packet) {
            [queuedPkts addObject:packet];
        }
        //Printing each datagram--lets not do this
        //for (NSData *datagram in queuedPkts) {
        //    NELogD(@"Non DNS UDP Count to write= %lu", [datagram length]);
        //    NELogS(@"--> Datagram: %@", [datagram description]);
        //}

        for (NSData *datagram in queuedPkts) {
            [_udpSocket sendData:datagram toHost:[endpoint hostname] port:[endpoint.port intValue] withTimeout:-1 tag:_tag];
            _tag++;
        }
 
    } else if (packet) {
        NELogD(@"Datagram queued for endpoint %@", [MIUdpHandler endpointName:endpoint]);
        [self enqueuePacket:packet toEndpoint:endpoint];
    }
}

+ (NWHostEndpoint*) fromGCDAddress: (NSData*) address {
    NSString *host = nil;
    uint16_t port = 0;
    [GCDAsyncUdpSocket getHost:&host port:&port fromAddress:address];
    return [NWHostEndpoint endpointWithHostname:host port:@(port).stringValue];
}

- (void)udpSocket:(GCDAsyncUdpSocket *)sock didSendDataWithTag:(long)tag
{
    self.connection.pktCountUDPDeliveredToEndPoint++;
    NELogD(@"Datagram written");
    NELogD(@"%@ Non DNS UDP sent %lu", _connName, tag);
}

- (void)udpSocket:(GCDAsyncUdpSocket *)sock didNotSendDataWithTag:(long)tag dueToError:(NSError *)error
{
    // You could add checks here
    NELogError(error, @"%@ Could not send datagram %lu", _connName, tag);
}

- (void)udpSocket:(GCDAsyncUdpSocket *)sock didReceiveData:(NSData *)data
                                            fromAddress:(NSData *)address
                                            withFilterContext:(nullable id)filterContext
{
    NELogS(@"%@ Read UDP Datagrams", _connName);
    NWHostEndpoint * remoteEP = [MIUdpDirect fromGCDAddress:address];
    NSString * remoteEpName = [MIUdpHandler endpointName:remoteEP];
    NELogS(@"%@ Read UDP Datagrams from remote %@", _connName, remoteEpName);
    self.connection.idleCounter = 0;
    self.connection.pktCountUDPReceivedFromEndPoint++;
    [self.connection respond:data fromEndpoint:remoteEP];
}

- (void)udpSocketDidClose:(GCDAsyncUdpSocket *)sock withError:(NSError  * _Nullable)error
{
    NELogS(@"%@ Closed UDP Socket", _connName);
}

@end
