// ScreenCaptureKit for mcap on macOS, behind a plain C interface (capture.h): ScreenCaptureKit is
// asynchronous, with completion blocks and its own queues, which is simpler to wait on here than from
// Jai. first.jai compiles this into capture.a beside it. Stills for the screen capture overlay, and
// a stream for recording (modules/Screen_Record/macos.jai): its frames and sound, on host time.

#import "capture.h"
#import <ScreenCaptureKit/ScreenCaptureKit.h>
#import <CoreGraphics/CoreGraphics.h>
#import <CoreMedia/CoreMedia.h>
#import <CoreVideo/CoreVideo.h>
#include <stdlib.h>
#include <string.h>

static void set_error(char *error, int error_size, NSString *text) {
    if (!error || error_size <= 0) return;
    const char *utf8 = text.UTF8String ? text.UTF8String : "unknown error";
    strncpy(error, utf8, (size_t) error_size - 1);
    error[error_size - 1] = 0;
}

// The image as packed RGBA, opaque, in sRGB (which the capture is already in): malloc'd.
static uint8_t *rgba_from_image(CGImageRef image, int32_t width, int32_t height) {
    size_t row = (size_t) width * 4;
    uint8_t *pixels = malloc(row * (size_t) height);
    if (!pixels) return NULL;
    CGColorSpaceRef srgb = CGColorSpaceCreateWithName(kCGColorSpaceSRGB);
    CGContextRef bitmap = CGBitmapContextCreate(pixels, (size_t) width, (size_t) height, 8, row, srgb,
                                                kCGBitmapByteOrder32Big | kCGImageAlphaNoneSkipLast);
    CGColorSpaceRelease(srgb);
    if (!bitmap) { free(pixels); return NULL; }
    CGContextSetBlendMode(bitmap, kCGBlendModeCopy);
    CGContextDrawImage(bitmap, CGRectMake(0, 0, width, height), image);
    CGContextRelease(bitmap);
    for (size_t i = 0; i < (size_t) width * (size_t) height; i++) pixels[i * 4 + 3] = 255;
    return pixels;
}

int mcap_sck_screenshots(Mcap_Sck_Shot *shots, int count, char *error, int error_size) {
    @autoreleasepool {
        for (int i = 0; i < count; i++) shots[i].pixels = NULL;

        __block SCShareableContent *content = nil;
        __block NSError *content_error = nil;
        dispatch_semaphore_t got = dispatch_semaphore_create(0);
        [SCShareableContent getShareableContentExcludingDesktopWindows:NO onScreenWindowsOnly:YES
                                                      completionHandler:^(SCShareableContent *c, NSError *e) {
            content = c;
            content_error = e;
            dispatch_semaphore_signal(got);
        }];
        if (dispatch_semaphore_wait(got, dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC)) != 0) {
            set_error(error, error_size, @"ScreenCaptureKit did not list the screens in time");
            return 0;
        }
        if (!content) {
            set_error(error, error_size, content_error ? content_error.localizedDescription : @"ScreenCaptureKit listed no screens");
            return 0;
        }

        // Every screen at once; each completion fills its own entry.
        dispatch_group_t group = dispatch_group_create();
        __block NSString *failure = nil;
        NSLock *lock = [NSLock new];
        for (int i = 0; i < count; i++) {
            SCDisplay *display = nil;
            for (SCDisplay *d in content.displays) if (d.displayID == shots[i].display_id) display = d;
            if (!display) {
                failure = [NSString stringWithFormat:@"ScreenCaptureKit has no display %u", shots[i].display_id];
                continue;
            }
            SCContentFilter *filter = [[SCContentFilter alloc] initWithDisplay:display excludingWindows:@[]];
            SCStreamConfiguration *configuration = [SCStreamConfiguration new];
            configuration.width = (size_t) shots[i].width;
            configuration.height = (size_t) shots[i].height;
            configuration.showsCursor = NO;
            configuration.colorSpaceName = kCGColorSpaceSRGB;
            configuration.captureResolution = SCCaptureResolutionBest;
            Mcap_Sck_Shot *shot = &shots[i];
            dispatch_group_enter(group);
            [SCScreenshotManager captureImageWithFilter:filter configuration:configuration
                                      completionHandler:^(CGImageRef image, NSError *e) {
                if (image) {
                    int32_t w = (int32_t) CGImageGetWidth(image), h = (int32_t) CGImageGetHeight(image);
                    shot->width = w;
                    shot->height = h;
                    shot->pixels = rgba_from_image(image, w, h);
                }
                if (!shot->pixels) {
                    [lock lock];
                    failure = e ? e.localizedDescription : @"ScreenCaptureKit gave no image";
                    [lock unlock];
                }
                dispatch_group_leave(group);
            }];
        }
        if (dispatch_group_wait(group, dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)) != 0) {
            failure = @"ScreenCaptureKit did not copy the screens in time";
        }
        if (failure) {
            // Entries still being filled after a timeout are left to leak rather than raced.
            set_error(error, error_size, failure);
            return 0;
        }
        return 1;
    }
}

