import UIKit
import AVFoundation

// MARK: - CaptureIntermediate

/// 截图中间结果
///
/// 主线程只做视图访问和像素 Buffer 提取（轻量）；
/// CIContext 重度渲染（AVPlayerLayer 视频帧路径）由调用方搬到 `TaskStore.ioQueue` 执行，
/// 对齐 Android `BitmapFrameEngine` 将 Bitmap 分配与渲染拆离主线程的设计。
enum CaptureIntermediate {
    /// 已完成渲染的 UIImage（`drawHierarchy` 路径）
    case image(UIImage)
    /// 从 AVPlayerLayer 提取的原始视频帧 Buffer（需在 ioQueue 上调用 `toImage()` 完成 CIContext 渲染）
    case pixelBuffer(CVPixelBuffer)

    /// 将 CaptureIntermediate 转换为 UIImage
    ///
    /// - `.image` 直接返回；`.pixelBuffer` 通过 `CIContext` 渲染（**可在任意线程调用**）
    func toImage() -> UIImage? {
        switch self {
        case .image(let img):
            return img
        case .pixelBuffer(let buf):
            let ciImage = CIImage(cvPixelBuffer: buf)
            let context = CIContext()
            guard let cgImage = context.createCGImage(ciImage, from: ciImage.extent) else { return nil }
            return UIImage(cgImage: cgImage)
        }
    }
}

// MARK: - ImageProcessor

/// 截图、缩放、压缩、落盘
///
/// ## 线程模型（对齐 Android BitmapFrameEngine）
///
/// | 阶段 | 线程 | 原因 |
/// |---|---|---|
/// | `captureViewIntermediate` | **主线程** | UIKit / CALayer 访问，drawHierarchy 强制主线程 |
/// | `CaptureIntermediate.toImage()` | ioQueue（子线程） | CIContext 渲染为纯 CPU/GPU 计算，无 UIKit 依赖 |
/// | `compressAndScale` / pHash / 落盘 | ioQueue（子线程） | IO + 计算密集，避免阻塞 UI |
///
/// `drawHierarchy` 路径返回 `.image`；
/// AVPlayerLayer（视频广告）路径返回 `.pixelBuffer`，CIContext 渲染从主线程移至 ioQueue，
/// 减轻主线程压力。
enum ImageProcessor {

    // MARK: - Screenshot（主线程）

    /// 截取指定 View 的画面，返回中间结果
    ///
    /// 按视图类型分流：
    /// - `AVPlayerLayer` → 提取 `CVPixelBuffer` → `.pixelBuffer`（CIContext 在调用方 ioQueue 完成）
    /// - 普通 `UIView`   → `drawHierarchy` → `.image`
    ///
    /// - Parameters:
    ///   - view: 目标视图
    ///   - excludeViews: 需要遮盖的子视图（隐私区域）
    /// - Returns: `CaptureIntermediate`；失败返回 `nil`
    /// - Note: **必须在主线程调用**（UIKit 强约束）
    static func captureViewIntermediate(_ view: UIView,
                                        excludeViews: [UIView]) -> CaptureIntermediate? {
        // 隐私遮盖：临时隐藏敏感视图
        let hiddenViews = collectSensitiveViews(in: view, excludeViews: excludeViews)
        hiddenViews.forEach { $0.isHidden = true }
        defer { hiddenViews.forEach { $0.isHidden = false } }

        // AVPlayerLayer：只提取 CVPixelBuffer（轻量），CIContext 渲染留给 ioQueue
        if let playerLayer = findPlayerLayer(in: view.layer), let laver_image = playerLayer.takeCurrentFrame() {
            if let window_image = view.snapshotWithVideo(player: playerLayer.player, videoRect: playerLayer.videoRect) {
                print(playerLayer)
                return .image(window_image)
            }
        }
        
        // 通用 drawHierarchy（UIKit 强制主线程）
        let renderer = UIGraphicsImageRenderer(bounds: view.bounds)
        let img = renderer.image { _ in
            view.drawHierarchy(in: view.bounds, afterScreenUpdates: true)
        }
        return .image(img)
    }

    /// 兼容旧调用路径：截取 View 并同步返回 UIImage
    ///
    /// 内部调用 `captureViewIntermediate` + `toImage()`；
    /// AVPlayerLayer 路径的 CIContext 仍在当前线程执行（需主线程），建议直接使用
    /// `captureViewIntermediate` + ioQueue `toImage()` 以获得更好性能。
    ///
    /// - Note: 必须在主线程调用
    static func captureView(_ view: UIView, excludeViews: [UIView]) -> UIImage? {
        captureViewIntermediate(view, excludeViews: excludeViews)?.toImage()
    }

