import CoreGraphics
import CoreImage
import CoreMedia
import CoreVideo
import Foundation
import ImageIO
import ScreenCaptureKit

@available(macOS 12.3, *)
final class JpegFrameOutput: NSObject, SCStreamOutput {
    private let stdout = FileHandle.standardOutput
    private let context = CIContext()
    private let quality: CGFloat

    init(quality: CGFloat) {
        self.quality = quality
    }

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

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        let rect = CGRect(x: 0, y: 0, width: width, height: height)
        let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
        guard let cgImage = context.createCGImage(ciImage, from: rect),
              let jpeg = encodeJpeg(cgImage) else {
            return
        }

        guard let header = "FRAME \(width) \(height) \(jpeg.count)\n".data(using: .utf8) else { return }
        stdout.write(header)
        stdout.write(jpeg)
    }

    private func encodeJpeg(_ image: CGImage) -> Data? {
        let data = NSMutableData()
        guard let destination = CGImageDestinationCreateWithData(
            data,
            "public.jpeg" as CFString,
            1,
            nil
        ) else {
            return nil
        }
        let options = [
            kCGImageDestinationLossyCompressionQuality: quality,
            kCGImagePropertyJFIFDictionary: [
                kCGImagePropertyJFIFIsProgressive: false
            ]
        ] as CFDictionary
        CGImageDestinationAddImage(destination, image, options)
        guard CGImageDestinationFinalize(destination) else { return nil }
        return data as Data
    }
}

@available(macOS 12.3, *)
func runJpegCapture() async throws -> Never {
    let args = CommandLine.arguments
    let requestedWidth = Int(args.dropFirst().first ?? "") ?? 960
    let requestedFps = Int(args.dropFirst(2).first ?? "") ?? 30
    let requestedQuality = Int(args.dropFirst(3).first ?? "") ?? 52

    let content = try await SCShareableContent.excludingDesktopWindows(false, onScreenWindowsOnly: true)
    guard let display = content.displays.first else {
        FileHandle.standardError.write(Data("No display is available.\n".utf8))
        exit(2)
    }

    let outputWidth = max(480, min(1920, requestedWidth))
    let scale = Double(outputWidth) / Double(max(display.width, 1))
    let outputHeight = max(270, Int(Double(display.height) * scale))
    let fps = max(1, min(30, requestedFps))
    let quality = CGFloat(max(25, min(90, requestedQuality))) / 100.0

    let filter = SCContentFilter(display: display, excludingWindows: [])
    let configuration = SCStreamConfiguration()
    configuration.width = outputWidth
    configuration.height = outputHeight
    configuration.pixelFormat = kCVPixelFormatType_32BGRA
    configuration.queueDepth = 3
    configuration.minimumFrameInterval = CMTime(value: 1, timescale: CMTimeScale(fps))
    configuration.showsCursor = true

    let output = JpegFrameOutput(quality: quality)
    let stream = SCStream(filter: filter, configuration: configuration, delegate: nil)
    try stream.addStreamOutput(output, type: .screen, sampleHandlerQueue: DispatchQueue(label: "her.screen.mjpeg"))
    try await stream.startCapture()
    while true {
        try await Task.sleep(nanoseconds: 60_000_000_000)
    }
}

if #available(macOS 12.3, *) {
    Task {
        do {
            try await runJpegCapture()
        } catch {
            FileHandle.standardError.write(Data("ScreenCaptureKit JPEG stream failed: \(error)\n".utf8))
            exit(1)
        }
    }
    RunLoop.main.run()
} else {
    FileHandle.standardError.write(Data("ScreenCaptureKit requires macOS 12.3 or later.\n".utf8))
    exit(1)
}
