kapture/Sources/Kapture/CaptureController.swift

import AppKit
import CoreGraphics
import ScreenCaptureKit
import ApplicationServices

final class CaptureController: NSObject {
    static let shared = CaptureController()

    private var overlays: [OverlayWindow] = []
    private(set) var isSelecting = false
    private var previousApp: NSRunningApplication?
    private var overlayIDs: [Int] = []

    // Selection state (driven by the event tap; AppKit global coordinates)
    private var marqueeStart: NSPoint?
    private(set) var marqueeRect: NSRect?
    private(set) var hoverWindow: (id: CGWindowID, frame: NSRect)?
    var hoverFrame: NSRect? { hoverWindow?.frame }
    private var trackedWindows: [(id: CGWindowID, frame: NSRect)] = []
    private var lastWindowScan = Date.distantPast
    private var eventTap: CFMachPort?

    /// Session event log for debugging (dumped via Copy Diagnostics).
    /// Also appended to /tmp/kapture-debug.log so it survives relaunches.
    static var sessionLog: [String] = []
    static func log(_ line: String) {
        let stamped = String(format: "%@ %@", DateFormatter.localizedString(from: Date(), dateStyle: .none, timeStyle: .medium), line)
        sessionLog.append(stamped)
        if sessionLog.count > 200 { sessionLog.removeFirst() }
        NSLog("[Kapture] \(line)")
        let path = "/tmp/kapture-debug.log"
        if let handle = FileHandle(forWritingAtPath: path) {
            handle.seekToEndOfFile()
            handle.write((stamped + "\n").data(using: .utf8)!)
            handle.closeFile()
        } else {
            FileManager.default.createFile(atPath: path, contents: (stamped + "\n").data(using: .utf8))
        }
    }

