#import "ABOrderedKeyArray.h"

@interface ABOrderedKeyArray ()

@property(nonatomic, strong) NSMutableArray<id> *array;
@property(nonatomic, strong) NSMutableDictionary<id, id> *dictionary;
@property(nonatomic, assign) ABOrderedKeyArrayComparator comparator;

@end

@implementation ABOrderedKeyArray

- (instancetype)initWithComparator:(ABOrderedKeyArrayComparator)comparator {
  self = [super init];
  if (self) {
    self.comparator = comparator;
  }
  return self;
}

- (nullable id)objectForKey:(id)key {
  return [self.dictionary objectForKey:key];
}

- (void)insertObject:(id)object forKey:(id)key {
  __auto_type index = [self indexForObject:object];
  if (index == NSNotFound) {
    return;
  }
  [self.array insertObject:object atIndex:index];
}

- (BOOL)removeObjectForKey:(id)key {
  id object = [self.dictionary objectForKey:key];
  if (!object) {
    return NO;
  }
  __auto_type index = [self indexForObject:object];
  if (index == NSNotFound) {
    NSAssert(false, @"Object found in dictionary but not found in array");
    return NO;
  }
  [self.dictionary removeObjectForKey:key];
  [self.array removeObjectAtIndex:index];
  return YES;
}

- (NSUInteger)indexForObject:(id)object {
  return [self.array indexOfObject:object
                     inSortedRange:NSMakeRange(0, self.array.count)
                           options:NSBinarySearchingInsertionIndex |
                                   NSBinarySearchingLastEqual
                   usingComparator:self.comparator];
}

#pragma mark - NSFastEnumeration

- (NSUInteger)
    countByEnumeratingWithState:(nonnull NSFastEnumerationState *)state
                        objects:
                            (__unsafe_unretained id _Nullable *_Nonnull)buffer
                          count:(NSUInteger)len {
  return [self.array countByEnumeratingWithState:state
                                         objects:buffer
                                           count:len];
}

@end