#import "IMCVideoStreamServer.h"
#import "IMCCanvasView.h"
#import <arpa/inet.h>
#import <netinet/in.h>
#import <sys/socket.h>
#import <string.h>
#import <unistd.h>

static const NSInteger IMCFrameSize = 15000;
static const NSInteger IMCUDPChunkSize = 1000;
static const NSInteger IMCUDPChunkCount = 15;
static const NSInteger IMCColorStreamPort = 8083;
static const NSUInteger IMCColorHeaderSize = 16;

@interface IMCVideoStreamServer ()
@property (nonatomic, weak) IMCCanvasView *canvasView;
@property (nonatomic, assign) BOOL running;
@property (nonatomic, assign) BOOL debugEnabled;
@property (nonatomic, strong) NSDate *startedAt;
@property (nonatomic, strong) NSDate *lastFrameAt;
@property (nonatomic, copy) NSString *lastProtocol;
@property (nonatomic, assign) uint64_t tcpConnections;
@property (nonatomic, assign) uint64_t tcpFrames;
@property (nonatomic, assign) uint64_t tcpBytes;
@property (nonatomic, assign) uint64_t udpFrames;
@property (nonatomic, assign) uint64_t udpPackets;
@property (nonatomic, assign) uint64_t udpBytes;
@property (nonatomic, assign) uint64_t invalidUDPPackets;
@property (nonatomic, assign) uint64_t discoveryRequests;
@property (nonatomic, assign) uint64_t colorConnections;
@property (nonatomic, assign) uint64_t colorFrames;
@property (nonatomic, assign) uint64_t colorBytes;
@property (nonatomic, assign) NSUInteger lastColorWidth;
@property (nonatomic, assign) NSUInteger lastColorHeight;
@end

@implementation IMCVideoStreamServer

- (instancetype)initWithCanvasView:(IMCCanvasView *)canvasView {
    self = [super init];
    if (self) {
        _canvasView = canvasView;
        _startedAt = [NSDate date];
        _lastProtocol = @"none";
    }
    return self;
}

- (void)start {
    if (self.running) {
        return;
    }
    self.running = YES;
    [NSThread detachNewThreadSelector:@selector(tcpLoop) toTarget:self withObject:nil];
    [NSThread detachNewThreadSelector:@selector(udpLoop) toTarget:self withObject:nil];
    [NSThread detachNewThreadSelector:@selector(discoveryLoop) toTarget:self withObject:nil];
    [NSThread detachNewThreadSelector:@selector(colorTCPLoop) toTarget:self withObject:nil];
}

- (BOOL)readExactly:(NSUInteger)length fromSocket:(int)socketFD intoData:(NSMutableData *)data {
    [data setLength:length];
    NSUInteger received = 0;
    while (self.running && received < length) {
        ssize_t count = recv(socketFD,
                             (uint8_t *)data.mutableBytes + received,
                             length - received,
                             0);
        if (count <= 0) {
            return NO;
        }
        received += (NSUInteger)count;
    }
    return received == length;
}

- (void)colorTCPLoop {
    @autoreleasepool {
        int serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            return;
        }

        int yes = 1;
        setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
        struct sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons((uint16_t)IMCColorStreamPort);
        if (bind(serverSocket, (struct sockaddr *)&address, sizeof(address)) < 0 ||
            listen(serverSocket, 1) < 0) {
            close(serverSocket);
            return;
        }

        while (self.running) {
            @autoreleasepool {
                int clientSocket = accept(serverSocket, NULL, NULL);
                if (clientSocket < 0) {
                    continue;
                }
                @synchronized (self) {
                    self.colorConnections += 1;
                }

                NSMutableData *header = [NSMutableData data];
                while ([self readExactly:IMCColorHeaderSize
                              fromSocket:clientSocket
                                intoData:header]) {
                    const uint8_t *bytes = header.bytes;
                    if (memcmp(bytes, "IMCR", 4) != 0 || bytes[4] != 1 || bytes[5] != 1) {
                        break;
                    }
                    NSUInteger width = ((NSUInteger)bytes[6] << 8) | bytes[7];
                    NSUInteger height = ((NSUInteger)bytes[8] << 8) | bytes[9];
                    uint32_t payloadLength = ((uint32_t)bytes[10] << 24) |
                                             ((uint32_t)bytes[11] << 16) |
                                             ((uint32_t)bytes[12] << 8) |
                                             bytes[13];
                    if (width == 0 || height == 0 || width > 2048 || height > 2048 ||
                        payloadLength != width * height * 2 || payloadLength > 8 * 1024 * 1024) {
                        break;
                    }

                    NSMutableData *payload = [NSMutableData data];
                    if (![self readExactly:payloadLength
                                fromSocket:clientSocket
                                  intoData:payload]) {
                        break;
                    }
                    NSData *snapshot = [payload copy];
                    @synchronized (self) {
                        self.colorFrames += 1;
                        self.colorBytes += payloadLength;
                        self.lastColorWidth = width;
                        self.lastColorHeight = height;
                        self.lastProtocol = @"color_tcp";
                        self.lastFrameAt = [NSDate date];
                    }
                    dispatch_async(dispatch_get_main_queue(), ^{
                        [self.canvasView displayRGB565FrameBuffer:snapshot width:width height:height];
                    });
                }
                close(clientSocket);
            }
        }
        close(serverSocket);
    }
}

