/*
 * 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.
 */

#import "Firestore/Source/Remote/FSTRemoteEvent.h"

#import "Firestore/Source/Core/FSTSnapshotVersion.h"
#import "Firestore/Source/Model/FSTDocument.h"
#import "Firestore/Source/Model/FSTDocumentKey.h"
#import "Firestore/Source/Remote/FSTWatchChange.h"
#import "Firestore/Source/Util/FSTAssert.h"
#import "Firestore/Source/Util/FSTClasses.h"
#import "Firestore/Source/Util/FSTLogger.h"

#include "Firestore/core/src/firebase/firestore/model/document_key.h"

NS_ASSUME_NONNULL_BEGIN

#pragma mark - FSTTargetMapping

@interface FSTTargetMapping ()

/** Private mutator method to add a document key to the mapping */
- (void)addDocumentKey:(FSTDocumentKey *)documentKey;

/** Private mutator method to remove a document key from the mapping */
- (void)removeDocumentKey:(FSTDocumentKey *)documentKey;

@end

@implementation FSTTargetMapping

- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
  @throw FSTAbstractMethodException();  // NOLINT
}

- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
  @throw FSTAbstractMethodException();  // NOLINT
}

@end

#pragma mark - FSTResetMapping

@interface FSTResetMapping ()
@property(nonatomic, strong) FSTDocumentKeySet *documents;
@end

@implementation FSTResetMapping

+ (instancetype)mappingWithDocuments:(NSArray<FSTDocument *> *)documents {
  FSTResetMapping *mapping = [[FSTResetMapping alloc] init];
  for (FSTDocument *doc in documents) {
    mapping.documents = [mapping.documents setByAddingObject:doc.key];
  }
  return mapping;
}

- (instancetype)init {
  self = [super init];
  if (self) {
    _documents = [FSTDocumentKeySet keySet];
  }
  return self;
}

- (BOOL)isEqual:(id)other {
  if (other == self) {
    return YES;
  }
  if (![other isMemberOfClass:[FSTResetMapping class]]) {
    return NO;
  }

  FSTResetMapping *otherMapping = (FSTResetMapping *)other;
  return [self.documents isEqual:otherMapping.documents];
}

- (NSUInteger)hash {
  return self.documents.hash;
}

- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
  self.documents = [self.documents setByAddingObject:documentKey];
}

- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
  self.documents = [self.documents setByRemovingObject:documentKey];
}

@end

#pragma mark - FSTUpdateMapping

@interface FSTUpdateMapping ()
@property(nonatomic, strong) FSTDocumentKeySet *addedDocuments;
@property(nonatomic, strong) FSTDocumentKeySet *removedDocuments;
@end

@implementation FSTUpdateMapping

+ (FSTUpdateMapping *)mappingWithAddedDocuments:(NSArray<FSTDocument *> *)added
                               removedDocuments:(NSArray<FSTDocument *> *)removed {
  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
  for (FSTDocument *doc in added) {
    mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
  }
  for (FSTDocument *doc in removed) {
    mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
  }
  return mapping;
}

- (instancetype)init {
  self = [super init];
  if (self) {
    _addedDocuments = [FSTDocumentKeySet keySet];
    _removedDocuments = [FSTDocumentKeySet keySet];
  }
  return self;
}

- (BOOL)isEqual:(id)other {
  if (other == self) {
    return YES;
  }
  if (![other isMemberOfClass:[FSTUpdateMapping class]]) {
    return NO;
  }

  FSTUpdateMapping *otherMapping = (FSTUpdateMapping *)other;
  return [self.addedDocuments isEqual:otherMapping.addedDocuments] &&
         [self.removedDocuments isEqual:otherMapping.removedDocuments];
}

- (NSUInteger)hash {
  return self.addedDocuments.hash * 31 + self.removedDocuments.hash;
}

- (FSTDocumentKeySet *)applyTo:(FSTDocumentKeySet *)keys {
  __block FSTDocumentKeySet *result = keys;
  [self.addedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
    result = [result setByAddingObject:key];
  }];
  [self.removedDocuments enumerateObjectsUsingBlock:^(FSTDocumentKey *key, BOOL *stop) {
    result = [result setByRemovingObject:key];
  }];
  return result;
}

- (void)addDocumentKey:(FSTDocumentKey *)documentKey {
  self.addedDocuments = [self.addedDocuments setByAddingObject:documentKey];
  self.removedDocuments = [self.removedDocuments setByRemovingObject:documentKey];
}

- (void)removeDocumentKey:(FSTDocumentKey *)documentKey {
  self.addedDocuments = [self.addedDocuments setByRemovingObject:documentKey];
  self.removedDocuments = [self.removedDocuments setByAddingObject:documentKey];
}

