/*
 * 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 "FIndexedNode.h"

#import "FChildrenNode.h"
#import "FImmutableSortedSet.h"
#import "FIndex.h"
#import "FKeyIndex.h"
#import "FPriorityIndex.h"

static FImmutableSortedSet *FALLBACK_INDEX;

@interface FIndexedNode ()

@property(nonatomic, strong) id<FNode> node;
/**
 * The indexed set is initialized lazily to prevent creation when it is not
 * needed
 */
@property(nonatomic, strong) FImmutableSortedSet *indexed;
@property(nonatomic, strong) id<FIndex> index;

@end

@implementation FIndexedNode

+ (FImmutableSortedSet *)fallbackIndex {
    static FImmutableSortedSet *fallbackIndex;
    static dispatch_once_t once;
    dispatch_once(&once, ^{
      fallbackIndex = [[FImmutableSortedSet alloc] init];
    });
    return fallbackIndex;
}

+ (FIndexedNode *)indexedNodeWithNode:(id<FNode>)node {
    return [[FIndexedNode alloc] initWithNode:node
                                        index:[FPriorityIndex priorityIndex]];
}

+ (FIndexedNode *)indexedNodeWithNode:(id<FNode>)node index:(id<FIndex>)index {
    return [[FIndexedNode alloc] initWithNode:node index:index];
}

- (id)initWithNode:(id<FNode>)node index:(id<FIndex>)index {
    // Initialize indexed lazily
    return [self initWithNode:node index:index indexed:nil];
}

- (id)initWithNode:(id<FNode>)node
             index:(id<FIndex>)index
           indexed:(FImmutableSortedSet *)indexed {
    self = [super init];
    if (self != nil) {
        self->_node = node;
        self->_index = index;
        self->_indexed = indexed;
    }
    return self;
}

- (void)ensureIndexed {
    if (!self.indexed) {
        if ([self.index isEqual:[FKeyIndex keyIndex]]) {
            self.indexed = [FIndexedNode fallbackIndex];
        } else {
            __block BOOL sawChild = NO;
            [self.node enumerateChildrenUsingBlock:^(
                           NSString *key, id<FNode> node, BOOL *stop) {
              sawChild = sawChild || [self.index isDefinedOn:node];
              *stop = sawChild;
            }];
            if (sawChild) {
                NSMutableDictionary *dict = [NSMutableDictionary dictionary];
                [self.node enumerateChildrenUsingBlock:^(
                               NSString *key, id<FNode> node, BOOL *stop) {
                  FNamedNode *namedNode =
                      [[FNamedNode alloc] initWithName:key andNode:node];
                  dict[namedNode] = [NSNull null];
                }];
                // Make sure to assign index here, because the comparator will
                // be retained and using self will cause a cycle
                id<FIndex> index = self.index;
                self.indexed = [FImmutableSortedSet
                    setWithKeysFromDictionary:dict
                                   comparator:^NSComparisonResult(
                                       FNamedNode *namedNode1,
                                       FNamedNode *namedNode2) {
                                     return [index compareNamedNode:namedNode1
                                                        toNamedNode:namedNode2];
                                   }];
            } else {
                self.indexed = [FIndexedNode fallbackIndex];
            }
        }
    }
}

- (BOOL)hasIndex:(id<FIndex>)index {
    return [self.index isEqual:index];
}

- (FIndexedNode *)updateChild:(NSString *)key
                 withNewChild:(id<FNode>)newChildNode {
    id<FNode> newNode = [self.node updateImmediateChild:key
                                           withNewChild:newChildNode];
    if (self.indexed == [FIndexedNode fallbackIndex] &&
        ![self.index isDefinedOn:newChildNode]) {
        // doesn't affect the index, no need to create an index
        return [[FIndexedNode alloc] initWithNode:newNode
                                            index:self.index
                                          indexed:[FIndexedNode fallbackIndex]];
    } else if (!self.indexed || self.indexed == [FIndexedNode fallbackIndex]) {
        // No need to index yet, index lazily
        return [[FIndexedNode alloc] initWithNode:newNode index:self.index];
    } else {
        id<FNode> oldChild = [self.node getImmediateChild:key];
        FImmutableSortedSet *newIndexed = [self.indexed
            removeObject:[FNamedNode nodeWithName:key node:oldChild]];
        if (![newChildNode isEmpty]) {
            newIndexed = [newIndexed
                addObject:[FNamedNode nodeWithName:key node:newChildNode]];
        }
        return [[FIndexedNode alloc] initWithNode:newNode
                                            index:self.index
                                          indexed:newIndexed];
    }
}

- (FIndexedNode *)updatePriority:(id<FNode>)priority {
    return
        [[FIndexedNode alloc] initWithNode:[self.node updatePriority:priority]
                                     index:self.index
                                   indexed:self.indexed];
}

- (FNamedNode *)firstChild {
    if (![self.node isKindOfClass:[FChildrenNode class]]) {
        return nil;
    } else {
        [self ensureIndexed];
        if (self.indexed == [FIndexedNode fallbackIndex]) {
            return [((FChildrenNode *)self.node) firstChild];
        } else {
            return self.indexed.firstObject;
        }
    }
}

- (FNamedNode *)lastChild {
    if (![self.node isKindOfClass:[FChildrenNode class]]) {
        return nil;
    } else {
        [self ensureIndexed];
        if (self.indexed == [FIndexedNode fallbackIndex]) {
            return [((FChildrenNode *)self.node) lastChild];
        } else {
            return self.indexed.lastObject;
        }
    }
}

- (NSString *)predecessorForChildKey:(NSString *)childKey
                           childNode:(id<FNode>)childNode
                               index:(id<FIndex>)index {
    if (![self.index isEqual:index]) {
        [NSException raise:NSInvalidArgumentException
                    format:@"Index not available in IndexedNode!"];
    }
    [self ensureIndexed];
    if (self.indexed == [FIndexedNode fallbackIndex]) {
        return [self.node predecessorChildKey:childKey];
    } else {
        FNamedNode *node = [self.indexed
            predecessorEntry:[FNamedNode nodeWithName:childKey node:childNode]];
        return node.name;
    }
}

- (void)enumerateChildrenReverse:(BOOL)reverse
                      usingBlock:
                          (void (^)(NSString *, id<FNode>, BOOL *))block {
    [self ensureIndexed];
    if (self.indexed == [FIndexedNode fallbackIndex]) {
        [self.node enumerateChildrenReverse:reverse usingBlock:block];
    } else {
        [self.indexed
            enumerateObjectsReverse:reverse
                         usingBlock:^(FNamedNode *namedNode, BOOL *stop) {
                           block(namedNode.name, namedNode.node, stop);
                         }];
    }
}

- (NSEnumerator *)childEnumerator {
    [self ensureIndexed];
    if (self.indexed == [FIndexedNode fallbackIndex]) {
        return [self.node childEnumerator];
    } else {
        return [self.indexed objectEnumerator];
    }
}

@end
