#import "ScreenAnimationController.h"
#import "DisplayLinkAnimator.h"
#import "RNNScreenTransitionsCreator.h"
#import "SharedElementAnimator.h"
#import "UIViewController+LayoutProtocol.h"

#ifdef RCT_NEW_ARCH_ENABLED
#include <React/RCTSurfacePresenter.h>
#include <React/RCTMountingManager.h>
#endif

@implementation ScreenAnimationController {
	RCTBridge *_bridge;
#ifdef RCT_NEW_ARCH_ENABLED
	RCTHost *_host;
#endif
	id<UIViewControllerContextTransitioning> _transitionContext;
	SharedElementAnimator *_sharedElementAnimator;
	BOOL _animate;
	CGFloat _duration;
}

- (instancetype)initWithContentTransition:(RNNEnterExitAnimation *)contentTransition
					   elementTransitions:(NSArray<ElementTransitionOptions *> *)elementTransitions
				 sharedElementTransitions:
					 (NSArray<SharedElementTransitionOptions *> *)sharedElementTransitions
								 duration:(CGFloat)duration
								   bridge:(RCTBridge *)bridge {
	self = [super init];
	_bridge = bridge;
	_content = contentTransition;
	_elementTransitions = elementTransitions;
	_sharedElementTransitions = sharedElementTransitions;
	_duration = duration;
	return self;
}

#ifdef RCT_NEW_ARCH_ENABLED
- (instancetype)initWithContentTransition:(RNNEnterExitAnimation *)contentTransition
					   elementTransitions:(NSArray<ElementTransitionOptions *> *)elementTransitions
				 sharedElementTransitions:
					 (NSArray<SharedElementTransitionOptions *> *)sharedElementTransitions
								 duration:(CGFloat)duration
									 host:(RCTHost *)host {
	self = [super init];
	_host = host;
	_content = contentTransition;
	_elementTransitions = elementTransitions;
	_sharedElementTransitions = sharedElementTransitions;
	_duration = duration;
	return self;
}
#endif

- (void)animateTransition:(id<UIViewControllerContextTransitioning>)transitionContext {
#ifdef RCT_NEW_ARCH_ENABLED
	if (_host != nil) {
		[_host.surfacePresenter addObserver:self];
	} else {
		[_bridge.uiManager.observerCoordinator addObserver:self];
	}
#else
	[_bridge.uiManager.observerCoordinator addObserver:self];
#endif
	
	_animate = YES;
	_transitionContext = transitionContext;
	[self prepareTransitionContext:transitionContext];

	UIViewController *fromVC =
		[transitionContext viewControllerForKey:UITransitionContextFromViewControllerKey];
	if (![fromVC.navigationController.childViewControllers containsObject:fromVC]) {
		[self performAnimationOnce];
	}
}

- (void)prepareTransitionContext:(id<UIViewControllerContextTransitioning>)transitionContext {
    UIViewController *toViewController =
    [transitionContext viewControllerForKey:UITransitionContextToViewControllerKey];
    
    UIView *fromView = [transitionContext viewForKey:UITransitionContextFromViewKey];
    UIView *toView = [transitionContext viewForKey:UITransitionContextToViewKey];
    
    if (toView) {
        toViewController.view.alpha = 0;
        if (fromView) {
            [transitionContext.containerView addSubview:fromView];
        }
        [transitionContext.containerView addSubview:toViewController.view];
    }
    
    [toViewController prepareForTransition];
}

- (NSArray *)createTransitionsFromVC:(UIViewController *)fromVC
								toVC:(UIViewController *)toVC
					   containerView:(UIView *)containerView {
	NSArray *transitions =
		[RNNScreenTransitionsCreator createTransitionsFromVC:fromVC
														toVC:toVC
											   containerView:containerView
										   contentTransition:self.content
										  elementTransitions:self.elementTransitions];
	_sharedElementAnimator =
		[[SharedElementAnimator alloc] initWithTransitions:self.sharedElementTransitions
													fromVC:fromVC
													  toVC:toVC
											 containerView:containerView];
	return [transitions arrayByAddingObjectsFromArray:[_sharedElementAnimator create]];
}