    // MARK: - Compress & Scale（IO 串行队列）

    /// 缩放 + JPEG 压缩
    ///
    /// - Parameters:
    ///   - image: 源图像
    ///   - options: 含 `maxLongEdge`（落盘 JPEG 最长边像素上限）和 `maxSizeKb`（压缩目标）
    /// - Returns: 满足大小约束的 JPEG Data；失败返回 `nil`
    static func compressAndScale(_ image: UIImage, options: FrameOptions) -> Data? {
        let scaled = scale(image, maxLongEdge: options.maxLongEdge)
        return compress(scaled, maxSizeKb: options.maxSizeKb)
    }

    // MARK: - Private helpers

    /// 位图像素尺寸（`size × scale`）；落盘 JPEG 以像素为准，不能只用 `UIImage.size`（pt）
    private static func pixelSize(of image: UIImage) -> CGSize {
        CGSize(
            width: image.size.width * image.scale,
            height: image.size.height * image.scale
        )
    }

    /// `scale = 1`，使 `UIImage.size` 与 JPEG 像素尺寸一致
    private static func pixelAlignedRendererFormat() -> UIGraphicsImageRendererFormat {
        let format = UIGraphicsImageRendererFormat()
        format.scale = 1.0
        return format
    }

    private static func renderAtPixelSize(_ image: UIImage, pixelSize: CGSize) -> UIImage {
        let format = pixelAlignedRendererFormat()
        let renderer = UIGraphicsImageRenderer(size: pixelSize, format: format)
        return renderer.image { _ in
            image.draw(in: CGRect(origin: .zero, size: pixelSize))
        }
    }

    /// 按**像素**最长边限制等比缩小（`maxLongEdge` 对应落盘 JPEG 的像素边长）
    private static func scale(_ image: UIImage, maxLongEdge: Int) -> UIImage {
        let pixels = pixelSize(of: image)
        let longest = max(pixels.width, pixels.height)
        let limit = CGFloat(maxLongEdge)
        guard longest > limit else { return image }

        let ratio = limit / longest
        let newPixelSize = CGSize(
            width: floor(pixels.width * ratio),
            height: floor(pixels.height * ratio)
        )
        return renderAtPixelSize(image, pixelSize: newPixelSize)
    }

    private static func compress(_ image: UIImage, maxSizeKb: Int) -> Data? {
        let maxBytes = maxSizeKb * 1024
        // 从质量 0.7 开始，每次降 0.1，最低 0.3
        var quality: CGFloat = 0.7
        while quality >= 0.3 {
            if let data = image.jpegData(compressionQuality: quality),
               data.count <= maxBytes {
                return data
            }
            quality -= 0.1
        }
        // 二次缩放后再压缩一次（像素尺寸缩至 75%）
        let pixels = pixelSize(of: image)
        let smallerPixelSize = CGSize(
            width: floor(pixels.width * 0.75),
            height: floor(pixels.height * 0.75)
        )
        let smaller = renderAtPixelSize(image, pixelSize: smallerPixelSize)
        return smaller.jpegData(compressionQuality: 0.5)
    }

    // MARK: AVPlayerLayer（主线程：仅提取 PixelBuffer；CIContext 渲染在 ioQueue 完成）

    /// 从 AVPlayerLayer 提取当前帧的 CVPixelBuffer（主线程调用）
    ///
    /// CIContext.createCGImage 是 CPU/GPU 密集操作，不在此处执行；
    /// 调用方应在 `TaskStore.ioQueue` 调用 `CaptureIntermediate.toImage()` 完成转换。
    ///
    /// ## 返回 nil 的常见原因
    /// 1. output 刚 add，decoder 还没解出第一帧（最常见 — 现已通过短轮询规避）
    /// 2. currentItem.status != .readyToPlay（item 仍在加载或加载失败）
    /// 3. 视频含 DRM 保护（FairPlay 等加密内容禁止抽帧）
    /// 4. 播放器 stall / buffer empty（网络问题）
    /// 5. 视频未真正用 AVPlayerLayer 渲染（广告 SDK 多用 AVSampleBufferDisplayLayer / Metal）
    ///
    /// ## 主线程阻塞预算
    /// 最长 ~100ms（6 次轮询 × 16ms vsync），位于 delayMs 之后的截图阶段，整体可接受。
    private static func extractPlayerPixelBuffer(in view: UIView) -> CVPixelBuffer? {
        guard let playerLayer = findPlayerLayer(in: view.layer),
              let player = playerLayer.player,
              let currentItem = player.currentItem else { return nil }
        
        
//        if let frameImg = playerLayer.takeCurrentFrame() {
//            // 使用图片：展示/保存/合成截图
//            print(frameImg)
//        }
        
        // 状态预检：item 未 ready 时再轮询也是浪费主线程
        guard currentItem.status == .readyToPlay else {
            mmLog("⚠️ AVPlayer 抽帧跳过：item.status=\(currentItem.status.rawValue) (非 readyToPlay)")
            return nil
        }

        // 显式指定输出格式 + IOSurface，避免 nil attributes 导致 copyPixelBuffer 不一致
        let attrs: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String:    kCVPixelFormatType_32BGRA,
            kCVPixelBufferIOSurfacePropertiesKey as String: [:]
        ]
        let output = AVPlayerItemVideoOutput(pixelBufferAttributes: attrs)
        currentItem.add(output)
        defer { currentItem.remove(output) }