@end

#pragma mark - FSTTargetChange

@interface FSTTargetChange ()
@property(nonatomic, assign) FSTCurrentStatusUpdate currentStatusUpdate;
@property(nonatomic, strong, nullable) FSTTargetMapping *mapping;
@property(nonatomic, strong) FSTSnapshotVersion *snapshotVersion;
@property(nonatomic, strong) NSData *resumeToken;
@end

@implementation FSTTargetChange

- (instancetype)init {
  if (self = [super init]) {
    _currentStatusUpdate = FSTCurrentStatusUpdateNone;
    _resumeToken = [NSData data];
  }
  return self;
}

+ (instancetype)changeWithDocuments:(NSArray<FSTMaybeDocument *> *)docs
                currentStatusUpdate:(FSTCurrentStatusUpdate)currentStatusUpdate {
  FSTUpdateMapping *mapping = [[FSTUpdateMapping alloc] init];
  for (FSTMaybeDocument *doc in docs) {
    if ([doc isKindOfClass:[FSTDeletedDocument class]]) {
      mapping.removedDocuments = [mapping.removedDocuments setByAddingObject:doc.key];
    } else {
      mapping.addedDocuments = [mapping.addedDocuments setByAddingObject:doc.key];
    }
  }
  FSTTargetChange *change = [[FSTTargetChange alloc] init];
  change.mapping = mapping;
  change.currentStatusUpdate = currentStatusUpdate;
  return change;
}

+ (instancetype)changeWithMapping:(FSTTargetMapping *)mapping
                  snapshotVersion:(FSTSnapshotVersion *)snapshotVersion
              currentStatusUpdate:(FSTCurrentStatusUpdate)currentStatusUpdate {
  FSTTargetChange *change = [[FSTTargetChange alloc] init];
  change.mapping = mapping;
  change.snapshotVersion = snapshotVersion;
  change.currentStatusUpdate = currentStatusUpdate;
  return change;
}

- (FSTTargetMapping *)mapping {
  if (!_mapping) {
    // Create an FSTUpdateMapping by default, since resets are always explicit
    _mapping = [[FSTUpdateMapping alloc] init];
  }
  return _mapping;
}

/**
 * Sets the resume token but only when it has a new value. Empty resumeTokens are
 * discarded.
 */
- (void)setResumeToken:(NSData *)resumeToken {
  if (resumeToken.length > 0) {
    _resumeToken = resumeToken;
  }
}

@end

#pragma mark - FSTRemoteEvent

@interface FSTRemoteEvent () {
  NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *_documentUpdates;
  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *_targetChanges;
}

