// screen-streamer.swift
// Persistent child process that captures all displays continuously using
// ScreenCaptureKit and encodes frames as H.264 via VideoToolbox.
//
// Output (stdout, binary):
//   On startup: 4-byte magic "RDSM" + JSON line (screen info) + newline
//   Per frame:  [4B length BE] [1B flags] [H.264 NAL data]
//               flags bit0 = keyframe
//
// Control (stdin, JSON lines):
//   {"fps": 30}           u2013 change frame rate
//   {"bitrate": 5000000}  u2013 change bitrate
//   {"keyframe": true}    u2013 force next frame to be a keyframe

import Foundation
import ScreenCaptureKit
import CoreGraphics
import CoreMedia
import VideoToolbox
import CoreVideo

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

var targetFPS: Double = 15
var targetBitrate: Int = 3_000_000
var forceKeyframe = false

// Global references to prevent ARC from deallocating SCStream objects
var activeStream: SCStream?
var activeOutput: AnyObject?
var activeDelegate: AnyObject?

// Currently captured display ID (0 = not yet started)
var currentDisplayID: CGDirectDisplayID = 0

// Cached SCShareableContent for display switching
var cachedContent: SCShareableContent?

// ---------------------------------------------------------------------------
// Binary output helpers
// ---------------------------------------------------------------------------

let stdoutFD = FileHandle.standardOutput
let stderrFD = FileHandle.standardError

// No prefix: the TypeScript side already prefixes every stderr line with
// "[screen-streamer] ", so prefixing here would double it up (e.g.
// "[screen-streamer] [screen-streamer] Frame #600 received").
func log(_ msg: String) {
    stderrFD.write(Data("\(msg)\n".utf8))
}

func writeFrameToStdout(nalData: Data, isKeyframe: Bool) {
    let length = UInt32(nalData.count + 1) // +1 for flags byte
    var lenBE = length.bigEndian
    var header = Data(bytes: &lenBE, count: 4)
    let flags: UInt8 = isKeyframe ? 1 : 0
    header.append(flags)
    stdoutFD.write(header)
    stdoutFD.write(nalData)
}

func writeScreenInfo(_ info: [[String: Any]]) {
    let magic = Data("RDSM".utf8)
    stdoutFD.write(magic)
    if let json = try? JSONSerialization.data(withJSONObject: info),
       let str = String(data: json, encoding: .utf8) {
        stdoutFD.write(Data("\(str)\n".utf8))
    }
}

// ---------------------------------------------------------------------------
// VideoToolbox H.264 encoder
// ---------------------------------------------------------------------------

var compressionSession: VTCompressionSession?
var frameCount: Int64 = 0
var sessionWidth: Int32 = 0
var sessionHeight: Int32 = 0