void mcap_sck_free(void *pointer) {
    free(pointer);
}

// ---------------------------------------------------------------- recording

#define AUDIO_RATE 48000

// Whether this is macOS 15 or later, for the microphone in the stream. Asked of NSProcessInfo rather
// than with @available, whose check comes from clang's runtime library, which Jai's link leaves out;
// the warnings for the calls it guards are off below.
static BOOL macos_15(void) {
    static int known = -1;
    if (known < 0) known = [[NSProcessInfo processInfo] isOperatingSystemAtLeastVersion:(NSOperatingSystemVersion) {15, 0, 0}] ? 1 : 0;
    return known == 1;
}
#pragma clang diagnostic ignored "-Wunguarded-availability-new"

// A source's packets, waiting to be taken: interleaved stereo float at AUDIO_RATE.
@interface McapAudioPacket : NSObject
@property (nonatomic) NSData *samples;
@property (nonatomic) int64_t hostTime;
@end
@implementation McapAudioPacket
@end

@interface McapStream : NSObject <SCStreamOutput, SCStreamDelegate>
@property (nonatomic) SCStream *stream;
@property (nonatomic) SCStreamConfiguration *configuration;
@property (nonatomic) dispatch_queue_t queue;
@property (nonatomic) NSCondition *condition;   // over everything below
@property (nonatomic) NSMutableData *frame;      // the newest frame, packed BGRA
@property (nonatomic) int64_t frameTime;
@property (nonatomic) BOOL fresh;                // a frame not yet taken
@property (nonatomic) int32_t width, height;
@property (nonatomic) NSString *failure;         // the stream stopped by itself
@property (nonatomic) NSMutableArray<McapAudioPacket *> *system, *microphone;
@end

static int64_t host_nanoseconds(CMTime time) {
    return (int64_t) (CMTimeGetSeconds(time) * 1e9);
}

@implementation McapStream

- (void)stream:(SCStream *)stream didOutputSampleBuffer:(CMSampleBufferRef)buffer ofType:(SCStreamOutputType)type {
    if (!CMSampleBufferIsValid(buffer)) return;
    if (type == SCStreamOutputTypeScreen) [self takeFrame:buffer];
    else if (type == SCStreamOutputTypeAudio) [self takeAudio:buffer into:self.system];
    else if (macos_15()) { if (type == SCStreamOutputTypeMicrophone) [self takeAudio:buffer into:self.microphone]; }
}

- (void)stream:(SCStream *)stream didStopWithError:(NSError *)error {
    [self.condition lock];
    self.failure = error.localizedDescription ?: @"the screen stopped being recorded";
    [self.condition broadcast];
    [self.condition unlock];
}

