import UIKit
final class ContactListViewControllerAnimatedTransitioning: NSObject, UIViewControllerAnimatedTransitioning {
    private weak var buttonView: UIView?

    init(buttonView: UIView) { self.buttonView = buttonView }

    func transitionDuration(using transitionContext: UIViewControllerContextTransitioning?) -> TimeInterval { 0.350 }

    func animateTransition(using transitionContext: UIViewControllerContextTransitioning) {
        guard
            let fromViewController = transitionContext.viewController(forKey: .from),
            let toViewController = transitionContext.viewController(forKey: .to) as? UINavigationController,
            let toChildViewController = toViewController.viewControllers.first,
            // iOS 27 returns nil: UIKit has just reparented this view into UITransitionView and
            // that move is still in an uncommitted transaction. iOS 26 still returns a snapshot.
            let fromViewSnapshot = fromViewController.view.snapshotView(afterScreenUpdates: false),
            let buttonView = buttonView
        else {
            // toViewController.view is never added to the container, and UIKit removes the
            // presenting view from the window once the transition completes -> black screen.
            transitionContext.completeTransition(!transitionContext.transitionWasCancelled)
            return
        }
        transitionContext.containerView.addSubview(fromViewSnapshot)
        let splitedSnapshotViews = splitSnapshot(of: toViewController)
        let originalYs = splitedSnapshotViews.map { $0.frame.origin.y }
        let rowSnapshotViews = Array(splitedSnapshotViews.dropFirst(2))
        guard
            let navigationBarSnapshotView = splitedSnapshotViews[safe: 0],
            let inputSnapshotView = splitedSnapshotViews[safe: 1],
            !rowSnapshotViews.isEmpty
        else {
            transitionContext.completeTransition(!transitionContext.transitionWasCancelled)
            return
        }
        let coveringButtonView = UIView()
        coveringButtonView.backgroundColor = .white
        coveringButtonView.frame = buttonView.frameInWindow
        transitionContext.containerView.addSubview(coveringButtonView)
        let overlayView = UIView()
        overlayView.backgroundColor = .white
        overlayView.frame = transitionContext.finalFrame(for: toViewController)
        overlayView.alpha = 0
        transitionContext.containerView.addSubview(overlayView)
        let buttonSnapshotView = UIImageView(image: buttonView.snapshot)
        buttonSnapshotView.frame = buttonView.frameInWindow
        transitionContext.containerView.addSubview(buttonSnapshotView)
        transitionContext.containerView.addSubview(inputSnapshotView)
        inputSnapshotView.frame.origin.y = buttonView.frameInWindow.origin.y
        rowSnapshotViews.enumerated().forEach {
            $0.element.frame.origin.y += (buttonView.frameInWindow.maxY - originalYs[1])
                * CGFloat(0.25) + CGFloat(24 * $0.offset)
            transitionContext.containerView.addSubview($0.element)
        }
        var coveringNavigationBarSnapshotView: UIView?
        if navigationBarSnapshotView.frame.maxY >= buttonView.frameInWindow.minY,
           let cgImage = fromViewController.view.snapshot?.cgImage {
            coveringNavigationBarSnapshotView = crop(cgImage: cgImage, to: navigationBarSnapshotView.frame)
        } else {
            navigationBarSnapshotView.frame.origin.y = (buttonView.frameInWindow.minY - navigationBarSnapshotView.frame.minY) / 2
        }
        coveringNavigationBarSnapshotView.map(transitionContext.containerView.addSubview)
        let coveringViews = [coveringNavigationBarSnapshotView, buttonSnapshotView]
        transitionContext.containerView.addSubview(navigationBarSnapshotView)
        splitedSnapshotViews.forEach { $0.alpha = 0 }
        transitionContext.containerView.addSubview(toViewController.view)
        toViewController.view.isHidden = true
        UIView.animate(withDuration: 0.125, delay: 0, options: .curveEaseOut, animations: { overlayView.alpha = 1 })
        UIView.animate(withDuration: 0.225, delay: 0, options: .curveEaseOut, animations: {
            splitedSnapshotViews.enumerated().prefix(2).forEach {
                $0.element.alpha = 1
                $0.element.frame.origin.y = originalYs[$0.offset]
            }
            coveringViews.forEach { $0?.alpha = 0 }
            buttonSnapshotView.frame.origin.y = originalYs[1]
        }, completion: { _ in _ = toChildViewController.becomeFirstResponder() })
        DispatchQueue.main.asyncAfter(deadline: .now() + 0.125) {
            UIView.animate(withDuration: 0.175, delay: 0, options: .curveEaseOut, animations: {
                rowSnapshotViews.forEach { $0.alpha = 1 }
            })
            UIView.animate(withDuration: 0.225, delay: 0, options: .curveEaseOut, animations: {
                splitedSnapshotViews.enumerated().dropFirst(2).forEach {
                    $0.element.frame.origin.y = originalYs[$0.offset]
                }
            }, completion: { _ in
                toViewController.view.isHidden = false
                transitionContext.completeTransition(!transitionContext.transitionWasCancelled)
            })
        }
    }

