import AppKit
import ScreenCaptureKit

/// Captures window thumbnails using ScreenCaptureKit.
enum ThumbnailCapture {
    /// Max thumbnail dimension — large enough for sharp display at overlay sizes,
    /// small enough for fast capture. 800px wide covers most overlay thumbnail widths at 2x.
    private static let maxDimension: Int = 800

    /// Pre-warm ScreenCaptureKit so the first real capture is fast.
    static func warmUp() async {
        _ = try? await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: true)
    }

    /// Capture thumbnails for all windows in parallel.
    static func captureThumbnails(for windows: inout [WindowInfo]) async {
        let t0 = CFAbsoluteTimeGetCurrent()
        guard let content = try? await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: true) else {
            return
        }
        let t1 = CFAbsoluteTimeGetCurrent()

        // Build a lookup from windowID to SCWindow
        var scWindowsByID: [CGWindowID: SCWindow] = [:]
        for scWindow in content.windows {
            scWindowsByID[scWindow.windowID] = scWindow
        }

        // Capture all windows in parallel
        let captures: [(Int, CGImage)] = await withTaskGroup(of: (Int, CGImage?).self) { group in
            for i in windows.indices {
                guard let scWindow = scWindowsByID[windows[i].windowID] else { continue }
                let windowFrame = windows[i].frame
                let idx = i

                group.addTask {
                    let filter = SCContentFilter(desktopIndependentWindow: scWindow)
                    let config = SCStreamConfiguration()

                    // Scale down to maxDimension while preserving aspect ratio
                    let aspect = windowFrame.width / max(windowFrame.height, 1)
                    if aspect >= 1 {
                        config.width = maxDimension
                        config.height = max(1, Int(Double(maxDimension) / aspect))
                    } else {
                        config.height = maxDimension
                        config.width = max(1, Int(Double(maxDimension) * aspect))
                    }
                    config.scalesToFit = true
                    config.showsCursor = false

                    let image = try? await SCScreenshotManager.captureImage(
                        contentFilter: filter,
                        configuration: config
                    )
                    return (idx, image)
                }
            }

            var results: [(Int, CGImage)] = []
            for await (idx, image) in group {
                if let image = image {
                    results.append((idx, image))
                }
            }
            return results
        }

        for (idx, image) in captures {
            windows[idx].thumbnail = image
        }

        let t2 = CFAbsoluteTimeGetCurrent()
        mctLog("[MCT] capture: shareable=%.0fms parallel=%.0fms total=%.0fms (%d windows)",
               (t1 - t0) * 1000, (t2 - t1) * 1000, (t2 - t0) * 1000, captures.count)
    }
}