//
//  MBIMPingPongWebSocket.m
//  kordophoned
//
//  Created by James Magahern on 6/13/25.
//  Copyright © 2025 James Magahern. All rights reserved.
//

#import "MBIMPingPongWebSocket.h"
#import "GCDAsyncSocket.h"

// WebSocket opcodes (from the base class)
#define TIMEOUT_NONE          -1
#define TIMEOUT_REQUEST_BODY  10

#define TAG_HTTP_REQUEST_BODY      100
#define TAG_HTTP_RESPONSE_HEADERS  200
#define TAG_HTTP_RESPONSE_BODY     201

#define TAG_PREFIX                 300
#define TAG_MSG_PLUS_SUFFIX        301
#define TAG_MSG_WITH_LENGTH        302
#define TAG_MSG_MASKING_KEY        303
#define TAG_PAYLOAD_PREFIX         304
#define TAG_PAYLOAD_LENGTH         305
#define TAG_PAYLOAD_LENGTH16       306
#define TAG_PAYLOAD_LENGTH64       307

#define WS_OP_CONTINUATION_FRAME   0
#define WS_OP_TEXT_FRAME           1
#define WS_OP_BINARY_FRAME         2
#define WS_OP_CONNECTION_CLOSE     8
#define WS_OP_PING                 9
#define WS_OP_PONG                 10

// Additional tags for ping/pong handling
#define TAG_PING_PAYLOAD           400
#define TAG_PONG_SENT             401

@interface WebSocket ()
- (void)socket:(GCDAsyncSocket *)asyncSocket didReadData:(NSData *)data withTag:(long)tag;
@end

@implementation MBIMPingPongWebSocket
{
    NSData *pendingPingPayload;
}

#pragma mark - Ping/Pong Frame Construction

- (NSData *)createPongFrameWithPayload:(NSData *)pingPayload {
    NSMutableData *frame = [NSMutableData data];
    
    // First byte: FIN (1) + RSV (000) + Opcode (1010 = 0xA for Pong)
    uint8_t firstByte = 0x8A; // 10001010 in binary
    [frame appendBytes:&firstByte length:1];
    
    // Second byte: MASK (0) + Payload Length
    NSUInteger payloadLength = [pingPayload length];
    
    if (payloadLength <= 125) {
        // Short payload length (0-125 bytes)
        uint8_t secondByte = (uint8_t)payloadLength; // MASK bit is 0 for server-to-client
        [frame appendBytes:&secondByte length:1];
    }
    else if (payloadLength <= 65535) {
        // Medium payload length (126-65535 bytes)
        uint8_t secondByte = 126;
        [frame appendBytes:&secondByte length:1];
        
        // Add 16-bit length in network byte order (big-endian)
        uint16_t extendedLength = htons((uint16_t)payloadLength);
        [frame appendBytes:&extendedLength length:2];
    }
    else {
        // Large payload length (65536+ bytes) - rarely needed for pings
        uint8_t secondByte = 127;
        [frame appendBytes:&secondByte length:1];
        
        // Add 64-bit length in network byte order (big-endian)
        uint64_t extendedLength = OSSwapHostToBigInt64(payloadLength);
        [frame appendBytes:&extendedLength length:8];
    }
    
    // Append the payload (copied exactly from ping)
    if (payloadLength > 0) {
        [frame appendData:pingPayload];
    }
    
    return [frame copy];
}

- (void)sendPongWithPayload:(NSData *)payload socket:(GCDAsyncSocket *)asyncSocket {
    NSData *pongFrame = [self createPongFrameWithPayload:payload];
    
    // Send the pong frame directly through the socket
    [asyncSocket writeData:pongFrame withTimeout:TIMEOUT_NONE tag:TAG_PONG_SENT];
    
    NSLog(@"Sent pong frame with payload length: %lu", (unsigned long)[payload length]);
}

#pragma mark - Override AsyncSocket Delegate