// A complete frame (not an idle one, which repeats the last), packed into self.frame.
- (void)takeFrame:(CMSampleBufferRef)buffer {
    CFArrayRef attachments = CMSampleBufferGetSampleAttachmentsArray(buffer, false);
    if (attachments && CFArrayGetCount(attachments) > 0) {
        NSDictionary *info = (__bridge NSDictionary *) CFArrayGetValueAtIndex(attachments, 0);
        NSNumber *status = info[SCStreamFrameInfoStatus];
        if (status && status.integerValue != SCFrameStatusComplete) return;
    }
    CVPixelBufferRef image = CMSampleBufferGetImageBuffer(buffer);
    if (!image) return;
    CVPixelBufferLockBaseAddress(image, kCVPixelBufferLock_ReadOnly);
    const uint8_t *base = CVPixelBufferGetBaseAddress(image);
    size_t stride = CVPixelBufferGetBytesPerRow(image);
    size_t w = MIN((size_t) self.width, CVPixelBufferGetWidth(image));
    size_t h = MIN((size_t) self.height, CVPixelBufferGetHeight(image));
    [self.condition lock];
    uint8_t *out = self.frame.mutableBytes;
    size_t row = (size_t) self.width * 4;
    if (base) for (size_t y = 0; y < h; y++) memcpy(out + y * row, base + y * stride, w * 4);
    self.frameTime = host_nanoseconds(CMSampleBufferGetPresentationTimeStamp(buffer));
    self.fresh = YES;
    [self.condition signal];
    [self.condition unlock];
    CVPixelBufferUnlockBaseAddress(image, kCVPixelBufferLock_ReadOnly);
}

// A sound buffer as interleaved stereo float at AUDIO_RATE, onto the source's queue.
- (void)takeAudio:(CMSampleBufferRef)buffer into:(NSMutableArray<McapAudioPacket *> *)queue {
    CMFormatDescriptionRef format = CMSampleBufferGetFormatDescription(buffer);
    const AudioStreamBasicDescription *asbd = format ? CMAudioFormatDescriptionGetStreamBasicDescription(format) : NULL;
    if (!asbd || !(asbd->mFormatFlags & kAudioFormatFlagIsFloat) || asbd->mBitsPerChannel != 32) return;
    CMItemCount frames = CMSampleBufferGetNumSamples(buffer);
    if (frames <= 0) return;
    size_t needed = 0;
    CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(buffer, &needed, NULL, 0, NULL, NULL, 0, NULL);
    AudioBufferList *list = malloc(needed);
    CMBlockBufferRef block = NULL;
    OSStatus status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(buffer, NULL, list, needed, NULL, NULL,
                                                                             kCMSampleBufferFlag_AudioBufferList_Assure16ByteAlignment, &block);
    if (status != noErr) { free(list); return; }
    int channels = (int) asbd->mChannelsPerFrame;
    BOOL planar = (asbd->mFormatFlags & kAudioFormatFlagIsNonInterleaved) != 0;
    // Each frame's left and right: a mono source on both.
    float (^sample)(CMItemCount, int) = ^float(CMItemCount i, int c) {
        int channel = channels >= 2 ? c : 0;
        if (planar) return ((float *) list->mBuffers[channel < (int) list->mNumberBuffers ? channel : 0].mData)[i];
        return ((float *) list->mBuffers[0].mData)[i * channels + channel];
    };
    double ratio = asbd->mSampleRate > 0 ? AUDIO_RATE / asbd->mSampleRate : 1;
    CMItemCount out_frames = (CMItemCount) (frames * ratio + 0.5);
    NSMutableData *samples = [NSMutableData dataWithLength:(size_t) out_frames * 2 * sizeof(float)];
    float *out = samples.mutableBytes;
    for (CMItemCount i = 0; i < out_frames; i++) {
        // Linearly between the source's frames, when its rate is not AUDIO_RATE.
        double at = ratio == 1 ? (double) i : i / ratio;
        CMItemCount a = (CMItemCount) at;
        CMItemCount b = MIN(a + 1, frames - 1);
        float t = (float) (at - (double) a);
        for (int c = 0; c < 2; c++) out[i * 2 + c] = sample(MIN(a, frames - 1), c) * (1 - t) + sample(b, c) * t;
    }
    CFRelease(block);
    free(list);
    McapAudioPacket *packet = [McapAudioPacket new];
    packet.samples = samples;
    packet.hostTime = host_nanoseconds(CMSampleBufferGetPresentationTimeStamp(buffer));
    [self.condition lock];
    [queue addObject:packet];
    [self.condition unlock];
}

@end

static void configure_audio(SCStreamConfiguration *configuration, int system_audio, int microphone) {
    configuration.capturesAudio = system_audio ? YES : NO;
    configuration.sampleRate = AUDIO_RATE;
    configuration.channelCount = 2;
    configuration.excludesCurrentProcessAudio = YES;
    if (macos_15()) configuration.captureMicrophone = microphone ? YES : NO;
}

