/*
 * Copyright 2017 Google
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Targeted compilation is ONLY for testing. UIKit is weak-linked in actual
// release build.

#import <Foundation/Foundation.h>

#import "FirebaseCore/Extension/FirebaseCoreInternal.h"
#import "FirebaseDatabase/Sources/Api/Private/FIRDatabase_Private.h"
#import "FirebaseDatabase/Sources/Constants/FConstants.h"
#import "FirebaseDatabase/Sources/Public/FirebaseDatabase/FIRDatabaseReference.h"
#import "FirebaseDatabase/Sources/Realtime/FWebSocketConnection.h"
#import "FirebaseDatabase/Sources/Utilities/FStringUtilities.h"

#if TARGET_OS_IOS || TARGET_OS_TV ||                                           \
    (defined(TARGET_OS_VISION) && TARGET_OS_VISION)
#import <UIKit/UIKit.h>
#endif // TARGET_OS_IOS || TARGET_OS_TV || (defined(TARGET_OS_VISION) &&
       // TARGET_OS_VISION)

#if TARGET_OS_WATCH
#import <Network/Network.h>
#import <WatchKit/WatchKit.h>
#endif // TARGET_OS_WATCH

static NSString *const kAppCheckTokenHeader = @"X-Firebase-AppCheck";
static NSString *const kUserAgentHeader = @"User-Agent";
static NSString *const kGoogleAppIDHeader = @"X-Firebase-GMPID";

@interface FWebSocketConnection () {
    NSMutableString *frame;
    BOOL everConnected;
    BOOL isClosed;
    dispatch_source_t keepAlive;
    NSDate *keepAliveFireDate;
}

- (void)shutdown;
- (void)onClosed;
- (void)closeIfNeverConnected;
- (void)cancelKeepAlive;

#if TARGET_OS_WATCH
@property(nonatomic, strong) NSURLSessionWebSocketTask *webSocketTask;
#else
@property(nonatomic, strong) FSRWebSocket *webSocket;
#endif // TARGET_OS_WATCH
@property(nonatomic, strong) NSNumber *connectionId;
@property(nonatomic, readwrite) int totalFrames;
@property(nonatomic, readonly) BOOL buffering;
@property(nonatomic, readonly) NSString *userAgent;
@property(nonatomic) dispatch_queue_t dispatchQueue;

- (void)nop;

@end

@implementation FWebSocketConnection

@synthesize delegate;

- (void)dealloc {
    if (keepAlive) {
        dispatch_source_cancel(keepAlive);
    }
}
#if !TARGET_OS_WATCH
@synthesize webSocket;
#endif // !TARGET_OS_WATCH
@synthesize connectionId;

- (instancetype)initWith:(FRepoInfo *)repoInfo
                andQueue:(dispatch_queue_t)queue
             googleAppID:(NSString *)googleAppID
           lastSessionID:(NSString *)lastSessionID
           appCheckToken:(nullable NSString *)appCheckToken {
    self = [super init];
    if (self) {
        everConnected = NO;
        isClosed = NO;
        self.connectionId = [FUtilities LUIDGenerator];
        self.totalFrames = 0;
        self.dispatchQueue = queue;
        frame = nil;

        NSString *userAgent = [self userAgent];
        NSString *connectionURL =
            [repoInfo connectionURLWithLastSessionID:lastSessionID];

        FFLog(@"I-RDB083001", @"(wsc:%@) Connecting to: %@ as %@",
              self.connectionId, connectionURL, userAgent);

        NSURLRequest *req = [[self class] createRequestWithURL:connectionURL
                                                     userAgent:userAgent
                                                   googleAppID:googleAppID
                                                 appCheckToken:appCheckToken];
        if (req.URL == nil) {
            // NSURL rejects the connection URL. Don't create a web socket;
            // open reports a failed connection attempt instead. Don't log the
            // URL, since it can include the session ID.
            FFWarn(@"I-RDB083022",
                   @"(wsc:%@) Couldn't parse the connection URL for host %@",
                   self.connectionId, repoInfo.internalHost);
            return self;
        }
#if TARGET_OS_WATCH
        // Regular NSURLSession websocket.
        NSOperationQueue *opQueue = [[NSOperationQueue alloc] init];
        opQueue.maxConcurrentOperationCount = 1;
        opQueue.underlyingQueue = queue;
        NSURLSession *session = [NSURLSession
            sessionWithConfiguration:[NSURLSessionConfiguration
                                         defaultSessionConfiguration]
                            delegate:self
                       delegateQueue:opQueue];
        NSURLSessionWebSocketTask *task =
            [session webSocketTaskWithRequest:req];
        self.webSocketTask = task;

        [[NSNotificationCenter defaultCenter]
            addObserverForName:WKApplicationWillResignActiveNotification
                        object:nil
                         queue:opQueue
                    usingBlock:^(NSNotification *_Nonnull note) {
                      FFLog(@"I-RDB083015",
                            @"Received watchOS background notification, "
                            @"closing web socket.");
                      [self onClosed];
                    }];
#else
        // TODO(mmaksym): Remove googleAppID and userAgent from FSRWebSocket as
        // they are passed via NSURLRequest.
        self.webSocket = [[FSRWebSocket alloc] initWithURLRequest:req
                                                            queue:queue
                                                      googleAppID:googleAppID
                                                     andUserAgent:userAgent];
        [self.webSocket setDelegateDispatchQueue:queue];
        self.webSocket.delegate = self;
#endif // TARGET_OS_WATCH
    }
    return self;
}

- (NSString *)userAgent {
    NSString *systemVersion;
    NSString *deviceName;
    BOOL hasUiDeviceClass = NO;

// Targeted compilation is ONLY for testing. UIKit is weak-linked in actual
// release build.
#if TARGET_OS_IOS || TARGET_OS_TV || TARGET_OS_VISION
    Class uiDeviceClass = NSClassFromString(@"UIDevice");
    if (uiDeviceClass) {
        systemVersion = [uiDeviceClass currentDevice].systemVersion;
        deviceName = [uiDeviceClass currentDevice].model;
        hasUiDeviceClass = YES;
    }
#endif // TARGET_OS_IOS || TARGET_OS_TV || (defined(TARGET_OS_VISION) &&
       // TARGET_OS_VISION)

    if (!hasUiDeviceClass) {
        NSDictionary *systemVersionDictionary = [NSDictionary
            dictionaryWithContentsOfFile:
                @"/System/Library/CoreServices/SystemVersion.plist"];
        systemVersion =
            [systemVersionDictionary objectForKey:@"ProductVersion"];
        deviceName = [systemVersionDictionary objectForKey:@"ProductName"];
    }

    NSString *bundleIdentifier = [[NSBundle mainBundle] bundleIdentifier];

    // Sanitize '/'s in deviceName and bundleIdentifier for stats
    deviceName = [FStringUtilities sanitizedForUserAgent:deviceName];
    bundleIdentifier =
        [FStringUtilities sanitizedForUserAgent:bundleIdentifier];

    // Firebase/5/<semver>_<build date>_<git hash>/<os version>/{device model /
    // os (Mac OS X, iPhone, etc.}_<bundle id>
    NSString *ua = [NSString
        stringWithFormat:@"Firebase/%@/%@/%@/%@_%@", kWebsocketProtocolVersion,
                         [FIRDatabase buildVersion], systemVersion, deviceName,
                         bundleIdentifier];
    return ua;
}

- (BOOL)buffering {
    return frame != nil;
}

#pragma mark -
#pragma mark Public FWebSocketConnection methods

- (void)open {
    FFLog(@"I-RDB083002", @"(wsc:%@) FWebSocketConnection open.",
          self.connectionId);
    assert(delegate);
    everConnected = NO;
#if TARGET_OS_WATCH
    BOOL hasWebSocket = self.webSocketTask != nil;
#else
    BOOL hasWebSocket = self.webSocket != nil;
#endif // TARGET_OS_WATCH
    if (!hasWebSocket) {
        // There's no web socket because the connection URL is invalid. Report
        // a failed connection attempt instead, asynchronously like the web
        // socket would, so that the delegate falls back to the configured host
        // and retries with backoff. There's nothing for the connect timeout to
        // close.
        dispatch_async(self.dispatchQueue, ^{
          [self onClosed];
        });
        return;
    }
#if TARGET_OS_WATCH
    [self.webSocketTask resume];
    // We need to request data from the web socket in order for it to start
    // sending data.
    [self receiveWebSocketData];
#else
    [self.webSocket open];
#endif // TARGET_OS_WATCH
    dispatch_time_t when = dispatch_time(
        DISPATCH_TIME_NOW, kWebsocketConnectTimeout * NSEC_PER_SEC);
    dispatch_after(when, self.dispatchQueue, ^{
      [self closeIfNeverConnected];
    });
}

- (void)close {
    FFLog(@"I-RDB083003", @"(wsc:%@) FWebSocketConnection is being closed.",
          self.connectionId);
    isClosed = YES;
    [self cancelKeepAlive];
#if TARGET_OS_WATCH
    [self.webSocketTask
        cancelWithCloseCode:NSURLSessionWebSocketCloseCodeNormalClosure
                     reason:nil];
#else
    [self.webSocket close];
#endif // TARGET_OS_WATCH
}

- (void)start {
    // Start is a no-op for websockets.
}

- (void)send:(NSDictionary *)dictionary {

    [self resetKeepAlive];

    NSData *jsonData = [NSJSONSerialization dataWithJSONObject:dictionary
                                                       options:kNilOptions
                                                         error:nil];

    NSString *data = [[NSString alloc] initWithData:jsonData
                                           encoding:NSUTF8StringEncoding];

    NSArray *dataSegs = [FUtilities splitString:data
                                    intoMaxSize:kWebsocketMaxFrameSize];

    // First send the header so the server knows how many segments are
    // forthcoming
    if (dataSegs.count > 1) {
        NSString *formattedData =
            [NSString stringWithFormat:@"%u", (unsigned int)dataSegs.count];
        [self sendStringToWebSocket:formattedData];
    }

    // Then, actually send the segments.
    for (NSString *segment in dataSegs) {
        [self sendStringToWebSocket:segment];
    }
}

- (void)nop {
    if (!isClosed) {
        FFLog(@"I-RDB083004", @"(wsc:%@) nop", self.connectionId);
        // Note: the backend is expecting a string "0" here, not any special
        // ping/pong from build in websocket APIs.
        [self sendStringToWebSocket:@"0"];
    }
}

- (void)handleNewFrameCount:(int)numFrames {
    self.totalFrames = numFrames;
    frame = [[NSMutableString alloc] initWithString:@""];
    FFLog(@"I-RDB083006", @"(wsc:%@) handleNewFrameCount: %d",
          self.connectionId, self.totalFrames);
}

- (NSString *)extractFrameCount:(NSString *)message {
    if ([message length] <= 4) {
        int frameCount = [message intValue];
        if (frameCount > 0) {
            [self handleNewFrameCount:frameCount];
            return nil;
        }
    }
    [self handleNewFrameCount:1];
    return message;
}

- (void)appendFrame:(NSString *)message {
    [frame appendString:message];
    self.totalFrames = self.totalFrames - 1;

    if (self.totalFrames == 0) {
        // Call delegate and pass an immutable version of the frame
        NSDictionary *json = [NSJSONSerialization
            JSONObjectWithData:[frame dataUsingEncoding:NSUTF8StringEncoding]
                       options:kNilOptions
                         error:nil];
        frame = nil;
        FFLog(@"I-RDB083007",
              @"(wsc:%@) handleIncomingFrame sending complete frame: %d",
              self.connectionId, self.totalFrames);

        @autoreleasepool {
            [self.delegate onMessage:self withMessage:json];
        }
    }
}

- (void)handleIncomingFrame:(NSString *)message {
    // The realtime protocol only uses text frames. SocketRocket delivers an
    // NSData for a binary frame (and the watchOS task delivers a nil string),
    // neither of which responds to the NSString selectors used below, so a
    // server sending such a frame would crash the client. Ignore anything that
    // is not a text frame.
    if (![message isKindOfClass:[NSString class]]) {
        FFWarn(@"I-RDB083021",
               @"(wsc:%@) Ignoring non-text frame received on stream.",
               self.connectionId);
        return;
    }
    [self resetKeepAlive];
    if (self.buffering) {
        [self appendFrame:message];
    } else {
        NSString *remaining = [self extractFrameCount:message];
        if (remaining) {
            [self appendFrame:remaining];
        }
    }
}

#pragma mark -
#pragma mark URLSessionWebSocketDelegate watchOS implementation
#if TARGET_OS_WATCH

- (void)URLSession:(NSURLSession *)session
          webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didOpenWithProtocol:(NSString *)protocol {
    [self webSocketDidOpen];
}

- (void)URLSession:(NSURLSession *)session
       webSocketTask:(NSURLSessionWebSocketTask *)webSocketTask
    didCloseWithCode:(NSURLSessionWebSocketCloseCode)closeCode
              reason:(NSData *)reason {
    FFLog(@"I-RDB083011", @"(wsc:%@) didCloseWithCode: %ld %@",
          self.connectionId, (long)closeCode, reason);
    [self onClosed];
}

- (void)receiveWebSocketData {
    __weak __auto_type weakSelf = self;
    [self.webSocketTask receiveMessageWithCompletionHandler:^(
                            NSURLSessionWebSocketMessage *_Nullable message,
                            NSError *_Nullable error) {
      __auto_type strongSelf = weakSelf;
      if (strongSelf == nil || strongSelf->isClosed) {
          return;
      }

      if (message) {
          [strongSelf handleIncomingFrame:message.string];
      } else if (error) {
          FFWarn(@"I-RDB083020",
                 @"Error received from web socket, closing the connection. %@",
                 error);
          [strongSelf shutdown];
          return;
      }

      [strongSelf receiveWebSocketData];
    }];
}

#else

#pragma mark SRWebSocketDelegate implementation

- (void)webSocket:(FSRWebSocket *)webSocket didReceiveMessage:(id)message {
    [self handleIncomingFrame:message];
}

- (void)webSocket:(FSRWebSocket *)webSocket didFailWithError:(NSError *)error {
    FFLog(@"I-RDB083010", @"(wsc:%@) didFailWithError didFailWithError: %@",
          self.connectionId, [error description]);
    [self onClosed];
}

- (void)webSocket:(FSRWebSocket *)webSocket
    didCloseWithCode:(NSInteger)code
              reason:(NSString *)reason
            wasClean:(BOOL)wasClean {
    FFLog(@"I-RDB083011", @"(wsc:%@) didCloseWithCode: %ld %@",
          self.connectionId, (long)code, reason);
    [self onClosed];
}

#endif // TARGET_OS_WATCH

// Common to both SRWebSocketDelegate and URLSessionWebSocketDelegate.

- (void)webSocketDidOpen {
    FFLog(@"I-RDB083008", @"(wsc:%@) webSocketDidOpen", self.connectionId);

    everConnected = YES;

    if (isClosed || keepAlive) {
        return;
    }
    keepAlive = dispatch_source_create(DISPATCH_SOURCE_TYPE_TIMER, 0, 0,
                                       self.dispatchQueue);
    keepAliveFireDate =
        [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
    dispatch_source_set_timer(
        keepAlive,
        dispatch_time(DISPATCH_TIME_NOW,
                      kWebsocketKeepaliveInterval * NSEC_PER_SEC),
        kWebsocketKeepaliveInterval * NSEC_PER_SEC, NSEC_PER_SEC);
    __weak FWebSocketConnection *weakSelf = self;
    dispatch_source_set_event_handler(keepAlive, ^{
      FWebSocketConnection *strongSelf = weakSelf;
      if (!strongSelf || strongSelf->isClosed) {
          return;
      }
      strongSelf->keepAliveFireDate =
          [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
      [strongSelf nop];
    });
    dispatch_resume(keepAlive);
    FFLog(@"I-RDB083009", @"(wsc:%@) nop timer kicked off", self.connectionId);
}

#pragma mark -
#pragma mark Private methods

/** Sends a string through the open web socket. */
- (void)sendStringToWebSocket:(NSString *)string {
#if TARGET_OS_WATCH
    // Use built-in URLSessionWebSocket functionality.
    [self.webSocketTask sendMessage:[[NSURLSessionWebSocketMessage alloc]
                                        initWithString:string]
                  completionHandler:^(NSError *_Nullable error) {
                    if (error) {
                        FFWarn(@"I-RDB083016",
                               @"Error sending web socket data: %@.", error);
                        return;
                    }
                  }];
#else
    // Use existing SocketRocket implementation.
    [self.webSocket send:string];
#endif // TARGET_OS_WATCH
}