- (void)performAnimationOnce {
	if (_animate) {
		_animate = NO;
		RCTExecuteOnMainQueue(^{
		  id<UIViewControllerContextTransitioning> transitionContext = self->_transitionContext;
		  UIViewController *fromVC =
			  [transitionContext viewControllerForKey:UITransitionContextFromViewControllerKey];
		  UIViewController *toVC =
			  [transitionContext viewControllerForKey:UITransitionContextToViewControllerKey];
		  NSArray *transitions = [self createTransitionsFromVC:fromVC
														  toVC:toVC
												 containerView:transitionContext.containerView];
		  [self animateTransitions:transitions andTransitioningContext:transitionContext];
		});
	}
}

- (void)animateTransitions:(NSArray<id<DisplayLinkAnimatorDelegate>> *)animators
	andTransitioningContext:(id<UIViewControllerContextTransitioning>)transitionContext {
    UIView *toView = [transitionContext viewForKey:UITransitionContextToViewKey];
    UIViewController *fromVC = [transitionContext viewControllerForKey:UITransitionContextFromViewControllerKey];
    
    BOOL isDismiss = (toView == nil);
    
    if (isDismiss) {
        [UIView animateWithDuration:[self transitionDuration:transitionContext]
                         animations:^{
            fromVC.view.alpha = 0;
        } completion:^(BOOL finished) {
            [transitionContext completeTransition:![transitionContext transitionWasCancelled]];
            fromVC.view.alpha = 1;
        }];
        return;
    }
    
    DisplayLinkAnimator *displayLinkAnimator = [[DisplayLinkAnimator alloc]
                                                initWithDisplayLinkAnimators:animators
                                                duration:[self transitionDuration:transitionContext]];
    
    [displayLinkAnimator setOnStart:^{
        if (!isDismiss) {
            UIViewController *toVC = [transitionContext viewControllerForKey:UITransitionContextToViewControllerKey];
            toVC.view.alpha = 1.f;
        }
    }];
    
    [displayLinkAnimator setCompletion:^{
        if (![transitionContext transitionWasCancelled]) {
            [transitionContext completeTransition:![transitionContext transitionWasCancelled]];
        }
    }];
    
    [displayLinkAnimator start];
}

- (NSTimeInterval)transitionDuration:(id<UIViewControllerContextTransitioning>)transitionContext {
	return _duration;
}

// TODO: TEST THIS
#ifdef RCT_NEW_ARCH_ENABLED
- (void)didMountComponentsWithRootTag:(NSInteger)rootTag {
	[self performAnimationOnce];
}
#endif

- (void)uiManagerDidPerformMounting:(RCTUIManager *)manager {
	[self performAnimationOnce];
}

- (id<UIViewControllerAnimatedTransitioning>)
	animationControllerForPresentedController:(UIViewController *)presented
						 presentingController:(UIViewController *)presenting
							 sourceController:(UIViewController *)source {
	return self;
}

- (id<UIViewControllerAnimatedTransitioning>)animationControllerForDismissedController:
	(UIViewController *)dismissed {
	return self;
}

- (void)animationEnded:(BOOL)transitionCompleted {
	UIView *toView = [_transitionContext viewForKey:UITransitionContextToViewKey];
	UIView *fromView = [_transitionContext viewForKey:UITransitionContextFromViewKey];
	[_sharedElementAnimator animationEnded];
    if (toView) {
        toView.layer.transform = CATransform3DIdentity;
        toView.alpha = 1.f;
    }
    if (fromView) {
        fromView.layer.transform = CATransform3DIdentity;
    }
	_transitionContext = nil;
	_sharedElementAnimator = nil;
}

@end