// Copyright 2020 Google LLC
//
// 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.

#import "FirebasePerformance/Sources/AppActivity/FPRScreenTraceTracker.h"
#import "FirebasePerformance/Sources/AppActivity/FPRScreenTraceTracker+Private.h"

#import <Foundation/Foundation.h>
#import <UIKit/UIKit.h>

#import "FirebasePerformance/Sources/Common/FPRDiagnostics.h"

NSString *const kFPRPrefixForScreenTraceName = @"_st_";
NSString *const kFPRFrozenFrameCounterName = @"_fr_fzn";
NSString *const kFPRSlowFrameCounterName = @"_fr_slo";
NSString *const kFPRTotalFramesCounterName = @"_fr_tot";

// Note: This was previously 60 FPS, but that resulted in 90% +  of all frames collected to be
// flagged as slow frames, and so the threshold for iOS is 59 FPS.
// b/545164487 tracks any changes / next steps for ProMotion devices - where the frame rate can be
// as high as 120 FPS.
CFTimeInterval const kFPRSlowFrameThreshold = 1.0 / 59.0;  // Anything less than 59 FPS is slow.
CFTimeInterval const kFPRFrozenFrameThreshold = 700.0 / 1000.0;

#if TARGET_OS_TV
/** Default/fallback FPS value used when UIScreen.maximumFramesPerSecond is unavailable or invalid.
 */
static const NSInteger kFPRDefaultFPS = 60;

/** Epsilon value to avoid floating point comparison issues (e.g., 59.94 vs 60) on tvOS. */
static const CFTimeInterval kFPRSlowFrameEpsilon = 0.001;
#endif

/** Constant that indicates an invalid time. */
CFAbsoluteTime const kFPRInvalidTime = -1.0;

/** Returns the class name without the prefixed module name present in Swift classes
 * (e.g. MyModule.MyViewController -> MyViewController).
 */
static NSString *FPRUnprefixedClassName(Class theClass) {
  NSString *className = NSStringFromClass(theClass);
  NSRange periodRange = [className rangeOfString:@"." options:NSBackwardsSearch];
  if (periodRange.location == NSNotFound) {
    return className;
  }
  return periodRange.location < className.length - 1
             ? [className substringFromIndex:periodRange.location + 1]
             : className;
}

/** Returns the name for the screen trace for a given UIViewController. It does the following:
 *  - Removes module name from swift classes - (e.g. MyModule.MyViewController -> MyViewController)
 *  - Prepends "_st_" to the class name
 *  - Truncates the length if it exceeds the maximum trace length.
 *
 *  @param viewController The view controller whose screen trace name we want. Cannot be nil.
 *  @return An NSString containing the trace name, or a string containing an error if the
 *      class was nil.
 */
static NSString *FPRScreenTraceNameForViewController(UIViewController *viewController) {
  NSString *unprefixedClassName = FPRUnprefixedClassName([viewController class]);
  if (unprefixedClassName.length != 0) {
    NSString *traceName =
        [NSString stringWithFormat:@"%@%@", kFPRPrefixForScreenTraceName, unprefixedClassName];
    return traceName.length > kFPRMaxNameLength ? [traceName substringToIndex:kFPRMaxNameLength]
                                                : traceName;
  } else {
    // This is unlikely, but might happen if there's a regression on iOS where the class name
    // returned for a non-nil class is nil or empty.
    return @"_st_ERROR_NIL_CLASS_NAME";
  }
}

@implementation FPRScreenTraceHolder
@end

@implementation FPRScreenTraceTracker {
  /** Instance variable storing the total frames observed so far. */
  atomic_int_fast64_t _totalFramesCount;

  /** Instance variable storing the slow frames observed so far. */
  atomic_int_fast64_t _slowFramesCount;

  /** Instance variable storing the frozen frames observed so far. */
  atomic_int_fast64_t _frozenFramesCount;

#if TARGET_OS_TV
  /** Cached maximum frames per second from UIScreen on tvOS. */
  NSInteger _cachedMaxFPS;
#endif

  /** Cached slow frame budget. On iOS, uses kFPRSlowFrameThreshold (59 FPS).
   *  On tvOS, computed dynamically from UIScreen.maximumFramesPerSecond.
   */
  CFTimeInterval _cachedSlowBudget;
}