    /// Diagnostic snapshot for "Copy Diagnostics" menu item.
    func diagnostics() -> String {
        var lines: [String] = []
        lines.append("Kapture diagnostics — version \(Bundle.main.infoDictionary?["CFBundleShortVersionString"] ?? "?")")
        lines.append("accessibility=\(AXIsProcessTrusted()) eventTap=\(eventTap != nil)")
        lines.append("--- session log (newest last) ---")
        lines.append(contentsOf: Self.sessionLog)
        lines.append("--- screens ---")
        for (i, s) in NSScreen.screens.enumerated() {
            let scale = s.backingScaleFactor
            let cgID = (s.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? NSNumber)?.uint32Value ?? 0
            lines.append("screen[\(i)]: frame=\(s.frame) scale=\(scale) backingSize=\(s.frame.width*scale)x\(s.frame.height*scale) cgID=\(cgID)")
        }
        if #available(macOS 14.0, *) {
            Task {
                if let content = try? await SCShareableContent.current {
                    for (i, d) in content.displays.enumerated() {
                        lines.append("SCDisplay[\(i)]: frame=\(d.frame) width=\(d.width)x\(d.height) id=\(d.displayID)")
                    }
                } else {
                    lines.append("SCShareableContent: unavailable (no Screen Recording permission?)")
                }
                let text = lines.joined(separator: "\n")
                let pb = NSPasteboard.general
                pb.clearContents()
                pb.setString(text, forType: .string)
            }
        } else {
            let text = lines.joined(separator: "\n")
            let pb = NSPasteboard.general
            pb.clearContents()
            pb.setString(text, forType: .string)
        }
        return lines.joined(separator: "\n")
    }

    // MARK: - Session lifecycle

    func begin() {
        guard !isSelecting else { return }
        guard PermissionManager.isGranted else {
            PermissionManager.request()
            return
        }
        isSelecting = true
        previousApp = NSWorkspace.shared.frontmostApplication
        // Activation is async on macOS 26; wait a tick so the app is active
        // before the overlays appear (otherwise clicks pass through).
        NSApp.activate(ignoringOtherApps: true)
        CaptureController.log("begin: activate issued, active=\(NSApp.isActive)")
        DispatchQueue.main.asyncAfter(deadline: .now() + 0.05) { [weak self] in
            guard let self, self.isSelecting else { return }
            CaptureController.log("begin: active=\(NSApp.isActive) — showing overlays")
            self.overlays = NSScreen.screens.map { screen in
                let panel = OverlayWindow(screen: screen, controller: self)
                (panel.contentView as? OverlayView)?.controller = self
                return panel
            }
            self.overlays.forEach { $0.makeKeyAndOrderFront(nil) }
            self.overlayIDs = self.overlays.map { $0.windowNumber }
            CaptureController.log("begin: \(self.overlays.count) overlays, ids=\(self.overlayIDs), active=\(NSApp.isActive)")
            self.installEventTap()
        }
    }

    /// A window was picked from hover — capture it directly (real rounded
    /// corners come from the window's own alpha).
    private func finishWindowPick(_ windowID: CGWindowID, _ rect: NSRect) {
        let ids = overlayIDs
        hideOverlays()
        isSelecting = false
        restorePreviousApp()
        DispatchQueue.main.async { [self] in
            capture(rect, excludingWindowIDs: ids, windowID: windowID)
        }
    }

    private func finishSelection(_ rect: NSRect) {
        let ids = overlayIDs
        hideOverlays()
        guard rect.width >= 2, rect.height >= 2 else {
            cancelSelection()
            return
        }
        DispatchQueue.main.async { [self] in
            capture(rect, excludingWindowIDs: ids)
        }
    }

    private func cancelSelection() {
        hideOverlays()
        isSelecting = false
        restorePreviousApp()
    }

    private func hideOverlays() {
        removeEventTap()
        overlays.forEach { $0.orderOut(nil) }
        overlays = []
        marqueeRect = nil
        hoverWindow = nil
        marqueeStart = nil
    }

    private func restorePreviousApp() {
        previousApp?.activate(options: [.activateIgnoringOtherApps])
        previousApp = nil
    }

    private func renderOverlays() {
        for overlay in overlays {
            (overlay.contentView as? OverlayView)?.needsDisplay = true
        }
    }

    // MARK: - Event tap

    private func installEventTap() {
        guard eventTap == nil else { return }
        guard AXIsProcessTrusted() else {
            CaptureController.log("tap: accessibility not granted — falling back to panel events")
            PermissionManager.requestAccessibility()
            return
        }
        let mask: CGEventMask = (1 << CGEventType.mouseMoved.rawValue)
            | (1 << CGEventType.leftMouseDown.rawValue)
            | (1 << CGEventType.leftMouseDragged.rawValue)
            | (1 << CGEventType.leftMouseUp.rawValue)
            | (1 << CGEventType.keyDown.rawValue)
        guard let tap = CGEvent.tapCreate(
            tap: .cghidEventTap,
            place: .headInsertEventTap,
            options: .defaultTap,
            eventsOfInterest: CGEventMask(mask),
            callback: Self.tapCallback,
            userInfo: Unmanaged.passUnretained(self).toOpaque()
        ) else {
            CaptureController.log("tap: creation FAILED — falling back to panel events")
            return
        }
        eventTap = tap
        if let source = CFMachPortCreateRunLoopSource(nil, tap, 0) {
            CFRunLoopAddSource(CFRunLoopGetMain(), source, .commonModes)
        }
        CGEvent.tapEnable(tap: tap, enable: true)
        CaptureController.log("tap: installed")
    }

    private func removeEventTap() {
        guard let tap = eventTap else { return }
        CGEvent.tapEnable(tap: tap, enable: false)
        if let source = CFMachPortCreateRunLoopSource(nil, tap, 0) {
            CFRunLoopRemoveSource(CFRunLoopGetMain(), source, .commonModes)
        }
        CFMachPortInvalidate(tap)
        eventTap = nil
        CaptureController.log("tap: removed")
    }

    private static let tapCallback: CGEventTapCallBack = { _, type, event, refcon in
        guard let refcon else { return Unmanaged.passUnretained(event) }
        let controller = Unmanaged<CaptureController>.fromOpaque(refcon).takeUnretainedValue()
        if let result = controller.handleTap(type, event) {
            return Unmanaged.passUnretained(result)
        }
        return nil
    }

    /// Convert a CG global point (top-left origin) to AppKit global.
    private func cgPointToAppKit(_ p: NSPoint) -> NSPoint {
        guard let mainScreen = NSScreen.screens.first else { return p }
        return NSPoint(x: p.x, y: mainScreen.frame.height - p.y)
    }

    /// Scan on-screen windows (AppKit global frames), front to back.
    private func refreshTrackedWindows() {
        guard let list = CGWindowListCopyWindowInfo(
            [ .optionOnScreenOnly, .excludeDesktopElements ], kCGNullWindowID
        ) as? [[String: Any]] else { trackedWindows = []; return }

        let ownIDs = Set(overlayIDs)
        var found: [(id: CGWindowID, frame: NSRect)] = []
        for info in list {
            guard let cid = info[kCGWindowNumber as String] as? CGWindowID,
                  let bounds = info[kCGWindowBounds as String] as? [String: CGFloat],
                  let layer = info[kCGWindowLayer as String] as? Int,
                  layer == 0
            else { continue }
            if ownIDs.contains(Int(cid)) { continue }
            let cgRect = NSRect(
                x: bounds["X"] ?? 0,
                y: bounds["Y"] ?? 0,
                width: bounds["Width"] ?? 0,
                height: bounds["Height"] ?? 0
            )
            guard cgRect.width >= 80, cgRect.height >= 60 else { continue }
            let rect = Self.cgToAppKit(rect: cgRect)
            found.append((cid, rect))
        }
        trackedWindows = found
    }

    /// Convert a CGWindowList (top-left origin) frame into AppKit screen
    /// coordinates (bottom-left origin).
    private static func cgToAppKit(rect: NSRect) -> NSRect {
        guard let mainScreen = NSScreen.screens.first else { return rect }
        let y = mainScreen.frame.height - rect.origin.y - rect.height
        return NSRect(x: rect.origin.x, y: y, width: rect.width, height: rect.height)
    }

    private func topmostWindowContaining(_ appKitPoint: NSPoint) -> (id: CGWindowID, frame: NSRect)? {
        for entry in trackedWindows {
            if entry.frame.contains(appKitPoint) {
                return entry
            }
        }
        return nil
    }

    /// Returns the event to pass through, or nil to consume it.
    func handleTap(_ type: CGEventType, _ event: CGEvent) -> CGEvent? {
        guard isSelecting else { return event }

        switch type {
        case .tapDisabledByUserInput, .tapDisabledByTimeout:
            // A timed-out tap must be re-enabled; guard against a torn-down tap.
            if let tap = eventTap {
                CGEvent.tapEnable(tap: tap, enable: true)
                CaptureController.log("tap: re-enabled after \(type)")
            }
            return event
        case .keyDown:
            if event.getIntegerValueField(.keyboardEventKeycode) == 53 {
                CaptureController.log("tap: Esc — cancel")
                cancelSelection()
                return nil
            }
            return event
        case .mouseMoved, .leftMouseDragged:
            // Throttle window scans while moving
            if Date().timeIntervalSince(lastWindowScan) > 0.15 {
                lastWindowScan = Date()
                refreshTrackedWindows()
            }
            let appKitPoint = cgPointToAppKit(NSPoint(x: event.location.x, y: event.location.y))
            if type == .mouseMoved {
                let window = topmostWindowContaining(appKitPoint)
                if window?.frame != hoverFrame {
                    hoverWindow = window
                    renderOverlays()
                }
                return event
            }
            // Drag: grow the marquee
            guard let start = marqueeStart else { return nil }
            let rect = NSRect(
                x: min(start.x, appKitPoint.x),
                y: min(start.y, appKitPoint.y),
                width: abs(appKitPoint.x - start.x),
                height: abs(appKitPoint.y - start.y)
            )
            marqueeRect = rect
            hoverWindow = nil
            renderOverlays()
            return nil
        case .leftMouseDown:
            marqueeStart = cgPointToAppKit(NSPoint(x: event.location.x, y: event.location.y))
            marqueeRect = nil
            CaptureController.log("tap: mouseDown at \(marqueeStart.map { "\($0)" } ?? "?")")
            return nil
        case .leftMouseUp:
            defer { marqueeStart = nil; marqueeRect = nil }
            let appKitPoint = cgPointToAppKit(NSPoint(x: event.location.x, y: event.location.y))
            let dragDistance = marqueeStart.map { hypot($0.x - appKitPoint.x, $0.y - appKitPoint.y) } ?? 0
            if let rect = marqueeRect, dragDistance >= 4, rect.width >= 2, rect.height >= 2 {
                CaptureController.log("tap: marquee complete \(rect)")
                finishSelection(rect)
            } else if let hover = hoverWindow {
                CaptureController.log("tap: window pick id=\(hover.id) \(hover.frame)")
                finishWindowPick(hover.id, hover.frame)
            } else {
                CaptureController.log("tap: click with no hover/marquee — cancel")
                cancelSelection()
            }
            return nil
        default:
            return event
        }
    }

    // MARK: - Capture

    private func capture(_ rect: NSRect, excludingWindowIDs ids: [Int], roundedCorners: Bool = false, windowID: CGWindowID? = nil) {
        ScreenCapture.shared.capture(rect, excludingWindowIDs: ids, roundedCorners: roundedCorners, windowID: windowID) { [weak self] image in
            DispatchQueue.main.async {
                self?.isSelecting = false
                self?.restorePreviousApp()
                Task { @MainActor in
                    guard let image else {
                        ResultWindow.shared.show(image: ResultWindow.failureImage("Capture failed. Check Screen Recording permission."))
                        return
                    }
                    CaptureController.log("capture done: \(image.width)x\(image.height) — showing result window")
                    ResultWindow.shared.show(image: image)
                }
            }
        }
    }
}

