kapture/Sources/Kapture/ScreenRecorder.swift

import AppKit
import AVFoundation
import CoreMedia
import ScreenCaptureKit

/// Records a screen region (or a single window) to an H.264 .mov via
/// SCStream -> AVAssetWriter. Works on macOS 13+ (no SCScreenshotManager).
///
/// Video-only for now: the audio path is wired behind `includeAudio`
/// (off) — enabling sound later means setting the flag to true, which
/// turns on `capturesAudio`, adds an AAC track, and starts appending
/// audio sample buffers. That path is untested until enabled.
final class ScreenRecorder: NSObject, SCStreamOutput, SCStreamDelegate {
    /// Future: also record system audio into the .mov.
    var includeAudio = false

    private let queue = DispatchQueue(label: "dev.tzgyn.kapture.rec")
    private var stream: SCStream?
    private var writer: AVAssetWriter?
    private var videoInput: AVAssetWriterInput?
    private var audioInput: AVAssetWriterInput?
    private var sessionStarted = false
    private var outputURL: URL?

    /// `region` is in AppKit global points. When `windowID` is given the
    /// stream records that window directly (and follows it if it moves);
    /// otherwise it records `region` via a display filter + sourceRect.
    func start(region: NSRect, windowID: CGWindowID?, excludingWindowIDs: [Int], completion: @escaping (Bool) -> Void) {
        Task {
            do {
                try await prepare(region: region, windowID: windowID, excludingWindowIDs: excludingWindowIDs)
                completion(true)
            } catch {
                CaptureController.log("rec: start failed: \(error.localizedDescription)")
                completion(false)
            }
        }
    }

    private func prepare(region: NSRect, windowID: CGWindowID?, excludingWindowIDs: [Int]) async throws {
        let content = try await SCShareableContent.current

        let excludedSet = Set(excludingWindowIDs.map { CGWindowID($0) })
        let excluded = content.windows.filter { excludedSet.contains($0.windowID) }

        let config = SCStreamConfiguration()
        let filter: SCContentFilter
        let pixelWidth: Int
        let pixelHeight: Int

        let regionCentre = NSPoint(x: region.midX, y: region.midY)
        let screen = NSScreen.screens.first(where: { $0.frame.contains(regionCentre) }) ?? NSScreen.screens.first
        let scale = screen?.backingScaleFactor ?? 2

        if let windowID, let scWindow = content.windows.first(where: { $0.windowID == windowID }) {
            filter = SCContentFilter(desktopIndependentWindow: scWindow)
            pixelWidth = max(2, Int(scWindow.frame.width * scale) / 2 * 2)
            pixelHeight = max(2, Int(scWindow.frame.height * scale) / 2 * 2)
            CaptureController.log("rec: window id=\(windowID) \(pixelWidth)x\(pixelHeight)")
        } else {
            guard let screen else {
                throw NSError(domain: "dev.tzgyn.kapture", code: 1, userInfo: [NSLocalizedDescriptionKey: "No screen for region"])
            }
            let cgID = (screen.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? NSNumber)?.uint32Value ?? 0
            guard let display = content.displays.first(where: { $0.displayID == cgID }) ?? content.displays.first else {
                throw NSError(domain: "dev.tzgyn.kapture", code: 2, userInfo: [NSLocalizedDescriptionKey: "No display for region"])
            }
            filter = SCContentFilter(display: display, excludingWindows: excluded)
            // sourceRect is display-local POINTS with a TOP-LEFT origin
            // (same mapping cropFrozen uses for the frozen snapshots).
            let local = CGRect(
                x: region.minX - screen.frame.minX,
                y: screen.frame.maxY - region.maxY,
                width: region.width,
                height: region.height
            ).intersection(CGRect(x: 0, y: 0, width: display.frame.width, height: display.frame.height))
            guard local.width >= 2, local.height >= 2 else {
                throw NSError(domain: "dev.tzgyn.kapture", code: 3, userInfo: [NSLocalizedDescriptionKey: "Region outside display"])
            }
            config.sourceRect = local
            pixelWidth = max(2, Int(local.width * scale) / 2 * 2)
            pixelHeight = max(2, Int(local.height * scale) / 2 * 2)
            CaptureController.log("rec: region \(region) sourceRect=\(local) \(pixelWidth)x\(pixelHeight)")
        }

        config.width = pixelWidth
        config.height = pixelHeight
        config.minimumFrameInterval = CMTime(value: 1, timescale: 60)
        config.pixelFormat = kCVPixelFormatType_32BGRA
        // Recordings SHOULD include the pointer (unlike stills).
        config.showsCursor = true
        config.queueDepth = 6
        config.capturesAudio = includeAudio

        let movies = FileManager.default.urls(for: .moviesDirectory, in: .userDomainMask).first
            ?? URL(fileURLWithPath: NSHomeDirectory())
        let url = movies.appendingPathComponent("Kapture-\(Self.timestamp()).mov")
        let writer = try AVAssetWriter(url: url, fileType: .mov)

        let videoInput = AVAssetWriterInput(mediaType: .video, outputSettings: [
            AVVideoCodecKey: AVVideoCodecType.h264,
            AVVideoWidthKey: pixelWidth,
            AVVideoHeightKey: pixelHeight,
        ])
        videoInput.expectsMediaDataInRealTime = true
        writer.add(videoInput)

        var audioInput: AVAssetWriterInput?
        if includeAudio {
            let input = AVAssetWriterInput(mediaType: .audio, outputSettings: [
                AVFormatIDKey: kAudioFormatMPEG4AAC,
                AVSampleRateKey: 48000,
                AVNumberOfChannelsKey: 2,
            ])
            input.expectsMediaDataInRealTime = true
            writer.add(input)
            audioInput = input
        }

        guard writer.startWriting() else {
            throw writer.error ?? NSError(domain: "dev.tzgyn.kapture", code: 4, userInfo: [NSLocalizedDescriptionKey: "AVAssetWriter failed to start"])
        }

        let stream = SCStream(filter: filter, configuration: config, delegate: self)
        try stream.addStreamOutput(self, type: .screen, sampleHandlerQueue: queue)
        if includeAudio {
            try stream.addStreamOutput(self, type: .audio, sampleHandlerQueue: queue)
        }
        try await stream.startCapture()

        self.stream = stream
        self.writer = writer
        self.videoInput = videoInput
        self.audioInput = audioInput
        self.outputURL = url
        CaptureController.log("rec: started -> \(url.lastPathComponent)")
    }