@dynamic totalFramesCount;
@dynamic frozenFramesCount;
@dynamic slowFramesCount;

+ (instancetype)sharedInstance {
  static FPRScreenTraceTracker *instance;
  static dispatch_once_t onceToken;
  dispatch_once(&onceToken, ^{
    instance = [[self alloc] init];
  });
  return instance;
}

- (instancetype)init {
  self = [super init];
  if (self) {
    _activeScreenTraces = [[NSMutableDictionary alloc] init];
    _activeScreenTracesLock = [[NSRecursiveLock alloc] init];

    _previouslyVisibleViewControllers = nil;  // Will be set when there is data.
    _screenTraceTrackerSerialQueue =
        dispatch_queue_create("com.google.FPRScreenTraceTracker", DISPATCH_QUEUE_SERIAL);
    _screenTraceTrackerDispatchGroup = dispatch_group_create();

    atomic_store_explicit(&_totalFramesCount, 0, memory_order_relaxed);
    atomic_store_explicit(&_frozenFramesCount, 0, memory_order_relaxed);
    atomic_store_explicit(&_slowFramesCount, 0, memory_order_relaxed);

#if TARGET_OS_TV
    // Initialize slow budget and maxFPS values with defaults for tvOS.
    // On tvOS, refresh rate depends on the connected display mode (e.g. 50 Hz vs 60 Hz).
    _cachedMaxFPS = kFPRDefaultFPS;
    _cachedSlowBudget = 1.0 / (CFTimeInterval)kFPRDefaultFPS;

    if ([NSThread isMainThread]) {
      [self updateCachedSlowBudget];
    } else {
      dispatch_sync(dispatch_get_main_queue(), ^{
        [self updateCachedSlowBudget];
      });
    }
#else
    _cachedSlowBudget = kFPRSlowFrameThreshold;
#endif

    _displayLink = [CADisplayLink displayLinkWithTarget:self selector:@selector(displayLinkStep)];
    [_displayLink addToRunLoop:[NSRunLoop mainRunLoop] forMode:NSRunLoopCommonModes];

    // We don't receive background and foreground events from analytics and so we have to listen to
    // them ourselves.
    [[NSNotificationCenter defaultCenter] addObserver:self
                                             selector:@selector(appDidBecomeActiveNotification:)
                                                 name:UIApplicationDidBecomeActiveNotification
                                               object:[UIApplication sharedApplication]];

    [[NSNotificationCenter defaultCenter] addObserver:self
                                             selector:@selector(appWillResignActiveNotification:)
                                                 name:UIApplicationWillResignActiveNotification
                                               object:[UIApplication sharedApplication]];

#if TARGET_OS_TV
    // On tvOS, the refresh rate can change when the user switches display modes or connects to
    // different displays. Listen for mode changes to recompute the slow budget.
    [[NSNotificationCenter defaultCenter] addObserver:self
                                             selector:@selector(screenModeDidChangeNotification:)
                                                 name:UIScreenModeDidChangeNotification
                                               object:nil];
#endif
  }
  return self;
}

- (void)dealloc {
  [_displayLink invalidate];

  [[NSNotificationCenter defaultCenter] removeObserver:self
                                                  name:UIApplicationDidBecomeActiveNotification
                                                object:[UIApplication sharedApplication]];
  [[NSNotificationCenter defaultCenter] removeObserver:self
                                                  name:UIApplicationWillResignActiveNotification
                                                object:[UIApplication sharedApplication]];
#if TARGET_OS_TV
  [[NSNotificationCenter defaultCenter] removeObserver:self
                                                  name:UIScreenModeDidChangeNotification
                                                object:nil];
#endif
}