let encoderCallback: VTCompressionOutputCallback = {
    _, _, status, infoFlags, sampleBuffer in

    guard status == noErr, let sb = sampleBuffer else {
        log("Encoder error: \(status)")
        return
    }

    // Check if keyframe: if kCMSampleAttachmentKey_NotSync is absent or false,
    // the frame is a sync sample (keyframe). Default to true (keyframe) when
    // the key is missing, because VTCompressionSession omits it for IDR frames.
    var isKeyframe = true
    if let attachments = CMSampleBufferGetSampleAttachmentsArray(sb, createIfNecessary: false)
        as? [[CFString: Any]],
       let first = attachments.first,
       let notSync = first[kCMSampleAttachmentKey_NotSync] as? Bool {
        isKeyframe = !notSync
    }

    guard let dataBuffer = CMSampleBufferGetDataBuffer(sb) else { return }
    var totalLength = 0
    var dataPointer: UnsafeMutablePointer<Int8>?
    let ret = CMBlockBufferGetDataPointer(dataBuffer, atOffset: 0, lengthAtOffsetOut: nil,
                                          totalLengthOut: &totalLength, dataPointerOut: &dataPointer)
    guard ret == kCMBlockBufferNoErr, let ptr = dataPointer else { return }

    // For keyframes, prepend SPS/PPS from format description
    var nalData = Data()

    if isKeyframe {
        if let formatDesc = CMSampleBufferGetFormatDescription(sb) {
            // SPS
            var spsSize = 0
            var spsCount = 0
            var spsPtr: UnsafePointer<UInt8>?
            if CMVideoFormatDescriptionGetH264ParameterSetAtIndex(
                formatDesc, parameterSetIndex: 0,
                parameterSetPointerOut: &spsPtr, parameterSetSizeOut: &spsSize,
                parameterSetCountOut: &spsCount, nalUnitHeaderLengthOut: nil) == noErr,
               let sps = spsPtr {
                // Annex B start code
                nalData.append(contentsOf: [0x00, 0x00, 0x00, 0x01])
                nalData.append(UnsafeBufferPointer(start: sps, count: spsSize))
            }
            // PPS
            var ppsSize = 0
            var ppsPtr: UnsafePointer<UInt8>?
            if CMVideoFormatDescriptionGetH264ParameterSetAtIndex(
                formatDesc, parameterSetIndex: 1,
                parameterSetPointerOut: &ppsPtr, parameterSetSizeOut: &ppsSize,
                parameterSetCountOut: nil, nalUnitHeaderLengthOut: nil) == noErr,
               let pps = ppsPtr {
                nalData.append(contentsOf: [0x00, 0x00, 0x00, 0x01])
                nalData.append(UnsafeBufferPointer(start: pps, count: ppsSize))
            }
        }
    }

    // Convert AVCC (length-prefixed) NAL units to Annex B (start-code-prefixed)
    var offset = 0
    while offset < totalLength - 4 {
        var nalLength: UInt32 = 0
        memcpy(&nalLength, ptr + offset, 4)
        nalLength = nalLength.bigEndian
        offset += 4

        if offset + Int(nalLength) > totalLength { break }

        nalData.append(contentsOf: [0x00, 0x00, 0x00, 0x01])
        nalData.append(Data(bytes: ptr + offset, count: Int(nalLength)))
        offset += Int(nalLength)
    }

    writeFrameToStdout(nalData: nalData, isKeyframe: isKeyframe)
}

func createEncoder(width: Int32, height: Int32) {
    if let session = compressionSession {
        VTCompressionSessionInvalidate(session)
        compressionSession = nil
    }

    sessionWidth = width
    sessionHeight = height

    var session: VTCompressionSession?
    let status = VTCompressionSessionCreate(
        allocator: kCFAllocatorDefault,
        width: width,
        height: height,
        codecType: kCMVideoCodecType_H264,
        encoderSpecification: [
            kVTVideoEncoderSpecification_EnableHardwareAcceleratedVideoEncoder: true
        ] as CFDictionary,
        imageBufferAttributes: nil,
        compressedDataAllocator: nil,
        outputCallback: encoderCallback,
        refcon: nil,
        compressionSessionOut: &session
    )

    guard status == noErr, let session = session else {
        log("Failed to create encoder: \(status)")
        return
    }

    compressionSession = session

    // Configure encoder properties
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_RealTime, value: kCFBooleanTrue)
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel,
                         value: kVTProfileLevel_H264_Baseline_AutoLevel)
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_AverageBitRate,
                         value: targetBitrate as CFNumber)
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_MaxKeyFrameInterval,
                         value: Int(targetFPS * 2) as CFNumber)
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ExpectedFrameRate,
                         value: targetFPS as CFNumber)
    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_AllowFrameReordering,
                         value: kCFBooleanFalse) // No B-frames for low latency

    VTCompressionSessionPrepareToEncodeFrames(session)
    log("Encoder created: \(width)x\(height) @ \(targetBitrate/1000)kbps \(targetFPS)fps")
}