    /// Stops the stream and finalizes the file. Completion gets the URL on
    /// success (called on an internal queue — hop to main for UI).
    func stop(completion: @escaping (URL?) -> Void) {
        Task {
            try? await stream?.stopCapture()
            stream = nil
            queue.async {
                guard let writer = self.writer else {
                    completion(nil)
                    return
                }
                self.videoInput?.markAsFinished()
                self.audioInput?.markAsFinished()
                let url = self.outputURL
                self.writer = nil
                self.videoInput = nil
                self.audioInput = nil
                self.sessionStarted = false
                writer.finishWriting {
                    CaptureController.log("rec: finished status=\(writer.status.rawValue)")
                    completion(writer.status == .completed ? url : nil)
                }
            }
        }
    }

    // MARK: SCStreamOutput (called on `queue`)

    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) {
        guard CMSampleBufferDataIsReady(sampleBuffer), let writer, let videoInput else { return }
        switch type {
        case .screen:
            // Skip incomplete frames (macOS 13.1+ exposes a frame status).
            if #available(macOS 13.1, *) {
                guard let attachments = (CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, createIfNecessary: false) as? [NSDictionary])?.first,
                      let dict = attachments as? [String: Any],
                      let statusRaw = dict[SCStreamFrameInfo.status.rawValue] as? Int,
                      SCFrameStatus(rawValue: statusRaw) == .complete else { return }
            }
            guard videoInput.isReadyForMoreMediaData else { return }
            if !sessionStarted {
                sessionStarted = true
                writer.startSession(atSourceTime: CMSampleBufferGetPresentationTimeStamp(sampleBuffer))
            }
            videoInput.append(sampleBuffer)
        case .audio:
            guard includeAudio, sessionStarted, let audioInput, audioInput.isReadyForMoreMediaData else { return }
            audioInput.append(sampleBuffer)
        default:
            break
        }
    }

    // MARK: SCStreamDelegate

    func stream(_ stream: SCStream, didStopWithError error: Error) {
        CaptureController.log("rec: stream stopped with error: \(error.localizedDescription)")
    }

    private static func timestamp() -> String {
        let formatter = DateFormatter()
        formatter.dateFormat = "yyyyMMdd-HHmmss"
        return formatter.string(from: Date())
    }
}