/**
 * Note that the close / onClosed / shutdown cycle here is a little different
 * from the javascript client. In order to properly handle deallocation, no
 * close-related action is taken at a higher level until we have received
 * notification from the websocket itself that it is closed. Otherwise, we end
 * up deallocating this class and the FConnection class before the websocket has
 * a change to call some of its delegate methods. So, since close is the
 * external close handler, we just set a flag saying not to call our own
 * delegate method and close the websocket. That will trigger a callback into
 * this class that can then do things like clean up the keepalive timer.
 *
 * closeIfNeverConnected is the exception, since a websocket that never
 * connected may not report that it closed for a long time. That's safe because
 * the websocket's delegate is weak.
 */

- (void)closeIfNeverConnected {
    if (!everConnected && !isClosed) {
        FFLog(@"I-RDB083012", @"(wsc:%@) Websocket timed out on connect",
              self.connectionId);
#if TARGET_OS_WATCH
        [self.webSocketTask
            cancelWithCloseCode:NSURLSessionWebSocketCloseCodeNoStatusReceived
                         reason:nil];
#else
        [self.webSocket close];
        self.webSocket.delegate = nil;
#endif // TARGET_OS_WATCH

        // Don't wait for the websocket to report that it closed, and ignore
        // the report if it comes later. It may not come for hours, e.g. if the
        // websocket is stuck in the TLS handshake, and the client would stay
        // offline until then.
        // https://github.com/firebase/firebase-ios-sdk/issues/9682
        [self onClosed];
    }
}