        // 短轮询：decoder 需要 1~3 帧时间把首帧推进 output 缓冲区
        // - 用 output.itemTime(forHostTime:) 把宿主时间换算为 item 时间，比 player.currentTime() 更准确
        // - 每次轮询失败后睡眠一个 vsync (~16ms)，最多 6 轮 (~100ms)
        let maxAttempts = 6
        for attempt in 1...maxAttempts {
            let hostTime = CACurrentMediaTime()
            let itemTime = output.itemTime(forHostTime: hostTime)

            if output.hasNewPixelBuffer(forItemTime: itemTime),
               let pb = output.copyPixelBuffer(forItemTime: itemTime, itemTimeForDisplay: nil) {
                mmLog("✅ AVPlayer 抽帧成功 attempts=\(attempt)/\(maxAttempts)")
                return pb
            }
            if attempt < maxAttempts {
                Thread.sleep(forTimeInterval: 0.016)
            }
        }

        // 兜底：播放器暂停时 hostTime 路径会 stall，改用 player.currentTime() 再试一次
        let fallbackTime = player.currentTime()
        if output.hasNewPixelBuffer(forItemTime: fallbackTime),
           let pb = output.copyPixelBuffer(forItemTime: fallbackTime, itemTimeForDisplay: nil) {
            mmLog("✅ AVPlayer 抽帧成功（fallback currentTime）")
            return pb
        }

        // 失败诊断：打印关键状态便于排查
        let rate     = player.rate
        let isLikely = currentItem.isPlaybackLikelyToKeepUp
        let empty    = currentItem.isPlaybackBufferEmpty
        mmLog("""
        ❌ AVPlayer 抽帧失败 attempts=\(maxAttempts)
             player.rate                    = \(rate)
             item.status                    = \(currentItem.status.rawValue)
             item.isPlaybackLikelyToKeepUp  = \(isLikely)
             item.isPlaybackBufferEmpty     = \(empty)
             提示：若视频实际由 AVSampleBufferDisplayLayer / CAMetalLayer 渲染，AVPlayer 路径无法命中
        """)
        return nil
    }

    private static func findPlayerLayer(in layer: CALayer) -> AVPlayerLayer? {
        if let pl = layer as? AVPlayerLayer { return pl }
        for sub in layer.sublayers ?? [] {
            if let found = findPlayerLayer(in: sub) { return found }
        }
        return nil
    }

    // MARK: Privacy Filtering

    private static func collectSensitiveViews(in root: UIView, excludeViews: [UIView]) -> [UIView] {
        var result: [UIView] = []
        let firstResponder = UIApplication.shared.mm_firstResponder

        func traverse(_ view: UIView) {
            let isSensitive: Bool
            if let tf = view as? UITextField {
                isSensitive = tf.isSecureTextEntry || view === firstResponder
            } else if view is UITextView, view === firstResponder {
                isSensitive = true
            } else {
                isSensitive = excludeViews.contains { $0 === view }
            }
            if isSensitive { result.append(view) }
            view.subviews.forEach { traverse($0) }
        }
        traverse(root)
        return result
    }
}

// MARK: - UIApplication + helpers

extension UIApplication {