final class ScreenCapture {
    static let shared = ScreenCapture()

    private init() {}

    @available(macOS 14.0, *)
    private func crop(_ image: CGImage, to cgRect: CGRect, display: SCDisplay, roundedCorners: Bool = false) -> CGImage? {
        // cgRect is in SCDisplay space (bottom-left origin, Y flipped vs AppKit).
        // The captured image has a TOP-LEFT origin. So within the display,
        // we must flip Y: pixelY = displayHeight - (cgY - displayMinY) - cgHeight.
        let scaleX = CGFloat(image.width) / display.frame.width
        let scaleY = CGFloat(image.height) / display.frame.height
        guard scaleX > 0, scaleY > 0 else { return nil }
        let localMinX = cgRect.minX - display.frame.minX
        let localMaxY = cgRect.maxY - display.frame.minY
        let flippedY = display.frame.height - localMaxY
        let pxRect = CGRect(
            x: localMinX * scaleX,
            y: flippedY * scaleY,
            width: cgRect.width * scaleX,
            height: cgRect.height * scaleY
        )
        // Guard against negative offsets when the selection reaches beyond the
        // display's frame (e.g. secondary monitor coordinate quirks).
        let safeRect = pxRect.intersection(CGRect(x: 0, y: 0, width: CGFloat(image.width), height: CGFloat(image.height)))
        guard safeRect.width >= 1, safeRect.height >= 1 else { return nil }
        // Use a pixel-exact subimage crop instead of redrawing through a
        // CGContext, which resamples and loses quality on retina captures.
        let rounded = safeRect.integral
        guard rounded.width >= 1, rounded.height >= 1 else { return nil }
        guard var result = image.cropping(to: rounded) else { return nil }

        // Apply a rounded-corner alpha mask for window captures, matching
        // macOS window chrome (≈10pt radius, scaled to image pixels).
        if roundedCorners {
            let radius = 10 * scaleX
            let maskRect = CGRect(x: 0, y: 0, width: result.width, height: result.height)
            guard let maskCtx = CGContext(
                data: nil, width: result.width, height: result.height,
                bitsPerComponent: 8, bytesPerRow: 0,
                space: CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
            ) else { return result }
            maskCtx.setFillColor(CGColor(gray: 1, alpha: 1))
            maskCtx.fill(maskRect)
            maskCtx.setBlendMode(.clear)
            maskCtx.setFillColor(CGColor(gray: 0, alpha: 1))
            maskCtx.addPath(CGPath(roundedRect: maskRect, cornerWidth: radius, cornerHeight: radius, transform: nil))
            maskCtx.fillPath()
            maskCtx.setBlendMode(.normal)
            if let mask = maskCtx.makeImage(),
               let masked = result.masking(mask) {
                result = masked
            }
        }
        return result
    }

