shotscreen/ShotScreen/Sources/ScreenCaptureKitProvider.swift

import ScreenCaptureKit
import AppKit

@available(macOS 12.3, *)
class ScreenCaptureKitProvider {

    func getDesktopIconWindows() async -> [SCWindow] {
        // Check if desktop icon hiding is enabled in settings
        guard SettingsManager.shared.hideDesktopIconsDuringScreenshot else {
            // If setting is disabled, return empty array so no windows are excluded
            return []
        }
        
        do {
            let content = try await SCShareableContent.current
            let windows = content.windows

            let desktopIconWindows = windows.filter { window in
                guard let app = window.owningApplication, app.bundleIdentifier == "com.apple.finder" else {
                    return false
                }
                guard window.title == nil || window.title == "" else {
                    return false
                }
                // Layer for desktop icons on Sonoma/Sequoia seems to be -2147483603
                // This value might be kCGDesktopWindowLevel - 20 (or similar)
                // kCGDesktopWindowLevel is (CGWindowLevel) (kCGBaseWindowLevel + kCGDesktopWindowLevelKey)
                // kCGBaseWindowLevel is Int32.min
                // kCGDesktopWindowLevelKey is 20
                // So kCGDesktopWindowLevel is Int32.min + 20 = -2147483648 + 20 = -2147483628
                // The observed value -2147483603 might be kCGDesktopWindowLevel + 25 or similar constant.
                // Let's stick to the observed value for now as it seems consistent.
                return window.windowLayer == -2147483603 && window.isOnScreen && window.frame.size.width > 0 && window.frame.size.height > 0
            }
            return desktopIconWindows
        } catch {
            NSLog("Error fetching shareable content for desktop icons: \(error.localizedDescription)")
            return []
        }
    }

    func getDesktopWidgetWindows() async -> [SCWindow] {
        // Check if desktop widget hiding is enabled in settings
        guard SettingsManager.shared.hideDesktopWidgetsDuringScreenshot else {
            // If setting is disabled, return empty array so no windows are excluded
            return []
        }
        
        do {
            let content = try await SCShareableContent.current
            let windows = content.windows
            
            // Use DesktopIconManager to detect widgets
            let detectedWidgets = DesktopIconManager.shared.detectDesktopWidgets(from: windows)
            NSLog("🔍 ScreenCaptureKitProvider: Found \(detectedWidgets.count) desktop widgets to hide")
            return detectedWidgets
        } catch {
            NSLog("Error fetching shareable content for desktop widgets: \(error.localizedDescription)")
            return []
        }
    }

    func getAllWindowsToExclude() async -> [SCWindow] {
        // Combine both desktop icons and widgets into one exclusion list
        async let iconWindows = getDesktopIconWindows()
        async let widgetWindows = getDesktopWidgetWindows()
        
        let allWindows = await iconWindows + widgetWindows
        NSLog("🔍 ScreenCaptureKitProvider: Total windows to exclude: \(allWindows.count) (icons + widgets)")
        return allWindows
    }

