import Foundation
import SwiftUI
import UIKit

extension View {
    func screenshotView(_ closure: @escaping ScreenshotToolClosure) -> some View {
        let screenshotView = ScreenshotToolView(closure)
        return overlay(screenshotView.allowsHitTesting(false))
    }
}

typealias ScreenshotToolClosure = (ScreenshotTool) -> Void

struct ScreenshotToolView: UIViewRepresentable {
    let closure: ScreenshotToolClosure

    init(_ closure: @escaping ScreenshotToolClosure) {
        self.closure = closure
    }

    func makeUIView(context _: Context) -> ScreenshotTool {
        let view = ScreenshotTool(frame: CGRect.zero)
        return view
    }

    func updateUIView(_ uiView: ScreenshotTool, context _: Context) {
        DispatchQueue.main.async {
            closure(uiView)
        }
    }
}

class ScreenshotTool: UIView {
    func screenshot() -> UIImage? {
        guard let containerView = self.superview?.superview,
              let containerSuperview = containerView.superview else { return nil }
        let renderer = UIGraphicsImageRenderer(bounds: containerView.frame)
        let image = renderer.image { _ in
            containerSuperview.drawHierarchy(in: containerSuperview.bounds, afterScreenUpdates: true)
        }

        UIGraphicsBeginImageContextWithOptions(containerView.bounds.size, false, UIScreen.main.scale)
        UIBezierPath(roundedRect: containerView.bounds, cornerRadius: 18).addClip()
        image.draw(in: containerView.bounds)
        let roundedImage = UIGraphicsGetImageFromCurrentImageContext()
        UIGraphicsEndImageContext()
        return roundedImage
    }
}