func encodeFrame(_ pixelBuffer: CVPixelBuffer, timestamp: CMTime) {
    guard let session = compressionSession else { return }

    let w = Int32(CVPixelBufferGetWidth(pixelBuffer))
    let h = Int32(CVPixelBufferGetHeight(pixelBuffer))

    // Recreate encoder if resolution changed
    if w != sessionWidth || h != sessionHeight {
        log("Resolution changed: \(sessionWidth)x\(sessionHeight) -> \(w)x\(h)")
        createEncoder(width: w, height: h)
        guard compressionSession != nil else { return }
    }

    var props: CFDictionary?
    if forceKeyframe {
        props = [kVTEncodeFrameOptionKey_ForceKeyFrame: true] as CFDictionary
        forceKeyframe = false
    }

    let status = VTCompressionSessionEncodeFrame(
        session,
        imageBuffer: pixelBuffer,
        presentationTimeStamp: timestamp,
        duration: CMTime(value: 1, timescale: CMTimeScale(targetFPS)),
        frameProperties: props,
        sourceFrameRefcon: nil,
        infoFlagsOut: nil
    )

    if status != noErr {
        log("Encode error: \(status)")
    }

    frameCount += 1
}

// ---------------------------------------------------------------------------
// ScreenCaptureKit stream delegate
// ---------------------------------------------------------------------------

@available(macOS 14.0, *)
class StreamOutput: NSObject, SCStreamOutput {
    private var receivedFrames = 0

    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer,
                of type: SCStreamOutputType) {
        guard type == .screen else { return }

        receivedFrames += 1
        if receivedFrames <= 3 || receivedFrames % 300 == 0 {
            log("Frame #\(receivedFrames) received")
        }

        guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
            // SCStream sends frames without pixel data when screen is idle;
            // this is normal and can be silently ignored.
            return
        }

        let pts = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
        encodeFrame(pixelBuffer, timestamp: pts)
    }
}

@available(macOS 14.0, *)
class StreamDelegate: NSObject, SCStreamDelegate {
    func stream(_ stream: SCStream, didStopWithError error: Error) {
        log("SCStream stopped with error: \(error.localizedDescription)")
    }
}

// ---------------------------------------------------------------------------
// Screen info collection
// ---------------------------------------------------------------------------

func getScreenInfo() -> [[String: Any]] {
    var result: [[String: Any]] = []
    var displayCount: UInt32 = 0
    CGGetActiveDisplayList(0, nil, &displayCount)
    var displays = [CGDirectDisplayID](repeating: 0, count: Int(displayCount))
    CGGetActiveDisplayList(displayCount, &displays, &displayCount)

    let mainID = CGMainDisplayID()

    for id in displays {
        let bounds = CGDisplayBounds(id)
        let mode = CGDisplayCopyDisplayMode(id)
        let pw = mode?.pixelWidth ?? Int(bounds.width)
        let ph = mode?.pixelHeight ?? Int(bounds.height)
        let scale = mode.map { Double($0.pixelWidth) / Double($0.width) } ?? 2.0
        let rotation = CGDisplayRotation(id)

        result.append([
            "id": id,
            "x": Int(bounds.origin.x),
            "y": Int(bounds.origin.y),
            "w": Int(bounds.width),
            "h": Int(bounds.height),
            "pixelW": pw,
            "pixelH": ph,
            "scale": scale,
            "rotation": rotation,
            "main": id == mainID,
        ])
    }
    return result
}

// ---------------------------------------------------------------------------
// stdin control reader (runs on background thread)
// ---------------------------------------------------------------------------