- (instancetype)
initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
          targetChanges:(NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges
        documentUpdates:
            (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates;

@property(nonatomic, strong) FSTSnapshotVersion *snapshotVersion;

@end

@implementation FSTRemoteEvent

+ (instancetype)
eventWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
           targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
         documentUpdates:
             (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
  return [[FSTRemoteEvent alloc] initWithSnapshotVersion:snapshotVersion
                                           targetChanges:targetChanges
                                         documentUpdates:documentUpdates];
}

- (instancetype)
initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
          targetChanges:(NSMutableDictionary<NSNumber *, FSTTargetChange *> *)targetChanges
        documentUpdates:
            (NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
  self = [super init];
  if (self) {
    _snapshotVersion = snapshotVersion;
    _targetChanges = targetChanges;
    _documentUpdates = documentUpdates;
  }
  return self;
}

- (NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *)targetChanges {
  return static_cast<NSDictionary<FSTBoxedTargetID *, FSTTargetChange *> *>(_targetChanges);
}

- (NSDictionary<FSTDocumentKey *, FSTMaybeDocument *> *)documentUpdates {
  return static_cast<NSDictionary<FSTDocumentKey *, FSTMaybeDocument *> *>(_documentUpdates);
}

/** Adds a document update to this remote event */
- (void)addDocumentUpdate:(FSTMaybeDocument *)document {
  _documentUpdates[(FSTDocumentKey *)document.key] = document;
}

/** Handles an existence filter mismatch */
- (void)handleExistenceFilterMismatchForTargetID:(FSTBoxedTargetID *)targetID {
  // An existence filter mismatch will reset the query and we need to reset the mapping to contain
  // no documents and an empty resume token.
  //
  // Note:
  //   * The reset mapping is empty, specifically forcing the consumer of the change to
  //     forget all keys for this targetID;
  //   * The resume snapshot for this target must be reset
  //   * The target must be unacked because unwatching and rewatching introduces a race for
  //     changes.
  //
  // TODO(dimond): keep track of reset targets not to raise.
  FSTTargetChange *targetChange =
      [FSTTargetChange changeWithMapping:[[FSTResetMapping alloc] init]
                         snapshotVersion:[FSTSnapshotVersion noVersion]
                     currentStatusUpdate:FSTCurrentStatusUpdateMarkNotCurrent];
  _targetChanges[targetID] = targetChange;
}

@end

#pragma mark - FSTWatchChangeAggregator

@interface FSTWatchChangeAggregator ()

/** The snapshot version for every target change this creates. */
@property(nonatomic, strong, readonly) FSTSnapshotVersion *snapshotVersion;

/** Keeps track of the current target mappings */
@property(nonatomic, strong, readonly)
    NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *targetChanges;

/** Keeps track of document to update */
@property(nonatomic, strong, readonly)
    NSMutableDictionary<FSTDocumentKey *, FSTMaybeDocument *> *documentUpdates;

/** The set of open listens on the client */
@property(nonatomic, strong, readonly)
    NSDictionary<FSTBoxedTargetID *, FSTQueryData *> *listenTargets;

/** Whether this aggregator was frozen and can no longer be modified */
@property(nonatomic, assign) BOOL frozen;

@end

@implementation FSTWatchChangeAggregator {
  NSMutableDictionary<FSTBoxedTargetID *, FSTExistenceFilter *> *_existenceFilters;
}

- (instancetype)
initWithSnapshotVersion:(FSTSnapshotVersion *)snapshotVersion
          listenTargets:(NSDictionary<FSTBoxedTargetID *, FSTQueryData *> *)listenTargets
 pendingTargetResponses:(NSDictionary<FSTBoxedTargetID *, NSNumber *> *)pendingTargetResponses {
  self = [super init];
  if (self) {
    _snapshotVersion = snapshotVersion;

    _frozen = NO;
    _targetChanges = [NSMutableDictionary dictionary];
    _listenTargets = listenTargets;
    _pendingTargetResponses = [NSMutableDictionary dictionaryWithDictionary:pendingTargetResponses];

    _existenceFilters = [NSMutableDictionary dictionary];
    _documentUpdates = [NSMutableDictionary dictionary];
  }
  return self;
}

- (NSDictionary<FSTBoxedTargetID *, FSTExistenceFilter *> *)existenceFilters {
  return static_cast<NSDictionary<FSTBoxedTargetID *, FSTExistenceFilter *> *>(_existenceFilters);
}

- (FSTTargetChange *)targetChangeForTargetID:(FSTBoxedTargetID *)targetID {
  FSTTargetChange *change = self.targetChanges[targetID];
  if (!change) {
    change = [[FSTTargetChange alloc] init];
    change.snapshotVersion = self.snapshotVersion;
    self.targetChanges[targetID] = change;
  }
  return change;
}

- (void)addWatchChanges:(NSArray<FSTWatchChange *> *)watchChanges {
  FSTAssert(!self.frozen, @"Trying to modify frozen FSTWatchChangeAggregator");
  for (FSTWatchChange *watchChange in watchChanges) {
    [self addWatchChange:watchChange];
  }
}

- (void)addWatchChange:(FSTWatchChange *)watchChange {
  FSTAssert(!self.frozen, @"Trying to modify frozen FSTWatchChangeAggregator");
  if ([watchChange isKindOfClass:[FSTDocumentWatchChange class]]) {
    [self addDocumentChange:(FSTDocumentWatchChange *)watchChange];
  } else if ([watchChange isKindOfClass:[FSTWatchTargetChange class]]) {
    [self addTargetChange:(FSTWatchTargetChange *)watchChange];
  } else if ([watchChange isKindOfClass:[FSTExistenceFilterWatchChange class]]) {
    [self addExistenceFilterChange:(FSTExistenceFilterWatchChange *)watchChange];
  } else {
    FSTFail(@"Unknown watch change: %@", watchChange);
  }
}

- (void)addDocumentChange:(FSTDocumentWatchChange *)docChange {
  BOOL relevant = NO;

  for (FSTBoxedTargetID *targetID in docChange.updatedTargetIDs) {
    if ([self isActiveTarget:targetID]) {
      FSTTargetChange *change = [self targetChangeForTargetID:targetID];
      [change.mapping addDocumentKey:docChange.documentKey];
      relevant = YES;
    }
  }

  for (FSTBoxedTargetID *targetID in docChange.removedTargetIDs) {
    if ([self isActiveTarget:targetID]) {
      FSTTargetChange *change = [self targetChangeForTargetID:targetID];
      [change.mapping removeDocumentKey:docChange.documentKey];
      relevant = YES;
    }
  }

  // Only update the document if there is a new document to replace, this might be just a target
  // update instead.
  if (docChange.document && relevant) {
    self.documentUpdates[docChange.documentKey] = docChange.document;
  }
}

- (void)addTargetChange:(FSTWatchTargetChange *)targetChange {
  for (FSTBoxedTargetID *targetID in targetChange.targetIDs) {
    FSTTargetChange *change = [self targetChangeForTargetID:targetID];
    switch (targetChange.state) {
      case FSTWatchTargetChangeStateNoChange:
        if ([self isActiveTarget:targetID]) {
          // Creating the change above satisfies the semantics of no-change.
          change.resumeToken = targetChange.resumeToken;
        }
        break;
      case FSTWatchTargetChangeStateAdded:
        [self recordResponseForTargetID:targetID];
        if (![self.pendingTargetResponses objectForKey:targetID]) {
          // We have a freshly added target, so we need to reset any state that we had previously
          // This can happen e.g. when remove and add back a target for existence filter
          // mismatches.
          change.mapping = nil;
          change.currentStatusUpdate = FSTCurrentStatusUpdateNone;
          [_existenceFilters removeObjectForKey:targetID];
        }
        change.resumeToken = targetChange.resumeToken;
        break;
      case FSTWatchTargetChangeStateRemoved:
        // We need to keep track of removed targets to we can post-filter and remove any target
        // changes.
        [self recordResponseForTargetID:targetID];
        FSTAssert(!targetChange.cause, @"WatchChangeAggregator does not handle errored targets.");
        break;
      case FSTWatchTargetChangeStateCurrent:
        if ([self isActiveTarget:targetID]) {
          change.currentStatusUpdate = FSTCurrentStatusUpdateMarkCurrent;
          change.resumeToken = targetChange.resumeToken;
        }
        break;
      case FSTWatchTargetChangeStateReset:
        if ([self isActiveTarget:targetID]) {
          // Overwrite any existing target mapping with a reset mapping. Every subsequent update
          // will modify the reset mapping, not an update mapping.
          change.mapping = [[FSTResetMapping alloc] init];
          change.resumeToken = targetChange.resumeToken;
        }
        break;
      default:
        FSTWarn(@"Unknown target watch change type: %ld", (long)targetChange.state);
    }
  }
}

/**
 * Records that we got a watch target add/remove by decrementing the number of pending target
 * responses that we have.
 */
- (void)recordResponseForTargetID:(FSTBoxedTargetID *)targetID {
  NSNumber *count = [self.pendingTargetResponses objectForKey:targetID];
  int newCount = count ? [count intValue] - 1 : -1;
  if (newCount == 0) {
    [self.pendingTargetResponses removeObjectForKey:targetID];
  } else {
    [self.pendingTargetResponses setObject:[NSNumber numberWithInt:newCount] forKey:targetID];
  }
}

/**
 * Returns true if the given targetId is active. Active targets are those for which there are no
 * pending requests to add a listen and are in the current list of targets the client cares about.
 *
 * Clients can repeatedly listen and stop listening to targets, so this check is useful in
 * preventing in preventing race conditions for a target where events arrive but the server hasn't
 * yet acknowledged the intended change in state.
 */
- (BOOL)isActiveTarget:(FSTBoxedTargetID *)targetID {
  return [self.listenTargets objectForKey:targetID] &&
         ![self.pendingTargetResponses objectForKey:targetID];
}

- (void)addExistenceFilterChange:(FSTExistenceFilterWatchChange *)existenceFilterChange {
  FSTBoxedTargetID *targetID = @(existenceFilterChange.targetID);
  if ([self isActiveTarget:targetID]) {
    _existenceFilters[targetID] = existenceFilterChange.filter;
  }
}

- (FSTRemoteEvent *)remoteEvent {
  NSMutableDictionary<FSTBoxedTargetID *, FSTTargetChange *> *targetChanges = self.targetChanges;

  NSMutableArray *targetsToRemove = [NSMutableArray array];

  // Apply any inactive targets.
  for (FSTBoxedTargetID *targetID in [targetChanges keyEnumerator]) {
    if (![self isActiveTarget:targetID]) {
      [targetsToRemove addObject:targetID];
    }
  }

  [targetChanges removeObjectsForKeys:targetsToRemove];

  // Mark this aggregator as frozen so no further modifications are made.
  self.frozen = YES;
  return [FSTRemoteEvent eventWithSnapshotVersion:self.snapshotVersion
                                    targetChanges:targetChanges
                                  documentUpdates:self.documentUpdates];
}

@end

NS_ASSUME_NONNULL_END