Tusker/Packages/MatchedGeometryPresentation/Sources/MatchedGeometryPresentation/MatchedGeometryViewController.swift

//
//  MatchedGeometryViewController.swift
//  MatchGeom
//
//  Created by Shadowfacts on 4/24/23.
//

import SwiftUI
import Combine

private let mass: CGFloat = 1
private let presentStiffness: CGFloat = 300
private let presentDamping: CGFloat = 20
private let dismissStiffness: CGFloat = 200
private let dismissDamping: CGFloat = 20

public class MatchedGeometryState: ObservableObject {
    @Published var presentationID: AnyHashable?
    @Published var animating: Bool = false
    @Published public var mode: Mode = .presenting
    @Published var sources: [SourceKey: (() -> AnyView, CGRect)] = [:]
    @Published var currentFrames: [AnyHashable: CGRect] = [:]
    @Published var destinations: [AnyHashable: (AnyView, CGRect)] = [:]
    
    public enum Mode: Equatable {
        case presenting
        case idle
        case dismissing
    }
}

class MatchedGeometryViewController<Content: View>: UIViewController, UIViewControllerTransitioningDelegate {
    
    let presentationID: AnyHashable
    let content: Content
    let state: MatchedGeometryState
    let backgroundColor: UIColor
    var contentHost: UIHostingController<ContentContainerView>!
    var matchedHost: UIHostingController<MatchedContainerView>!
    
    init(presentationID: AnyHashable, content: Content, state: MatchedGeometryState, backgroundColor: UIColor) {
        self.presentationID = presentationID
        self.content = content
        self.state = state
        self.backgroundColor = backgroundColor
        
        super.init(nibName: nil, bundle: nil)
        
        modalPresentationStyle = .custom
        transitioningDelegate = self
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        contentHost = UIHostingController(rootView: ContentContainerView(content: content, state: state))
        contentHost.view.autoresizingMask = [.flexibleWidth, .flexibleHeight]
        contentHost.view.frame = view.bounds
        contentHost.view.backgroundColor = backgroundColor
        addChild(contentHost)
        view.addSubview(contentHost.view)
        contentHost.didMove(toParent: self)
    }
    
    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)
        
        state.presentationID = presentationID
    }
    
    var currentPresentationSources: [AnyHashable: (() -> AnyView, CGRect)] {
        Dictionary(uniqueKeysWithValues: state.sources.filter { $0.key.presentationID == presentationID }.map { ($0.key.matchedID, $0.value) })
    }
    
    func addMatchedHostingController() {
        let sources = currentPresentationSources.map { (id: $0.key, view: $0.value.0) }
        matchedHost = UIHostingController(rootView: MatchedContainerView(sources: sources, state: state))
        matchedHost.view.autoresizingMask = [.flexibleWidth, .flexibleHeight]
        matchedHost.view.frame = view.bounds
        matchedHost.view.backgroundColor = .clear
        matchedHost.view.layer.zPosition = 100
        addChild(matchedHost)
        view.addSubview(matchedHost.view)
        matchedHost.didMove(toParent: self)
    }
    
    struct ContentContainerView: View {
        let content: Content
        let state: MatchedGeometryState
        
        var body: some View {
            content
                .environmentObject(state)
        }
    }
    
    struct MatchedContainerView: View {
        let sources: [(id: AnyHashable, view: () -> AnyView)]
        @ObservedObject var state: MatchedGeometryState
        
        var body: some View {
            ZStack {
                ForEach(sources, id: \.id) { (id, view) in
                    matchedView(id: id, source: view)
                }
            }
        }
        
        @ViewBuilder
        func matchedView(id: AnyHashable, source: () -> AnyView) -> some View {
            if let frame = state.currentFrames[id],
               let dest = state.destinations[id]?.0 {
                ZStack {
                    source()
                    dest
                        .opacity(state.mode == .presenting ? (state.animating ? 1 : 0) : (state.animating ? 0 : 1))
                }
                .frame(width: frame.width, height: frame.height)
                .position(x: frame.midX, y: frame.midY)
                .ignoresSafeArea()
                .animation(.interpolatingSpring(mass: Double(mass), stiffness: Double(state.mode == .presenting ? presentStiffness : dismissStiffness), damping: Double(state.mode == .presenting ? presentDamping : dismissDamping), initialVelocity: 0), value: frame)
            }
        }
    }
    
    // MARK: UIViewControllerTransitioningDelegate
    func animationController(forPresented presented: UIViewController, presenting: UIViewController, source: UIViewController) -> UIViewControllerAnimatedTransitioning? {
        return MatchedGeometryPresentationAnimationController<Content>()
    }
    
    func animationController(forDismissed dismissed: UIViewController) -> UIViewControllerAnimatedTransitioning? {
        return MatchedGeometryDismissAnimationController<Content>()
    }
    
    func presentationController(forPresented presented: UIViewController, presenting: UIViewController?, source: UIViewController) -> UIPresentationController? {
        return MatchedGeometryPresentationController(presentedViewController: presented, presenting: presenting)
    }
    
}