int mcap_sck_microphone_supported(void) {
    if (macos_15()) return 1;
    return 0;
}

Mcap_Sck_Stream *mcap_sck_stream_start(uint32_t display_id, double x, double y, double w, double h,
                                       int32_t width, int32_t height, int32_t fps, int show_cursor,
                                       int system_audio, int microphone,
                                       const uint32_t *excluded, int excluded_count,
                                       char *error, int error_size) {
    @autoreleasepool {
        __block SCShareableContent *content = nil;
        __block NSError *content_error = nil;
        dispatch_semaphore_t got = dispatch_semaphore_create(0);
        [SCShareableContent getShareableContentExcludingDesktopWindows:NO onScreenWindowsOnly:NO
                                                      completionHandler:^(SCShareableContent *c, NSError *e) {
            content = c;
            content_error = e;
            dispatch_semaphore_signal(got);
        }];
        if (dispatch_semaphore_wait(got, dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC)) != 0 || !content) {
            set_error(error, error_size, content_error ? content_error.localizedDescription : @"ScreenCaptureKit did not list the screens");
            return NULL;
        }
        SCDisplay *display = nil;
        for (SCDisplay *d in content.displays) if (d.displayID == display_id) display = d;
        if (!display) { set_error(error, error_size, [NSString stringWithFormat:@"ScreenCaptureKit has no display %u", display_id]); return NULL; }
        NSMutableArray<SCWindow *> *leave_out = [NSMutableArray array];
        for (SCWindow *window in content.windows)
            for (int i = 0; i < excluded_count; i++) if (window.windowID == excluded[i]) [leave_out addObject:window];
        SCContentFilter *filter = [[SCContentFilter alloc] initWithDisplay:display excludingWindows:leave_out];

        SCStreamConfiguration *configuration = [SCStreamConfiguration new];
        configuration.sourceRect = CGRectMake(x, y, w, h);
        configuration.width = (size_t) width;
        configuration.height = (size_t) height;
        configuration.minimumFrameInterval = CMTimeMake(1, fps > 0 ? fps : 30);
        configuration.pixelFormat = kCVPixelFormatType_32BGRA;
        configuration.colorSpaceName = kCGColorSpaceSRGB;
        configuration.showsCursor = show_cursor ? YES : NO;
        configuration.queueDepth = 6;
        configure_audio(configuration, system_audio, microphone);

        McapStream *holder = [McapStream new];
        holder.configuration = configuration;
        holder.condition = [NSCondition new];
        holder.width = width;
        holder.height = height;
        holder.frame = [NSMutableData dataWithLength:(size_t) width * (size_t) height * 4];
        holder.system = [NSMutableArray array];
        holder.microphone = [NSMutableArray array];
        holder.queue = dispatch_queue_create("mcap.record", DISPATCH_QUEUE_SERIAL);
        holder.stream = [[SCStream alloc] initWithFilter:filter configuration:configuration delegate:holder];
        NSError *add_error = nil;
        BOOL added = [holder.stream addStreamOutput:holder type:SCStreamOutputTypeScreen sampleHandlerQueue:holder.queue error:&add_error];
        if (added) added = [holder.stream addStreamOutput:holder type:SCStreamOutputTypeAudio sampleHandlerQueue:holder.queue error:&add_error];
        if (added) if (macos_15()) added = [holder.stream addStreamOutput:holder type:SCStreamOutputTypeMicrophone sampleHandlerQueue:holder.queue error:&add_error];
        if (!added) { set_error(error, error_size, add_error.localizedDescription ?: @"ScreenCaptureKit would not take the stream's outputs"); return NULL; }

        __block NSError *start_error = nil;
        dispatch_semaphore_t started = dispatch_semaphore_create(0);
        [holder.stream startCaptureWithCompletionHandler:^(NSError *e) { start_error = e; dispatch_semaphore_signal(started); }];
        if (dispatch_semaphore_wait(started, dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)) != 0) {
            set_error(error, error_size, @"ScreenCaptureKit did not start recording in time");
            return NULL;
        }
        if (start_error) { set_error(error, error_size, start_error.localizedDescription); return NULL; }
        return (__bridge_retained Mcap_Sck_Stream *) holder;
    }
}