    private func splitSnapshot(of viewController: UIViewController) -> [UIView] {
        guard let navigationController = viewController as? UINavigationController,
              let cgImage = viewController.view.snapshot?.cgImage else { return [] }
        var imageViews = [UIImageView]()
        let navigationBarFrame = navigationController.navigationBar.frame
        crop(cgImage: cgImage, to: navigationBarFrame).map { imageViews.append($0) }
        let inputViewFrame = CGRect(x: 0, y: navigationBarFrame.maxY, width: navigationBarFrame.width, height: 72)
        crop(cgImage: cgImage, to: inputViewFrame).map { imageViews.append($0) }
        var lastY = inputViewFrame.maxY
        while lastY < viewController.view.frame.height {
            let rowFrame = CGRect(x: 0, y: lastY, width: navigationBarFrame.width, height: min(60, viewController.view.frame.height - lastY))
            crop(cgImage: cgImage, to: rowFrame).map { imageViews.append($0) }
            lastY += 60
        }
        return imageViews
    }

    private func crop(cgImage: CGImage, to rect: CGRect) -> UIImageView? {
        let scale = UIScreen.main.scale
        let scaledRect = CGRect(x: rect.origin.x * scale, y: rect.origin.y * scale, width: rect.width * scale, height: rect.height * scale)
        guard let croppedImage = cgImage.cropping(to: scaledRect) else { return nil }
        let imageView = UIImageView(image: UIImage(cgImage: croppedImage))
        imageView.frame = rect
        return imageView
    }
}

final class ViewController: UIViewController {
    private let openButton = UIButton(type: .system)

    override func viewDidLoad() {
        super.viewDidLoad()
        view.backgroundColor = .white
        navigationItem.title = "Snapshot sample"
        openButton.setTitle("Open contacts", for: .normal)
        openButton.setTitleColor(.white, for: .normal)
        openButton.backgroundColor = .systemBlue
        openButton.addTarget(self, action: #selector(openTapped), for: .touchUpInside)
        view.addSubview(openButton)
    }

    override func viewDidLayoutSubviews() {
        super.viewDidLayoutSubviews()
        openButton.frame = CGRect(x: (view.bounds.width - 220) / 2, y: view.bounds.midY - 24, width: 220, height: 48)
    }

    @objc private func openTapped() {
        let navigationController = UINavigationController(rootViewController: ContactListViewController())
        // .fullScreen means shouldRemovePresentersView == true, so UIKit reparents the presenting
        // view into UITransitionView. With .custom nothing is reparented and the bug disappears.
        navigationController.modalPresentationStyle = .fullScreen
        navigationController.transitioningDelegate = self
        present(navigationController, animated: true)
    }
}

extension ViewController: UIViewControllerTransitioningDelegate {
    func animationController(forPresented presented: UIViewController, presenting: UIViewController,
                             source: UIViewController) -> UIViewControllerAnimatedTransitioning? {
        ContactListViewControllerAnimatedTransitioning(buttonView: openButton)
    }
}

final class ContactListViewController: UIViewController, UITableViewDataSource {
    private let searchField = UITextField()
    private let closeButton = UIButton(type: .system)
    private let tableView = UITableView()
    private let contacts = (1...12).map { "Contact #\($0)" }

    override func viewDidLoad() {
        super.viewDidLoad()
        view.backgroundColor = .white
        navigationItem.title = "Contacts"
        searchField.backgroundColor = .systemGray6
        closeButton.setTitle("Close", for: .normal)
        closeButton.addTarget(self, action: #selector(closeTapped), for: .touchUpInside)
        tableView.dataSource = self
        tableView.register(UITableViewCell.self, forCellReuseIdentifier: "cell")
        tableView.rowHeight = 60
        [closeButton, searchField, tableView].forEach(view.addSubview)
    }

    override func viewDidLayoutSubviews() {
        super.viewDidLayoutSubviews()
        let top = view.safeAreaInsets.top + 12
        closeButton.frame = CGRect(x: view.bounds.width - 96, y: top, width: 80, height: 30)
        searchField.frame = CGRect(x: 16, y: closeButton.frame.maxY + 8, width: view.bounds.width - 32, height: 72)
        tableView.frame = CGRect(x: 0, y: searchField.frame.maxY + 8,
                                 width: view.bounds.width, height: view.bounds.height - searchField.frame.maxY - 8)
    }

    // Escape hatch: when the bug hits, this view never reaches the window and Close is unreachable.
    override func viewDidAppear(_ animated: Bool) {
        super.viewDidAppear(animated)
        guard view.window == nil else { return }
        DispatchQueue.main.asyncAfter(deadline: .now() + 3) { [weak self] in self?.dismiss(animated: false) }
    }

    @objc private func closeTapped() { dismiss(animated: true) }

    func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int { contacts.count }

    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        let cell = tableView.dequeueReusableCell(withIdentifier: "cell", for: indexPath)
        cell.textLabel?.text = contacts[indexPath.row]
        return cell
    }
}

private extension UIView {
    var snapshot: UIImage? {
        UIGraphicsImageRenderer(bounds: bounds).image { layer.render(in: $0.cgContext) }
    }

    var frameInWindow: CGRect { convert(bounds, to: nil) }
}

private extension Array {
    subscript(safe index: Int) -> Element? { indices.contains(index) ? self[index] : nil }
}