- (void)appDidBecomeActiveNotification:(NSNotification *)notification {
  // To get the most accurate numbers of total, frozen and slow frames, we need to capture them as
  // soon as we're notified of an event.
  int64_t currentTotalFrames = atomic_load_explicit(&_totalFramesCount, memory_order_relaxed);
  int64_t currentFrozenFrames = atomic_load_explicit(&_frozenFramesCount, memory_order_relaxed);
  int64_t currentSlowFrames = atomic_load_explicit(&_slowFramesCount, memory_order_relaxed);

  dispatch_group_async(self.screenTraceTrackerDispatchGroup, self.screenTraceTrackerSerialQueue, ^{
    for (id viewController in self.previouslyVisibleViewControllers) {
      [self startScreenTraceForViewController:viewController
                           currentTotalFrames:currentTotalFrames
                          currentFrozenFrames:currentFrozenFrames
                            currentSlowFrames:currentSlowFrames];
    }
    self.previouslyVisibleViewControllers = nil;
  });
}

- (void)appWillResignActiveNotification:(NSNotification *)notification {
  // To get the most accurate numbers of total, frozen and slow frames, we need to capture them as
  // soon as we're notified of an event.
  int64_t currentTotalFrames = atomic_load_explicit(&_totalFramesCount, memory_order_relaxed);
  int64_t currentFrozenFrames = atomic_load_explicit(&_frozenFramesCount, memory_order_relaxed);
  int64_t currentSlowFrames = atomic_load_explicit(&_slowFramesCount, memory_order_relaxed);

  dispatch_group_async(self.screenTraceTrackerDispatchGroup, self.screenTraceTrackerSerialQueue, ^{
    self.previouslyVisibleViewControllers = [NSPointerArray weakObjectsPointerArray];
    [self.activeScreenTracesLock lock];
    NSArray<FPRScreenTraceHolder *> *holders = [self.activeScreenTraces allValues];
    for (FPRScreenTraceHolder *holder in holders) {
      UIViewController *vc = holder.viewController;
      if (vc) {
        [self.previouslyVisibleViewControllers addPointer:(__bridge void *)(vc)];
      }
    }
    [self.activeScreenTracesLock unlock];

    for (id visibleViewController in self.previouslyVisibleViewControllers) {
      [self stopScreenTraceForViewController:visibleViewController
                          currentTotalFrames:currentTotalFrames
                         currentFrozenFrames:currentFrozenFrames
                           currentSlowFrames:currentSlowFrames];
    }
  });
}

#if TARGET_OS_TV
/** Handles the UIScreenModeDidChangeNotification. Recomputes the cached slow budget when the screen
 *  mode changes on tvOS.
 *
 *  @param notification The NSNotification object.
 */
- (void)screenModeDidChangeNotification:(NSNotification *)notification {
  dispatch_async(dispatch_get_main_queue(), ^{
    [self updateCachedSlowBudget];
  });
}
#endif

/** Updates the cached slow budget. On tvOS, recomputes from UIScreen.maximumFramesPerSecond.
 *  On iOS, maintains kFPRSlowFrameThreshold.
 *  This method must be called on the main thread - and is available on iOS only for verification.
 */
- (void)updateCachedSlowBudget {
  NSAssert([NSThread isMainThread], @"updateCachedSlowBudget must be called on main thread");
#if TARGET_OS_TV
  UIScreen *mainScreen = [UIScreen mainScreen];
  NSInteger maxFPS = 0;
  if (mainScreen) {
    maxFPS = mainScreen.maximumFramesPerSecond;
  }
  if (maxFPS > 0) {
    _cachedMaxFPS = maxFPS;
    _cachedSlowBudget = 1.0 / (CFTimeInterval)maxFPS;
  } else {
    // Fallback to default FPS if maximumFramesPerSecond is unavailable or invalid.
    _cachedMaxFPS = kFPRDefaultFPS;
    _cachedSlowBudget = 1.0 / (CFTimeInterval)kFPRDefaultFPS;
  }
#else
  // TODO(b/545164487): Updating this is likely the first step for ProMotion devices.
  // https://github.com/firebase/firebase-ios-sdk/issues/10220#issuecomment-1248632304
  _cachedSlowBudget = kFPRSlowFrameThreshold;
#endif
}

#pragma mark - Frozen, slow and good frames