- (void)tcpLoop {
    @autoreleasepool {
        int serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            return;
        }

        int yes = 1;
        setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

        struct sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons(8081);

        if (bind(serverSocket, (struct sockaddr *)&address, sizeof(address)) < 0) {
            close(serverSocket);
            return;
        }

        if (listen(serverSocket, 1) < 0) {
            close(serverSocket);
            return;
        }

        while (self.running) {
            @autoreleasepool {
                int clientSocket = accept(serverSocket, NULL, NULL);
                if (clientSocket < 0) {
                    continue;
                }
                [self noteTCPConnection];

                NSMutableData *frame = [NSMutableData dataWithLength:IMCFrameSize];
                NSUInteger received = 0;
                while (self.running) {
                    ssize_t count = recv(clientSocket,
                                         (uint8_t *)frame.mutableBytes + received,
                                         IMCFrameSize - received,
                                         0);
                    if (count <= 0) {
                        break;
                    }

                    received += (NSUInteger)count;
                    [self noteTCPBytes:(NSUInteger)count];
                    if (received == IMCFrameSize) {
                        [self displayFrame:frame protocol:@"tcp"];
                        received = 0;
                    }
                }

                close(clientSocket);
            }
        }

        close(serverSocket);
    }
}

- (void)udpLoop {
    @autoreleasepool {
        int udpSocket = socket(AF_INET, SOCK_DGRAM, 0);
        if (udpSocket < 0) {
            return;
        }

        int yes = 1;
        setsockopt(udpSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

        struct sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons(8082);

        if (bind(udpSocket, (struct sockaddr *)&address, sizeof(address)) < 0) {
            close(udpSocket);
            return;
        }

        NSMutableData *frame = [NSMutableData dataWithLength:IMCFrameSize];
        uint8_t currentFrameID = 0;
        uint16_t chunksReceived = 0;
        BOOL hasFrame = NO;
        uint8_t packet[1002];

        while (self.running) {
            ssize_t count = recvfrom(udpSocket, packet, sizeof(packet), 0, NULL, NULL);
            if (count < 2) {
                [self noteInvalidUDPPacket];
                continue;
            }

            uint8_t frameID = packet[0];
            uint8_t chunkIndex = packet[1];
            if (chunkIndex >= IMCUDPChunkCount || count < 2 + IMCUDPChunkSize) {
                [self noteInvalidUDPPacket];
                continue;
            }
            [self noteUDPPacketBytes:(NSUInteger)count];

            if (!hasFrame || frameID != currentFrameID) {
                currentFrameID = frameID;
                chunksReceived = 0;
                hasFrame = YES;
            }

            memcpy((uint8_t *)frame.mutableBytes + chunkIndex * IMCUDPChunkSize,
                   packet + 2,
                   IMCUDPChunkSize);
            chunksReceived |= (uint16_t)(1 << chunkIndex);

            if (chunksReceived == 0x7fff) {
                [self displayFrame:frame protocol:@"udp"];
                chunksReceived = 0;
            }
        }

        close(udpSocket);
    }
}

- (void)discoveryLoop {
    @autoreleasepool {
        int udpSocket = socket(AF_INET, SOCK_DGRAM, 0);
        if (udpSocket < 0) {
            return;
        }

        int yes = 1;
        setsockopt(udpSocket, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
        setsockopt(udpSocket, SOL_SOCKET, SO_BROADCAST, &yes, sizeof(yes));

        struct sockaddr_in address;
        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = INADDR_ANY;
        address.sin_port = htons(5000);

        if (bind(udpSocket, (struct sockaddr *)&address, sizeof(address)) < 0) {
            close(udpSocket);
            return;
        }

        char buffer[64];
        while (self.running) {
            struct sockaddr_in clientAddress;
            socklen_t clientLength = sizeof(clientAddress);
            ssize_t count = recvfrom(udpSocket,
                                     buffer,
                                     sizeof(buffer),
                                     0,
                                     (struct sockaddr *)&clientAddress,
                                     &clientLength);
            if (count == 15 && memcmp(buffer, "DISCOVER_SCREEN", 15) == 0) {
                [self noteDiscoveryRequest];
                const char *reply = "SCREEN_IP_8080";
                sendto(udpSocket,
                       reply,
                       strlen(reply),
                       0,
                       (struct sockaddr *)&clientAddress,
                       clientLength);
            }
        }

        close(udpSocket);
    }
}

- (void)displayFrame:(NSData *)frame protocol:(NSString *)protocol {
    [self noteFrameForProtocol:protocol];
    NSData *snapshot = [frame copy];
    dispatch_async(dispatch_get_main_queue(), ^{
        [self.canvasView displayMonochromeFrameBuffer:snapshot];
    });
}

- (NSDictionary *)streamStats {
    @synchronized (self) {
        NSTimeInterval uptime = [[NSDate date] timeIntervalSinceDate:self.startedAt ?: [NSDate date]];
        NSTimeInterval secondsSinceLastFrame = self.lastFrameAt ? [[NSDate date] timeIntervalSinceDate:self.lastFrameAt] : -1.0;
        return @{
            @"debug_enabled": @(self.debugEnabled),
            @"running": @(self.running),
            @"tcp_port": @8081,
            @"udp_port": @8082,
            @"color_tcp_port": @(IMCColorStreamPort),
            @"discovery_port": @5000,
            @"frame_size_bytes": @(IMCFrameSize),
            @"udp_chunk_size_bytes": @(IMCUDPChunkSize),
            @"udp_chunks_per_frame": @(IMCUDPChunkCount),
            @"uptime_sec": @((NSInteger)uptime),
            @"tcp_connections": @(self.tcpConnections),
            @"tcp_frames": @(self.tcpFrames),
            @"tcp_bytes": @(self.tcpBytes),
            @"udp_frames": @(self.udpFrames),
            @"udp_packets": @(self.udpPackets),
            @"udp_bytes": @(self.udpBytes),
            @"invalid_udp_packets": @(self.invalidUDPPackets),
            @"discovery_requests": @(self.discoveryRequests),
            @"color_connections": @(self.colorConnections),
            @"color_frames": @(self.colorFrames),
            @"color_bytes": @(self.colorBytes),
            @"last_color_width": @(self.lastColorWidth),
            @"last_color_height": @(self.lastColorHeight),
            @"last_protocol": self.lastProtocol ?: @"none",
            @"seconds_since_last_frame": @(secondsSinceLastFrame)
        };
    }
}

- (void)resetStreamStats {
    @synchronized (self) {
        self.startedAt = [NSDate date];
        self.lastFrameAt = nil;
        self.lastProtocol = @"none";
        self.tcpConnections = 0;
        self.tcpFrames = 0;
        self.tcpBytes = 0;
        self.udpFrames = 0;
        self.udpPackets = 0;
        self.udpBytes = 0;
        self.invalidUDPPackets = 0;
        self.discoveryRequests = 0;
        self.colorConnections = 0;
        self.colorFrames = 0;
        self.colorBytes = 0;
        self.lastColorWidth = 0;
        self.lastColorHeight = 0;
    }
}

- (void)setDebugEnabled:(BOOL)enabled {
    @synchronized (self) {
        _debugEnabled = enabled;
    }
}

- (void)noteTCPConnection {
    @synchronized (self) {
        self.tcpConnections += 1;
        if (self.debugEnabled) {
            NSLog(@"IPhoneMCP stream: TCP client connected");
        }
    }
}

- (void)noteTCPBytes:(NSUInteger)count {
    @synchronized (self) {
        self.tcpBytes += count;
    }
}

- (void)noteUDPPacketBytes:(NSUInteger)count {
    @synchronized (self) {
        self.udpPackets += 1;
        self.udpBytes += count;
    }
}

- (void)noteInvalidUDPPacket {
    @synchronized (self) {
        self.invalidUDPPackets += 1;
        if (self.debugEnabled) {
            NSLog(@"IPhoneMCP stream: invalid UDP packet");
        }
    }
}

- (void)noteDiscoveryRequest {
    @synchronized (self) {
        self.discoveryRequests += 1;
        if (self.debugEnabled) {
            NSLog(@"IPhoneMCP stream: discovery request");
        }
    }
}

- (void)noteFrameForProtocol:(NSString *)protocol {
    @synchronized (self) {
        if ([protocol isEqualToString:@"tcp"]) {
            self.tcpFrames += 1;
        } else if ([protocol isEqualToString:@"udp"]) {
            self.udpFrames += 1;
        }
        self.lastProtocol = protocol ?: @"none";
        self.lastFrameAt = [NSDate date];
        if (self.debugEnabled) {
            NSLog(@"IPhoneMCP stream: %@ frame displayed", self.lastProtocol);
        }
    }
}

@end