    func captureScreen(screen: NSScreen, excludingWindows: [SCWindow]? = nil) async -> NSImage? {
        do {
            let content = try await SCShareableContent.current
            guard let display = content.displays.first(where: { $0.displayID == screen.displayID }) else {
                NSLog("Error: Could not find SCDisplay matching NSScreen with ID: \(screen.displayID)")
                return nil
            }

            var windowsToExclude = excludingWindows ?? []
            if let ownBundleID = Bundle.main.bundleIdentifier {
                let ownWindows = content.windows.filter { $0.owningApplication?.bundleIdentifier == ownBundleID && $0.isOnScreen }
                windowsToExclude.append(contentsOf: ownWindows)
            }

            let filter = SCContentFilter(display: display, excludingWindows: windowsToExclude)
            let configuration = SCStreamConfiguration()

            configuration.width = display.width
            configuration.height = display.height
            configuration.showsCursor = SettingsManager.shared.windowCaptureIncludeCursor
            configuration.capturesAudio = false
            configuration.pixelFormat = kCVPixelFormatType_32BGRA
            configuration.colorSpaceName = CGColorSpace.sRGB
            
            let stream = SCStream(filter: filter, configuration: configuration, delegate: nil)
            try await stream.addStreamOutput(SingleFrameOutput.shared, type: .screen, sampleHandlerQueue: .main)
            try await stream.startCapture()
            // Wacht kort op een frame - dit is een placeholder, een betere synchronisatie is nodig
            try await Task.sleep(nanoseconds: 100_000_000) // 0.1 seconde
            let image = await SingleFrameOutput.shared.retrieveFrame()
            try await stream.stopCapture()
            return image
        } catch {
            NSLog("Error capturing screen with SCStream: \(error.localizedDescription)")
            return nil
        }
    }