- (void)displayLinkStep {
  static CFAbsoluteTime previousTimestamp = kFPRInvalidTime;
  CFAbsoluteTime currentTimestamp = self.displayLink.timestamp;
  // Use the cached slow budget where needed - currently only on tvOS.
  RecordFrameType(currentTimestamp, previousTimestamp, &_slowFramesCount, &_frozenFramesCount,
                  &_totalFramesCount, _cachedSlowBudget);
  previousTimestamp = currentTimestamp;
}

/** This function increments the relevant frame counters based on the current and previous
 *  timestamp provided by the displayLink.
 *
 *  @param currentTimestamp The current timestamp of the displayLink.
 *  @param previousTimestamp The previous timestamp of the displayLink.
 *  @param slowFramesCounter The value of the slowFramesCount before this function was called.
 *  @param frozenFramesCounter The value of the frozenFramesCount before this function was called.
 *  @param totalFramesCounter The value of the totalFramesCount before this function was called.
 *  @param slowBudget The frame duration threshold above which a frame is considered slow.
 */
FOUNDATION_STATIC_INLINE
void RecordFrameType(CFAbsoluteTime currentTimestamp,
                     CFAbsoluteTime previousTimestamp,
                     atomic_int_fast64_t *slowFramesCounter,
                     atomic_int_fast64_t *frozenFramesCounter,
                     atomic_int_fast64_t *totalFramesCounter,
                     CFTimeInterval slowBudget) {
  CFTimeInterval frameDuration = currentTimestamp - previousTimestamp;
  if (previousTimestamp == kFPRInvalidTime) {
    return;
  }
#if TARGET_OS_TV
  // On tvOS, use cached slowBudget with epsilon to avoid floating point comparison issues
  // (e.g., 59.94 vs 60 Hz displays).
  if (frameDuration > slowBudget + kFPRSlowFrameEpsilon) {
    atomic_fetch_add_explicit(slowFramesCounter, 1, memory_order_relaxed);
  }
#else
  if (frameDuration > kFPRSlowFrameThreshold) {
    atomic_fetch_add_explicit(slowFramesCounter, 1, memory_order_relaxed);
  }
#endif
  if (frameDuration > kFPRFrozenFrameThreshold) {
    atomic_fetch_add_explicit(frozenFramesCounter, 1, memory_order_relaxed);
  }
  atomic_fetch_add_explicit(totalFramesCounter, 1, memory_order_relaxed);
}

#pragma mark - Helper methods

/** Starts a screen trace for the given UIViewController instance if it doesn't exist. This method
 *  does NOT ensure thread safety - the caller is responsible for making sure that this is invoked
 *  in a thread safe manner.
 *
 *  @param viewController The UIViewController instance for which the trace is to be started.
 *  @param currentTotalFrames The value of the totalFramesCount before this method was called.
 *  @param currentFrozenFrames The value of the frozenFramesCount before this method was called.
 *  @param currentSlowFrames The value of the slowFramesCount before this method was called.
 */
- (void)startScreenTraceForViewController:(UIViewController *)viewController
                       currentTotalFrames:(int64_t)currentTotalFrames
                      currentFrozenFrames:(int64_t)currentFrozenFrames
                        currentSlowFrames:(int64_t)currentSlowFrames {
  if (![self shouldCreateScreenTraceForViewController:viewController]) {
    return;
  }

  [self.activeScreenTracesLock lock];
  [self cleanupStaleTraces];

  NSValue *key = [NSValue valueWithNonretainedObject:viewController];
  FPRScreenTraceHolder *holder = [self.activeScreenTraces objectForKey:key];
  if (holder && holder.viewController != viewController) {
    // Stale entry due to pointer reuse. Remove it.
    [self.activeScreenTraces removeObjectForKey:key];
    holder = nil;
  }

  if (!holder) {
    NSString *traceName = FPRScreenTraceNameForViewController(viewController);
    FIRTrace *newTrace = [[FIRTrace alloc] initInternalTraceWithName:traceName];
    [newTrace start];
    [newTrace setIntValue:currentTotalFrames forMetric:kFPRTotalFramesCounterName];
    [newTrace setIntValue:currentFrozenFrames forMetric:kFPRFrozenFrameCounterName];
    [newTrace setIntValue:currentSlowFrames forMetric:kFPRSlowFrameCounterName];

    holder = [[FPRScreenTraceHolder alloc] init];
    holder.viewController = viewController;
    holder.trace = newTrace;

    [self.activeScreenTraces setObject:holder forKey:key];
  }
  [self.activeScreenTracesLock unlock];
}