    /// 兼容 iOS 12 / 13+ 的 keyWindow 获取
    var mm_keyWindow: UIWindow? {
        if #available(iOS 13.0, *) {
            for scene in connectedScenes {
                guard let windowScene = scene as? UIWindowScene else { continue }
                if windowScene.activationState == .foregroundActive ||
                    windowScene.activationState == .foregroundInactive {
                    if let window = windowScene.windows.first(where: { $0.isKeyWindow }) {
                        return window
                    }
                    if let window = windowScene.windows.first {
                        return window
                    }
                }
            }
        }
        return keyWindow
    }

    /// 当前 firstResponder（用于隐私视图过滤）
    var mm_firstResponder: UIResponder? {
        UIResponder.mm_current
    }
}

// MARK: - UIResponder + firstResponder

private extension UIResponder {
    static weak var mm_current: UIResponder?

    @objc private func mm_findFirstResponder(_ sender: Any) {
        UIResponder.mm_current = self
    }

    static var mm_currentFirstResponder: UIResponder? {
        mm_current = nil
        UIApplication.shared.sendAction(
            #selector(mm_findFirstResponder(_:)), to: nil, from: nil, for: nil
        )
        return mm_current
    }
}

extension AVPlayerLayer {
    /// 截取当前播放帧（同步）
    func takeCurrentFrame() -> UIImage? {
        guard let player = self.player,
              let currentItem = player.currentItem else { return nil }
        
        let generator = AVAssetImageGenerator(asset: currentItem.asset)
        // 适配视频旋转、镜像
        generator.appliesPreferredTrackTransform = true
        // 精准取帧，消除时间偏差
        generator.requestedTimeToleranceBefore = .zero
        generator.requestedTimeToleranceAfter = .zero
        
        let currentTime = player.currentTime()
        do {
            let cgImage = try generator.copyCGImage(at: currentTime, actualTime: nil)
            return UIImage(cgImage: cgImage)
        } catch {
            print("AVPlayerLayer 截帧失败: \(error)")
            return nil
        }
    }
    
    /// 异步截帧（子线程执行，不卡UI）
    func asyncTakeCurrentFrame(completion: @escaping (UIImage?) -> Void) {
        guard let player = self.player,
              let currentItem = player.currentItem else {
            completion(nil)
            return
        }
        let currentTime = player.currentTime()
        
        DispatchQueue.global(qos: .utility).async {
            let generator = AVAssetImageGenerator(asset: currentItem.asset)
            generator.appliesPreferredTrackTransform = true
            generator.requestedTimeToleranceBefore = .zero
            generator.requestedTimeToleranceAfter = .zero
            
            var resultImg: UIImage?
            do {
                let cgImage = try generator.copyCGImage(at: currentTime, actualTime: nil)
                resultImg = UIImage(cgImage: cgImage)
            } catch {
                print("异步截帧失败: \(error)")
            }
            
            DispatchQueue.main.async {
                completion(resultImg)
            }
        }
    }
}


extension UIView {
    /// 普通截图（仅UI，视频区域黑屏，系统默认效果）
    func normalSnapshot() -> UIImage? {
        UIGraphicsBeginImageContextWithOptions(bounds.size, false, UIScreen.main.scale)
        defer { UIGraphicsEndImageContext() }
        drawHierarchy(in: bounds, afterScreenUpdates: true)
        return UIGraphicsGetImageFromCurrentImageContext()
    }
    
    /// 截取窗口 + 合成视频画面（解决视频黑屏问题）
    /// - Parameters:
    ///   - player: 视频播放器 AVPlayer
    ///   - videoRect: 视频视图在当前 Window 中的位置尺寸
    /// - Returns: 合成后的完整图片
    func snapshotWithVideo(player: AVPlayer?, videoRect: CGRect) -> UIImage? {
        // 1. 截取窗口UI层
        guard let uiImage = normalSnapshot() else { return nil }
        guard let player = player else { return nil }
        
        // 2. 提取视频当前帧
        guard let asset = player.currentItem?.asset else { return uiImage }
        let generator = AVAssetImageGenerator(asset: asset)
        generator.appliesPreferredTrackTransform = true
        generator.maximumSize = .zero // 不限制尺寸
        
        let currentTime = player.currentTime()
        guard let cgImage = try? generator.copyCGImage(at: currentTime, actualTime: nil) else {
            return uiImage
        }
        let videoImage = UIImage(cgImage: cgImage)
        
        // 3. 合成 UI + 视频帧
        UIGraphicsBeginImageContextWithOptions(uiImage.size, false, UIScreen.main.scale)
        defer { UIGraphicsEndImageContext() }
        
        uiImage.draw(at: .zero)
        videoImage.draw(in: videoRect)
        
        return UIGraphicsGetImageFromCurrentImageContext()
    }
}