    func captureSelection(selectionRectInPoints: CGRect, screen: NSScreen, excludingWindows: [SCWindow]? = nil) async -> NSImage? {
        do {
            let content = try await SCShareableContent.current
            guard let display = content.displays.first(where: { $0.displayID == screen.displayID }) else {
                NSLog("Error: Could not find SCDisplay matching NSScreen with ID: \(screen.displayID) for selection capture")
                return nil
            }

            let scale = screen.backingScaleFactor
            
            // selectionRectInPoints is de intersectie van de globale selectie met screen.frame,
            // nog steeds in globale AppKit coördinaten (Y-omhoog, oorsprong linksonder hoofdvenster).

            // Stap 1: Converteer globale AppKit selectie naar lokale AppKit punten voor dit scherm.
            // (Y-omhoog, oorsprong linksonder DIT scherm).
            let localOriginXInPoints = selectionRectInPoints.origin.x - screen.frame.origin.x
            let localOriginYInPoints = selectionRectInPoints.origin.y - screen.frame.origin.y
            // Breedte en hoogte (in punten) blijven hetzelfde als selectionRectInPoints.size.width/height

            // Stap 2: Converteer naar sourceRect in fysieke pixels, Y-omlaag, oorsprong linksboven DIT scherm.
            // AANNAMES GEWIJZIGD: selectionRectInPoints en screen.frame lijken al in de 'pixel'-eenheid te zijn
            // die SCDisplay verwacht, ondanks dat ze 'punten' worden genoemd. De 'scale' factor wordt hier dus niet gebruikt voor de conversie.
            let sourceRectXPixels = localOriginXInPoints
            let sourceRectYPixels = (screen.frame.size.height - (localOriginYInPoints + selectionRectInPoints.size.height))
            let sourceRectWidthPixels = selectionRectInPoints.size.width
            let sourceRectHeightPixels = selectionRectInPoints.size.height

            var sourceRect = CGRect(
                x: sourceRectXPixels,
                y: sourceRectYPixels,
                width: sourceRectWidthPixels,
                height: sourceRectHeightPixels
            )
            
            // Rond af naar dichtstbijzijnde gehele pixel om mogelijke SCK API problemen te voorkomen
            sourceRect = CGRect(
                x: round(sourceRect.origin.x),
                y: round(sourceRect.origin.y),
                width: round(sourceRect.size.width),
                height: round(sourceRect.size.height)
            )

            NSLog("🎯 CAPTURE SELECTION DEBUG V3 (Simplified V1 Logic):")
            NSLog("   Input selectionRectInPoints (AppKit Global Y-up): \(selectionRectInPoints)")
            NSLog("   Target NSScreen: \(screen.customLocalizedName), Frame (AppKit Y-up): \(screen.frame), Scale: \(scale)")
            NSLog("   Calculated localOriginInPoints (AppKit Y-up, screen-local): x=\(localOriginXInPoints), y=\(localOriginYInPoints)")
            NSLog("   SCDisplay: ID \(display.displayID), display.width (pixels): \(display.width), display.height (pixels): \(display.height)")
            NSLog("   Calculated sourceRect (Physical Pixels, screen-local, Y-down from screen top-left, rounded): \(sourceRect)")

            // Basis validatie
            guard sourceRect.width > 0 && sourceRect.height > 0 else {
                NSLog("Error V3: Calculated sourceRect has zero or negative rounded width/height. sourceRect: \(sourceRect)")
                return nil
            }

            // Strikte grenscontrole en eventuele clipping
            if !(sourceRect.origin.x >= 0 &&
                 sourceRect.origin.y >= 0 &&
                 sourceRect.maxX <= CGFloat(display.width) + 0.5 &&
                 sourceRect.maxY <= CGFloat(display.height) + 0.5) {
                
                NSLog("Warning V3: Calculated sourceRect \(sourceRect) is out of bounds for the SCDisplay [W:\(display.width), H:\(display.height)]. Attempting to clip.")
                // Log de individuele checks voor duidelijkheid
                NSLog("   Check: sourceRect.origin.x (\(sourceRect.origin.x)) >= 0")
                NSLog("   Check: sourceRect.origin.y (\(sourceRect.origin.y)) >= 0")
                NSLog("   Check: sourceRect.maxX (\(sourceRect.maxX)) <= display.width (\(CGFloat(display.width) + 0.5))")
                NSLog("   Check: sourceRect.maxY (\(sourceRect.maxY)) <= display.height (\(CGFloat(display.height) + 0.5))")

                let clippedRect = sourceRect.intersection(CGRect(x: 0, y: 0, width: CGFloat(display.width), height: CGFloat(display.height)))

                if clippedRect.width > 1 && clippedRect.height > 1 {
                    NSLog("   Successfully clipped sourceRect to: \(clippedRect)")
                    sourceRect = clippedRect
                } else {
                    NSLog("Error V3: Clipping failed or resulted in too small rect: \(clippedRect). Aborting.")
                    return nil
                }
            }
            
            var windowsToExclude = excludingWindows ?? []
            if let ownBundleID = Bundle.main.bundleIdentifier {
                let ownWindows = content.windows.filter { $0.owningApplication?.bundleIdentifier == ownBundleID && $0.isOnScreen }
                windowsToExclude.append(contentsOf: ownWindows)
            }

            let filter = SCContentFilter(display: display, excludingWindows: windowsToExclude)
            let configuration = SCStreamConfiguration()

            configuration.sourceRect = sourceRect 
            configuration.width = Int(sourceRect.width)
            configuration.height = Int(sourceRect.height)
            configuration.showsCursor = SettingsManager.shared.windowCaptureIncludeCursor
            configuration.capturesAudio = false
            configuration.pixelFormat = kCVPixelFormatType_32BGRA
            configuration.colorSpaceName = CGColorSpace.sRGB
            
            let stream = SCStream(filter: filter, configuration: configuration, delegate: nil)
            try await stream.addStreamOutput(SingleFrameOutput.shared, type: .screen, sampleHandlerQueue: .main)
            try await stream.startCapture()
            try await Task.sleep(nanoseconds: 100_000_000) // 0.1 seconde
            let image = await SingleFrameOutput.shared.retrieveFrame()
            try await stream.stopCapture()
            return image
        } catch {
            NSLog("Error capturing selection with SCStream: \(error.localizedDescription)")
            return nil
        }
    }