/** Stops a screen trace for the given UIViewController instance if it exist. This method does NOT
 *  ensure thread safety - the caller is responsible for making sure that this is invoked in a
 *  thread safe manner.
 *
 *  @param viewController The UIViewController instance for which the trace is to be stopped.
 *  @param currentTotalFrames The value of the totalFramesCount before this method was called.
 *  @param currentFrozenFrames The value of the frozenFramesCount before this method was called.
 *  @param currentSlowFrames The value of the slowFramesCount before this method was called.
 */
- (void)stopScreenTraceForViewController:(UIViewController *)viewController
                      currentTotalFrames:(int64_t)currentTotalFrames
                     currentFrozenFrames:(int64_t)currentFrozenFrames
                       currentSlowFrames:(int64_t)currentSlowFrames {
  if (viewController == nil) {
    return;
  }
  NSValue *key = [NSValue valueWithNonretainedObject:viewController];

  [self.activeScreenTracesLock lock];
  FPRScreenTraceHolder *holder = [self.activeScreenTraces objectForKey:key];
  if (holder) {
    [self.activeScreenTraces removeObjectForKey:key];
    if (holder.viewController != viewController) {
      // Stale entry due to pointer reuse.
      holder = nil;
    }
  }
  [self.activeScreenTracesLock unlock];

  if (!holder) {
    return;
  }

  FIRTrace *previousScreenTrace = holder.trace;

  // Get a diff between the counters now and what they were at trace start.
  int64_t actualTotalFrames =
      currentTotalFrames - [previousScreenTrace valueForIntMetric:kFPRTotalFramesCounterName];
  int64_t actualFrozenFrames =
      currentFrozenFrames - [previousScreenTrace valueForIntMetric:kFPRFrozenFrameCounterName];
  int64_t actualSlowFrames =
      currentSlowFrames - [previousScreenTrace valueForIntMetric:kFPRSlowFrameCounterName];

  // Update the values in the trace.
  if (actualTotalFrames != 0) {
    [previousScreenTrace setIntValue:actualTotalFrames forMetric:kFPRTotalFramesCounterName];
  } else {
    [previousScreenTrace deleteMetric:kFPRTotalFramesCounterName];
  }

  if (actualFrozenFrames != 0) {
    [previousScreenTrace setIntValue:actualFrozenFrames forMetric:kFPRFrozenFrameCounterName];
  } else {
    [previousScreenTrace deleteMetric:kFPRFrozenFrameCounterName];
  }

  if (actualSlowFrames != 0) {
    [previousScreenTrace setIntValue:actualSlowFrames forMetric:kFPRSlowFrameCounterName];
  } else {
    [previousScreenTrace deleteMetric:kFPRSlowFrameCounterName];
  }

  if (previousScreenTrace.numberOfCounters > 0) {
    [previousScreenTrace stop];
  } else {
    // The trace did not collect any data. Don't log it.
    [previousScreenTrace cancel];
  }
}

#pragma mark - Filtering for screen traces

/** Determines whether to create a screen trace for the given UIViewController instance.
 *
 *  @param viewController The UIViewController instance.
 *  @return YES if a screen trace should be created for the given UIViewController instance,
        NO otherwise.
 */