- (void)socket:(GCDAsyncSocket *)asyncSocket didReadData:(NSData *)data withTag:(long)tag {
    NSUInteger nextOpCode = [[self valueForKey:@"nextOpCode"] unsignedIntValue]; // sheesh.
    
    // Handle our custom ping/pong tags
    if (tag == TAG_PING_PAYLOAD) {
        // We've read the ping payload, now send the pong response
        [self sendPongWithPayload:data socket:asyncSocket];
        
        // Continue reading the next frame
        [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
        return;
    }
    
    if (tag == TAG_PONG_SENT) {
        // Pong was sent successfully, continue normal operation
        return;
    }
    
    // For TAG_PAYLOAD_PREFIX, we need to intercept and handle ping frames
    if (tag == TAG_PAYLOAD_PREFIX) {
        UInt8 *pFrame = (UInt8 *)[data bytes];
        UInt8 frame = *pFrame;
        
        // Check if this is a valid WebSocket frame
        if (![self isValidWebSocketFrame:frame]) {
            [self didClose];
            return;
        }
        
        NSUInteger opcode = frame & 0x0F;
        
        // Handle ping frames specially
        if (opcode == WS_OP_PING) {
            // Read the payload length byte next
            [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH];
            return;
        }
        
        // Handle pong frames (just log and continue)
        if (opcode == WS_OP_PONG) {
            NSLog(@"Received pong frame");
            // Read the payload length byte next, but we'll ignore pong payloads
            [asyncSocket readDataToLength:1 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH];
            return;
        }
    }
    
    // For TAG_PAYLOAD_LENGTH, check if we're handling a ping
    if (tag == TAG_PAYLOAD_LENGTH) {
        UInt8 frame = *(UInt8 *)[data bytes];
        BOOL masked = (frame & 0x80) != 0;
        NSUInteger length = frame & 0x7F;
        
        // If we were processing a ping frame, handle the payload reading
        if (nextOpCode == WS_OP_PING) {
            if (length <= 125) {
                if (masked) {
                    // Read masking key first, then payload
                    [asyncSocket readDataToLength:4 withTimeout:TIMEOUT_NONE tag:TAG_MSG_MASKING_KEY];
                }
                // Read ping payload
                [asyncSocket readDataToLength:length withTimeout:TIMEOUT_NONE tag:TAG_PING_PAYLOAD];
                return;
            }
            else if (length == 126) {
                [asyncSocket readDataToLength:2 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH16];
                return;
            }
            else {
                [asyncSocket readDataToLength:8 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH64];
                return;
            }
        }
        
        // If we were processing a pong frame, just skip it
        if (nextOpCode == WS_OP_PONG) {
            if (length <= 125) {
                if (masked) {
                    [asyncSocket readDataToLength:4 withTimeout:TIMEOUT_NONE tag:TAG_MSG_MASKING_KEY];
                }
                // Skip pong payload and continue to next frame
                [asyncSocket readDataToLength:length withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_PREFIX];
                return;
            }
            else if (length == 126) {
                [asyncSocket readDataToLength:2 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH16];
                return;
            }
            else {
                [asyncSocket readDataToLength:8 withTimeout:TIMEOUT_NONE tag:TAG_PAYLOAD_LENGTH64];
                return;
            }
        }
    }
    
    // Handle masking key reading for ping frames
    if (tag == TAG_MSG_MASKING_KEY && nextOpCode == WS_OP_PING) {
        // Store the masking key and read the actual payload
        pendingPingPayload = [data copy];
        // The payload length was already determined, read it
        // Note: This is a simplified version - you'd need to track the actual length
        return;
    }
    
    // For all other cases, call the parent implementation
    [super socket:asyncSocket didReadData:data withTag:tag];
}

#pragma mark - Helper Methods

// Expose the isValidWebSocketFrame method since it's private in the parent
- (BOOL)isValidWebSocketFrame:(UInt8)frame {
    NSUInteger rsv = frame & 0x70;
    NSUInteger opcode = frame & 0x0F;
    if (rsv || (3 <= opcode && opcode <= 7) || (0xB <= opcode && opcode <= 0xF)) {
        return NO;
    }
    return YES;
}

// Override didReceiveMessage to add logging
- (void)didReceiveMessage:(NSString *)msg {
    NSLog(@"Received message: %@", msg);
    [super didReceiveMessage:msg];
}

// Override didOpen to add logging
- (void)didOpen {
    NSLog(@"WebSocket opened with ping/pong support");
    [super didOpen];
}

- (void)didClose {
    NSLog(@"WebSocket closed");
    [super didClose];
}

@end