int mcap_sck_next_frame(Mcap_Sck_Stream *s, uint8_t *pixels, int64_t *host_time_ns, int timeout_ms, char *error, int error_size) {
    McapStream *holder = (__bridge McapStream *) s;
    [holder.condition lock];
    NSDate *until = [NSDate dateWithTimeIntervalSinceNow:timeout_ms / 1000.0];
    while (!holder.fresh && !holder.failure) if (![holder.condition waitUntilDate:until]) break;
    int result = 0;
    if (holder.failure) {
        set_error(error, error_size, holder.failure);
        result = -1;
    } else if (holder.fresh) {
        memcpy(pixels, holder.frame.bytes, holder.frame.length);
        *host_time_ns = holder.frameTime;
        holder.fresh = NO;
        result = 1;
    }
    [holder.condition unlock];
    return result;
}

int mcap_sck_next_audio(Mcap_Sck_Stream *s, int source, float *samples, int max_frames, int64_t *host_time_ns) {
    McapStream *holder = (__bridge McapStream *) s;
    NSMutableArray<McapAudioPacket *> *queue = source == MCAP_SCK_MICROPHONE ? holder.microphone : holder.system;
    [holder.condition lock];
    int frames = 0;
    if (queue.count) {
        McapAudioPacket *packet = queue[0];
        int available = (int) (packet.samples.length / (2 * sizeof(float)));
        frames = MIN(available, max_frames);
        memcpy(samples, packet.samples.bytes, (size_t) frames * 2 * sizeof(float));
        *host_time_ns = packet.hostTime;
        if (frames == available) [queue removeObjectAtIndex:0];
        else {   // the rest stays, its time moved on
            packet.samples = [packet.samples subdataWithRange:NSMakeRange((size_t) frames * 2 * sizeof(float), packet.samples.length - (size_t) frames * 2 * sizeof(float))];
            packet.hostTime += (int64_t) frames * 1000000000LL / AUDIO_RATE;
        }
    }
    [holder.condition unlock];
    return frames;
}

int mcap_sck_stream_set_audio(Mcap_Sck_Stream *s, int system_audio, int microphone, char *error, int error_size) {
    @autoreleasepool {
        McapStream *holder = (__bridge McapStream *) s;
        if (microphone && !mcap_sck_microphone_supported()) { set_error(error, error_size, @"the microphone needs macOS 15 or later"); return 0; }
        configure_audio(holder.configuration, system_audio, microphone);
        __block NSError *update_error = nil;
        dispatch_semaphore_t done = dispatch_semaphore_create(0);
        [holder.stream updateConfiguration:holder.configuration completionHandler:^(NSError *e) { update_error = e; dispatch_semaphore_signal(done); }];
        if (dispatch_semaphore_wait(done, dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC)) != 0) {
            set_error(error, error_size, @"ScreenCaptureKit did not change the sound in time");
            return 0;
        }
        if (update_error) { set_error(error, error_size, update_error.localizedDescription); return 0; }
        // A source turned off keeps nothing queued.
        [holder.condition lock];
        if (!system_audio) [holder.system removeAllObjects];
        if (!microphone) [holder.microphone removeAllObjects];
        [holder.condition unlock];
        return 1;
    }
}

void mcap_sck_stream_stop(Mcap_Sck_Stream *s) {
    if (!s) return;
    McapStream *holder = (__bridge_transfer McapStream *) s;
    @autoreleasepool {
        dispatch_semaphore_t stopped = dispatch_semaphore_create(0);
        [holder.stream stopCaptureWithCompletionHandler:^(NSError *e) { dispatch_semaphore_signal(stopped); }];
        dispatch_semaphore_wait(stopped, dispatch_time(DISPATCH_TIME_NOW, 5 * NSEC_PER_SEC));
        [holder.stream removeStreamOutput:holder type:SCStreamOutputTypeScreen error:nil];
        [holder.stream removeStreamOutput:holder type:SCStreamOutputTypeAudio error:nil];
        if (macos_15()) [holder.stream removeStreamOutput:holder type:SCStreamOutputTypeMicrophone error:nil];
        dispatch_sync(holder.queue, ^{});   // the last sample handler is done
    }
}