class MatchedGeometryPresentationAnimationController<Content: View>: NSObject, UIViewControllerAnimatedTransitioning {
    func transitionDuration(using transitionContext: UIViewControllerContextTransitioning?) -> TimeInterval {
        return 0.8
    }

    func animateTransition(using transitionContext: UIViewControllerContextTransitioning) {
        let matchedGeomVC = transitionContext.viewController(forKey: .to) as! MatchedGeometryViewController<Content>
        let container = transitionContext.containerView
        
        // add the VC to the container, which kicks off layout out the content hosting controller
        matchedGeomVC.view.autoresizingMask = [.flexibleWidth, .flexibleHeight]
        matchedGeomVC.view.frame = container.bounds
        container.addSubview(matchedGeomVC.view)
        
        // layout out the content hosting controller and having enough destinations may take a while
        // so listen for when it's ready, rather than trying to guess at the timing
        let cancellable = matchedGeomVC.state.$destinations
            .filter { destinations in matchedGeomVC.currentPresentationSources.allSatisfy { source in destinations.keys.contains(source.key) } }
            .first()
            .sink { destinations in
                matchedGeomVC.addMatchedHostingController()
                
                // setup the initial state for the animation
                matchedGeomVC.matchedHost.view.isHidden = true
                matchedGeomVC.state.mode = .presenting
                matchedGeomVC.state.currentFrames = matchedGeomVC.currentPresentationSources.mapValues(\.1)
                
                // wait one runloop iteration for the matched hosting controller to be setup
                DispatchQueue.main.async {
                    matchedGeomVC.matchedHost.view.isHidden = false
                    matchedGeomVC.state.animating = true
                    // get the now-current destinations, in case they've changed since the sunk value was published
                    matchedGeomVC.state.currentFrames = matchedGeomVC.state.destinations.mapValues(\.1)
                }
            }
        
        matchedGeomVC.contentHost.view.layer.opacity = 0
        let spring = UISpringTimingParameters(mass: mass, stiffness: presentStiffness, damping: presentDamping, initialVelocity: .zero)
        let animator = UIViewPropertyAnimator(duration: self.transitionDuration(using: transitionContext), timingParameters: spring)
        animator.addAnimations {
            matchedGeomVC.contentHost.view.layer.opacity = 1
        }
        animator.addCompletion { _ in
            transitionContext.completeTransition(true)
            matchedGeomVC.state.animating = false
            matchedGeomVC.state.mode = .idle
            
            matchedGeomVC.matchedHost?.view.removeFromSuperview()
            matchedGeomVC.matchedHost?.removeFromParent()
            cancellable.cancel()
        }
        animator.startAnimation()
    }
}

class MatchedGeometryDismissAnimationController<Content: View>: NSObject, UIViewControllerAnimatedTransitioning {
    func transitionDuration(using transitionContext: UIViewControllerContextTransitioning?) -> TimeInterval {
        return 0.8
    }
    
    func animateTransition(using transitionContext: UIViewControllerContextTransitioning) {
        let matchedGeomVC = transitionContext.viewController(forKey: .from) as! MatchedGeometryViewController<Content>
        
        // recreate the matched host b/c using the current destinations doesn't seem to update the existing one
        matchedGeomVC.addMatchedHostingController()
        matchedGeomVC.matchedHost.view.isHidden = true
        matchedGeomVC.state.mode = .dismissing
        matchedGeomVC.state.currentFrames = matchedGeomVC.state.destinations.mapValues(\.1)
        
        DispatchQueue.main.async {
            matchedGeomVC.matchedHost.view.isHidden = false
            matchedGeomVC.state.animating = true
            matchedGeomVC.state.currentFrames = matchedGeomVC.currentPresentationSources.mapValues(\.1)
        }
        
        let spring = UISpringTimingParameters(mass: mass, stiffness: dismissStiffness, damping: dismissDamping, initialVelocity: .zero)
        let animator = UIViewPropertyAnimator(duration: transitionDuration(using: transitionContext), timingParameters: spring)
        animator.addAnimations {
            matchedGeomVC.contentHost.view.layer.opacity = 0
        }
        animator.addCompletion { _ in
            transitionContext.completeTransition(true)
            matchedGeomVC.state.animating = false
            matchedGeomVC.state.mode = .idle
        }
        animator.startAnimation()
    }
}

class MatchedGeometryPresentationController: UIPresentationController {
    override func dismissalTransitionWillBegin() {
        super.dismissalTransitionWillBegin()
        delegate?.presentationControllerWillDismiss?(self)
    }
}