- (BOOL)shouldCreateScreenTraceForViewController:(UIViewController *)viewController {
  if (viewController == nil) {
    return NO;
  }

  // Ignore non-main bundle view controllers whose class or superclass is an internal iOS view
  // controller. This is borrowed from the logic for tracking screens in Firebase Analytics.
  NSBundle *bundle = [NSBundle bundleForClass:[viewController class]];
  if (bundle != [NSBundle mainBundle]) {
    NSString *className = FPRUnprefixedClassName([viewController class]);
    if ([className hasPrefix:@"_"]) {
      return NO;
    }
    NSString *superClassName = FPRUnprefixedClassName([viewController superclass]);
    if ([superClassName hasPrefix:@"_"]) {
      return NO;
    }
  }

  // We are not creating screen traces for these view controllers because they're container view
  // controllers. They always have a child view controller which will provide better context for a
  // screen trace. We are however capturing traces if a developer subclasses these as there may be
  // some context. Special case: We are not capturing screen traces for any input view
  // controllers.
  return !([viewController isMemberOfClass:[UINavigationController class]] ||
           [viewController isMemberOfClass:[UITabBarController class]] ||
           [viewController isMemberOfClass:[UISplitViewController class]] ||
           [viewController isMemberOfClass:[UIPageViewController class]] ||
           [viewController isKindOfClass:[UIInputViewController class]]);
}

- (void)cleanupStaleTraces {
  [self.activeScreenTracesLock lock];
  NSMutableArray *keysToRemove = [NSMutableArray array];
  [self.activeScreenTraces
      enumerateKeysAndObjectsUsingBlock:^(NSValue *key, FPRScreenTraceHolder *holder, BOOL *stop) {
        if (holder.viewController == nil) {
          [keysToRemove addObject:key];
        }
      }];
  [self.activeScreenTraces removeObjectsForKeys:keysToRemove];
  [self.activeScreenTracesLock unlock];
}

#pragma mark - Screen Traces swizzling hooks

- (void)viewControllerDidAppear:(UIViewController *)viewController {
  // To get the most accurate numbers of total, frozen and slow frames, we need to capture them as
  // soon as we're notified of an event.
  int64_t currentTotalFrames = atomic_load_explicit(&_totalFramesCount, memory_order_relaxed);
  int64_t currentFrozenFrames = atomic_load_explicit(&_frozenFramesCount, memory_order_relaxed);
  int64_t currentSlowFrames = atomic_load_explicit(&_slowFramesCount, memory_order_relaxed);

  dispatch_sync(self.screenTraceTrackerSerialQueue, ^{
    [self startScreenTraceForViewController:viewController
                         currentTotalFrames:currentTotalFrames
                        currentFrozenFrames:currentFrozenFrames
                          currentSlowFrames:currentSlowFrames];
  });
}

- (void)viewControllerDidDisappear:(id)viewController {
  // To get the most accurate numbers of total, frozen and slow frames, we need to capture them as
  // soon as we're notified of an event.
  int64_t currentTotalFrames = atomic_load_explicit(&_totalFramesCount, memory_order_relaxed);
  int64_t currentFrozenFrames = atomic_load_explicit(&_frozenFramesCount, memory_order_relaxed);
  int64_t currentSlowFrames = atomic_load_explicit(&_slowFramesCount, memory_order_relaxed);

  dispatch_sync(self.screenTraceTrackerSerialQueue, ^{
    [self stopScreenTraceForViewController:viewController
                        currentTotalFrames:currentTotalFrames
                       currentFrozenFrames:currentFrozenFrames
                         currentSlowFrames:currentSlowFrames];
  });
}

#pragma mark - Test Helper Methods

- (int_fast64_t)totalFramesCount {
  return atomic_load_explicit(&_totalFramesCount, memory_order_relaxed);
}

- (void)setTotalFramesCount:(int_fast64_t)count {
  atomic_store_explicit(&_totalFramesCount, count, memory_order_relaxed);
}

- (int_fast64_t)slowFramesCount {
  return atomic_load_explicit(&_slowFramesCount, memory_order_relaxed);
}

- (void)setSlowFramesCount:(int_fast64_t)count {
  atomic_store_explicit(&_slowFramesCount, count, memory_order_relaxed);
}

- (int_fast64_t)frozenFramesCount {
  return atomic_load_explicit(&_frozenFramesCount, memory_order_relaxed);
}

- (void)setFrozenFramesCount:(int_fast64_t)count {
  atomic_store_explicit(&_frozenFramesCount, count, memory_order_relaxed);
}

@end