    func captureWindow(window: SCWindow) async -> NSImage? {
        // Ensure the window is on screen and has a valid frame.
        // SCWindow.frame is in SCK global coordinates (Y-down, origin top-left of main display usually).
        // We need to ensure it has a non-zero size.
        guard window.isOnScreen, window.frame.width > 0, window.frame.height > 0 else {
            NSLog("Error: Window to capture is not on screen or has invalid frame: \(window.windowID), Title: \(window.title ?? "N/A"), Frame: \(window.frame)")
            return nil
        }

        // For capturing a single window, we don't need to exclude other windows explicitly in the filter,
        // as the filter will be configured to only include this specific window.
        // However, we DO need to exclude our own app's overlay windows if they happen to be on top of the target window.
        var windowsToExclude: [SCWindow] = []
        if let ownBundleID = Bundle.main.bundleIdentifier {
            do {
                let content = try await SCShareableContent.current
                let ownWindows = content.windows.filter { $0.owningApplication?.bundleIdentifier == ownBundleID && $0.isOnScreen && $0.windowID != window.windowID }
                windowsToExclude.append(contentsOf: ownWindows)
            } catch {
                 NSLog("Error fetching shareable content for own window exclusion: \(error.localizedDescription)")
            }
        }

        let filter = SCContentFilter(desktopIndependentWindow: window)
        let configuration = SCStreamConfiguration()

        // The frame of SCWindow is already in pixels (SCK coordinates).
        // The width and height should be set to the window's frame size.
        configuration.width = Int(window.frame.width)
        configuration.height = Int(window.frame.height)
        configuration.showsCursor = SettingsManager.shared.windowCaptureIncludeCursor
        configuration.capturesAudio = false
        configuration.pixelFormat = kCVPixelFormatType_32BGRA
        // For window capture, SCContentFilter is configured with a single window, so sourceRect is not needed.
        // The scaleFactor and pointPixelConversion properties on SCWindow might be useful if further coordinate transformations were needed,
        // but captureImage with a window filter typically handles this.

        do {
            let stream = SCStream(filter: filter, configuration: configuration, delegate: nil)
            try await stream.addStreamOutput(SingleFrameOutput.shared, type: .screen, sampleHandlerQueue: .main)
            try await stream.startCapture()
            try await Task.sleep(nanoseconds: 100_000_000) // 0.1 seconde
            let image = await SingleFrameOutput.shared.retrieveFrame()
            try await stream.stopCapture()
            return image
        } catch {
            NSLog("Error capturing window ID \(window.windowID) with SCStream: \(error.localizedDescription)")
            return nil
        }
    }

    // We will add other screenshot methods here later
}

// Helper class to capture a single frame from SCStream
@MainActor
class SingleFrameOutput: NSObject, SCStreamOutput {
    static let shared = SingleFrameOutput()
    private var capturedImage: NSImage?
    private var continuation: CheckedContinuation<NSImage?, Never>?

    // MOET NONISOLATED ZIJN VANWEGE PROTOCOL
    nonisolated func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) {
        // Moet terug naar MainActor voor UI updates/property access
        Task { @MainActor in
            guard type == .screen, CMSampleBufferIsValid(sampleBuffer), CMSampleBufferGetNumSamples(sampleBuffer) == 1 else {
                return
            }

            guard let cvPixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
                return
            }
            
            let ciImage = CIImage(cvPixelBuffer: cvPixelBuffer)
            let rep = NSCIImageRep(ciImage: ciImage)
            let nsImage = NSImage(size: rep.size)
            nsImage.addRepresentation(rep)
            
            self.capturedImage = nsImage
            self.continuation?.resume(returning: nsImage)
            self.continuation = nil
        }
    }
    
    // Deze functie wordt aangeroepen vanaf een andere actor (degene die capture aanroept)
    // maar interacteert met @MainActor properties via de continuation.
    func retrieveFrame() async -> NSImage? {
        if let image = capturedImage { // Lees direct als al beschikbaar (op MainActor)
            self.capturedImage = nil 
            return image
        }
        return await withCheckedContinuation { continuation in
            // De continuation zelf is Sendable.
            // De .resume() wordt aangeroepen vanuit de (nonisolated) stream functie,
            // maar de Task daarbinnen springt terug naar @MainActor voor de daadwerkelijke resume.
            self.continuation = continuation
        }
    }
}