    @available(macOS 14.0, *)
    private func displayContaining(_ rect: NSRect, in content: SCShareableContent) -> SCDisplay? {
        let mid = NSPoint(x: rect.midX, y: rect.midY)
        guard let screen = NSScreen.screens.first(where: { $0.frame.contains(mid) }) ?? NSScreen.screens.first else {
            return nil
        }
        let cgID = (screen.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? NSNumber)?.uint32Value ?? 0
        return content.displays.first(where: { $0.displayID == cgID }) ?? content.displays.first
    }

    func capture(_ rect: NSRect, excludingWindowIDs excludedIDs: [Int], roundedCorners: Bool = false, windowID: CGWindowID? = nil, completion: @escaping (CGImage?) -> Void) {
        if #available(macOS 14.0, *) {
            Task {
                guard let content = try? await SCShareableContent.current,
                      let display = displayContaining(rect, in: content) else {
                    completion(nil)
                    return
                }
                let excludedSet = Set(excludedIDs.map { CGWindowID($0) })
                let excluded = content.windows.filter { excludedSet.contains($0.windowID) }
                let filter = SCContentFilter(display: display, excludingWindows: excluded)
                let config = SCStreamConfiguration()
                // SCDisplay.width/height are in POINTS; SCStreamConfiguration
                // width/height are output PIXELS. Multiply by the backing scale
                // of the screen containing the selection so retina displays are
                // captured at full resolution and 1x displays stay 1x.
                guard let screen = NSScreen.screens.first(where: { $0.frame.contains(NSPoint(x: rect.midX, y: rect.midY)) }) else {
                    completion(nil)
                    return
                }
                let backingScale = screen.backingScaleFactor
                config.width = Int(CGFloat(display.width) * backingScale)
                config.height = Int(CGFloat(display.height) * backingScale)
                config.showsCursor = false

                // ScreenCaptureKit uses a different origin than AppKit for
                // screens stacked around the main display: per-display,
                // SCDisplay y = appKitY - NSScreen.maxY + SCDisplay.maxY.
                let scRect = CGRect(
                    x: rect.minX,
                    y: rect.minY - screen.frame.maxY + display.frame.maxY,
                    width: rect.width,
                    height: rect.height
                )

                // Window pick: capture the window DIRECTLY — the image comes
                // with the app's real rounded-corner alpha (no mask guessing).
                if let windowID,
                   let scWindow = content.windows.first(where: { $0.windowID == windowID }) {
                    var windowConfig = SCStreamConfiguration()
                    windowConfig.width = Int(scWindow.frame.width * backingScale)
                    windowConfig.height = Int(scWindow.frame.height * backingScale)
                            windowConfig.showsCursor = false
                    if let windowImage = try? await SCScreenshotManager.captureImage(
                        contentFilter: SCContentFilter(desktopIndependentWindow: scWindow),
                        configuration: windowConfig
                    ) {
                        CaptureController.log("capture: WINDOW id=\(windowID) image=\(windowImage.width)x\(windowImage.height) (native alpha corners)")
                        completion(windowImage)
                        return
                    }
                    CaptureController.log("capture: window filter failed for id=\(windowID) — falling back to crop+mask")
                }

                guard let image = try? await SCScreenshotManager.captureImage(contentFilter: filter, configuration: config) else {
                    completion(nil)
                    return
                }
                CaptureController.log("capture: rect=\(rect) screen=\(screen.frame) display=\(display.frame) backing=\(backingScale) image=\(image.width)x\(image.height) scRect=\(scRect)")
                completion(self.crop(image, to: scRect, display: display, roundedCorners: roundedCorners))
            }
        } else {
            guard let primary = NSScreen.screens.first(where: { $0.frame.origin == .zero }) ?? NSScreen.screens.first else {
                completion(nil)
                return
            }
            let cgRect = CGRect(
                x: rect.minX,
                y: primary.frame.height - rect.maxY,
                width: rect.width,
                height: rect.height
            )
            completion(CGWindowListCreateImage(cgRect, .optionOnScreenOnly, kCGNullWindowID, [.bestResolution, .boundsIgnoreFraming]))
        }
    }
}