- (void)shutdown {
    if (isClosed) {
        return;
    }
    isClosed = YES;
    [self cancelKeepAlive];

    // Call delegate methods
    [self.delegate onDisconnect:self wasEverConnected:everConnected];
}

- (void)onClosed {
    if (!isClosed) {
        FFLog(@"I-RDB083013", @"Websocket is closing itself");
        [self shutdown];
    }
#if TARGET_OS_WATCH
    self.webSocketTask = nil;
#else
    self.webSocket = nil;
#endif // TARGET_OS_WATCH
    [self cancelKeepAlive];
}

- (void)resetKeepAlive {
    // Connection callbacks and this timer share the worker queue.
    if (isClosed || !keepAlive) {
        return;
    }
    NSDate *newTime =
        [NSDate dateWithTimeIntervalSinceNow:kWebsocketKeepaliveInterval];
    // Calling dispatch_source_set_timer on every frame is unnecessary, so wait
    // at least 5 seconds before updating it.
    if ([newTime timeIntervalSinceDate:keepAliveFireDate] > 5) {
        FFLog(@"I-RDB083014", @"(wsc:%@) resetting keepalive, to %@ ; old: %@",
              self.connectionId, newTime, keepAliveFireDate);
        dispatch_source_set_timer(
            keepAlive,
            dispatch_time(DISPATCH_TIME_NOW,
                          kWebsocketKeepaliveInterval * NSEC_PER_SEC),
            kWebsocketKeepaliveInterval * NSEC_PER_SEC, NSEC_PER_SEC);
        keepAliveFireDate = newTime;
    }
}

- (void)cancelKeepAlive {
    if (keepAlive) {
        dispatch_source_cancel(keepAlive);
        keepAlive = nil;
    }
    keepAliveFireDate = nil;
}

+ (NSURLRequest *)createRequestWithURL:(NSString *)connectionURL
                             userAgent:(NSString *)userAgent
                           googleAppID:(NSString *)googleAppID
                         appCheckToken:(nullable NSString *)appCheckToken {

    NSMutableURLRequest *request = [[NSMutableURLRequest alloc]
        initWithURL:[[NSURL alloc] initWithString:connectionURL]];

    [request setValue:appCheckToken forHTTPHeaderField:kAppCheckTokenHeader];
    [request setValue:userAgent forHTTPHeaderField:kUserAgentHeader];
    [request setValue:googleAppID forHTTPHeaderField:kGoogleAppIDHeader];

    return [request copy];
}

@end