func startStdinReader() {
    DispatchQueue.global(qos: .utility).async {
        while let line = readLine() {
            guard let data = line.data(using: .utf8),
                  let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
            else {
                log("Invalid stdin JSON: \(line)")
                continue
            }

            if let fps = json["fps"] as? Double {
                targetFPS = max(1, min(60, fps))
                log("FPS changed to \(targetFPS)")
                // Recreate encoder with new fps
                if sessionWidth > 0 && sessionHeight > 0 {
                    createEncoder(width: sessionWidth, height: sessionHeight)
                }
            }
            if let bitrate = json["bitrate"] as? Int {
                targetBitrate = max(100_000, min(20_000_000, bitrate))
                log("Bitrate changed to \(targetBitrate)")
                if let session = compressionSession {
                    VTSessionSetProperty(session, key: kVTCompressionPropertyKey_AverageBitRate,
                                         value: targetBitrate as CFNumber)
                }
            }
            if let kf = json["keyframe"] as? Bool, kf {
                forceKeyframe = true
                log("Keyframe requested")
            }
            if let switchNum = json["switch"] as? Int {
                let switchID = UInt32(switchNum)
                log("Switch to display \(switchID) requested")
                if #available(macOS 14.0, *) {
                    Task {
                        do {
                            try await switchToDisplay(switchID)
                        } catch {
                            log("Switch display error: \(error.localizedDescription)")
                        }
                    }
                }
            }
        }
        // stdin closed u2014 parent process exited, we should exit too
        log("stdin closed, exiting")
        Foundation.exit(0)
    }
}

// ---------------------------------------------------------------------------
// Main entry point
// ---------------------------------------------------------------------------

/// Switch capture to a specific display (stops old stream if any)
@available(macOS 14.0, *)
func switchToDisplay(_ targetID: CGDirectDisplayID) async throws {
    guard let content = cachedContent else {
        log("No cached content available")
        return
    }

    // Find the display in SCShareableContent
    guard let display = content.displays.first(where: { $0.displayID == targetID }) else {
        log("Display \(targetID) not found in SCShareableContent")
        return
    }

    // Stop old stream if any
    if let oldStream = activeStream {
        log("Stopping old stream for display \(currentDisplayID)...")
        try? await oldStream.stopCapture()
        activeStream = nil
        activeOutput = nil
        activeDelegate = nil
    }

    currentDisplayID = targetID

    // Capture using nominal (logical) resolution.
    // 1:1 with macOS logical coords for correct mouse mapping.
    let filter = SCContentFilter(display: display, excludingWindows: [])
    let config = SCStreamConfiguration()
    config.showsCursor = true
    config.captureResolution = .nominal
    config.width = display.width
    config.height = display.height
    config.minimumFrameInterval = CMTime(value: 1, timescale: CMTimeScale(targetFPS))
    config.queueDepth = 3

    log("Capture config: \(config.width)x\(config.height) (nominal) display=\(targetID)")

    // Reset encoder for new resolution
    createEncoder(width: Int32(config.width), height: Int32(config.height))
    // Force a keyframe on the first frame of the new display
    forceKeyframe = true

    let delegate = StreamDelegate()
    let output = StreamOutput()
    let stream = SCStream(filter: filter, configuration: config, delegate: delegate)
    let queue = DispatchQueue(label: "screen-capture", qos: .userInteractive)
    try stream.addStreamOutput(output, type: .screen, sampleHandlerQueue: queue)

    // Store references globally to prevent ARC deallocation
    activeStream = stream
    activeOutput = output
    activeDelegate = delegate

    try await stream.startCapture()
    log("Capture started: \(config.width)x\(config.height) @ \(targetFPS)fps display=\(targetID)")
}

@available(macOS 14.0, *)
func startCapture() async throws {
    let content = try await SCShareableContent.excludingDesktopWindows(
        false, onScreenWindowsOnly: false
    )
    cachedContent = content

    guard content.displays.first(where: {
        $0.displayID == CGMainDisplayID()
    }) ?? content.displays.first != nil else {
        log("No display found")
        Foundation.exit(1)
    }

    // Output screen info (all displays)
    let screenInfo = getScreenInfo()
    writeScreenInfo(screenInfo)
    log("Displays: \(screenInfo.count)")

    // Start capturing main display
    try await switchToDisplay(CGMainDisplayID())

    // Start stdin reader for control commands
    startStdinReader()
}

if #available(macOS 14.0, *) {
    Task {
        do {
            try await startCapture()
        } catch {
            log("Fatal: \(error.localizedDescription)")
            Foundation.exit(1)
        }
    }
    dispatchMain()
} else {
    log("Requires macOS 14.0+")
    Foundation.exit(1)
}
