/***************************************************************************************
 *
 *  IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
 *
 *  By downloading, copying, installing or using the software you agree to this license.
 *  If you do not agree to this license, do not download, install, 
 *  copy or use the software.
 *
 *  Copyright (C) 2014-2022, Happytimesoft Corporation, all rights reserved.
 *
 *  Redistribution and use in binary forms, with or without modification, are permitted.
 *
 *  Unless required by applicable law or agreed to in writing, software distributed 
 *  under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 *  CONDITIONS OF ANY KIND, either express or implied. See the License for the specific
 *  language governing permissions and limitations under the License.
 *
****************************************************************************************/

#include "sys_inc.h"
#include "window_capture_avf.h"
#include <libavutil/pixfmt.h>
#import <AVFoundation/AVFoundation.h>

/***************************************************************************************/

typedef struct
{
    double              framerate;
    int                 width;
    int                 height;
    char                window_title[256];
    int                 window_id;
    
    enum AVPixelFormat  pixel_format;

    avf_window_callback callback;
    void              * userdata;
} AVFWindowContext;

/***************************************************************************************/

#define WINDOW_NAME     ((NSString *)kCGWindowName)
#define WINDOW_NUMBER   ((NSString *)kCGWindowNumber)
#define WINDOW_LAYER    ((NSString *)kCGWindowLayer)
#define OWNER_NAME      ((NSString *)kCGWindowOwnerName)
#define OWNER_PID       ((NSNumber *)kCGWindowOwnerPID)

static NSComparator win_info_cmp = ^(NSDictionary *o1, NSDictionary *o2) 
{
    NSComparisonResult res = [o1[OWNER_NAME] compare:o2[OWNER_NAME]];
    if (res != NSOrderedSame)
        return res;

    res = [o1[OWNER_PID] compare:o2[OWNER_PID]];
    if (res != NSOrderedSame)
        return res;

    res = [o1[WINDOW_NAME] compare:o2[WINDOW_NAME]];
    if (res != NSOrderedSame)
        return res;

    return [o1[WINDOW_NUMBER] compare:o2[WINDOW_NUMBER]];
};

static CGImageRef avf_window_get_image(int window_id)
{
    NSArray *arr = (NSArray *)CGWindowListCreate(
        kCGWindowListOptionIncludingWindow, window_id);
    [arr autorelease];

    if (!arr.count)
    {
        return NULL;
    }

    return CGWindowListCreateImage(CGRectNull,
    		       kCGWindowListOptionIncludingWindow,
    		       window_id, kCGWindowImageDefault);
}

NSArray * avf_window_enumerate()
{
    NSArray *arr = (NSArray *)CGWindowListCopyWindowInfo(
        kCGWindowListOptionOnScreenOnly | kCGWindowListExcludeDesktopElements, kCGNullWindowID);

    [arr autorelease];

    return [arr sortedArrayUsingComparator:win_info_cmp];
}

void avf_window_list()
{
    NSArray *arr = avf_window_enumerate();

    printf("\r\nAvailable window name : \r\n\r\n");

    for (NSDictionary *dict in arr) 
    {
        NSNumber *layer = (NSNumber *)dict[WINDOW_LAYER];
        if (0 != layer.intValue)
        {
            continue;
        }

        NSString *name = (NSString *)dict[WINDOW_NAME];

        printf("%s\r\n", name.UTF8String);
    }
}

void * avf_window_init(char * title, int width, int height, double framerate)
{
    AVFWindowContext * context = (AVFWindowContext *)malloc(sizeof(AVFWindowContext));
    if (NULL == context)
    {
        return NULL;
    }
    
    memset(context, 0, sizeof(AVFWindowContext));

    strncpy(context->window_title, title, sizeof(context->window_title));
    context->width = width;
    context->height = height;
    context->framerate = framerate;

    NSArray *arr = avf_window_enumerate();
    
    for (NSDictionary *dict in arr) 
    {
        NSNumber *layer = (NSNumber *)dict[WINDOW_LAYER];
        if (0 != layer.intValue)
        {
            continue;
        }
        
        NSString *name = (NSString *)dict[WINDOW_NAME];
        NSNumber *wid = (NSNumber *)dict[WINDOW_NUMBER];

        if (strncasecmp(name.UTF8String, title, strlen(title)) == 0)
        {
            context->window_id = wid.intValue;
            break;
        }
    }

    if (0 == context->window_id)
    {
        free(context);
        log_print(HT_LOG_ERR, "%s, not found window, %s\r\n", __FUNCTION__, title);
        return NULL;
    }

    CGImageRef img = avf_window_get_image(context->window_id);
    if (!img)
    {
        free(context);
        log_print(HT_LOG_ERR, "%s, avf_window_get_image failed\r\n", __FUNCTION__);
        return NULL;
    }

    size_t img_width = CGImageGetWidth(img);
    size_t img_height = CGImageGetHeight(img);

    if (!img_width || !img_height || CGImageGetBitsPerPixel(img) != 32 ||
        CGImageGetBitsPerComponent(img) != 8) 
    {
        CGImageRelease(img);
        free(context);
        log_print(HT_LOG_ERR, "%s, invalid image format!\r\n", __FUNCTION__);
        return NULL;
    }

    CGImageRelease(img);

    context->width = img_width;
    context->height = img_height;
    context->pixel_format = AV_PIX_FMT_BGR0;

    return context;
}

void avf_window_uninit(void * ctx)
{
    if (NULL == ctx)
    {
        return;
    }

    AVFWindowContext * context = (AVFWindowContext *)ctx;
    
    free(context);
}

void avf_window_set_callback(void * ctx, avf_window_callback cb, void * userdata)
{
    if (NULL == ctx)
    {
        return;
    }

    AVFWindowContext * context = (AVFWindowContext *)ctx;

    context->callback = cb;
    context->userdata = userdata;
}

int avf_window_get_width(void * ctx)
{
    if (NULL == ctx)
    {
        return 0;
    }

    AVFWindowContext * context = (AVFWindowContext *)ctx;

    return context->width;
}

int avf_window_get_height(void * ctx)
{
    if (NULL == ctx)
    {
        return 0;
    }

    AVFWindowContext * context = (AVFWindowContext *)ctx;

    return context->height;
}

int avf_window_get_pixfmt(void * ctx)
{
    if (NULL == ctx)
    {
        return AV_PIX_FMT_NONE;
    }

    AVFWindowContext * context = (AVFWindowContext *)ctx;

    return context->pixel_format;
}

BOOL avf_window_capture(AVFWindowContext * context)
{
    CGImageRef img = avf_window_get_image(context->window_id);
    if (!img)
    {
        return FALSE;
    }

    size_t width = CGImageGetWidth(img);
    size_t height = CGImageGetHeight(img);

    if (!width || !height || CGImageGetBitsPerPixel(img) != 32 ||
        CGImageGetBitsPerComponent(img) != 8) 
    {
        CGImageRelease(img);
        return FALSE;
    }

    CGDataProviderRef provider = CGImageGetDataProvider(img);
    CFDataRef data = CGDataProviderCopyData(provider);

    avf_window_data frame;

    memset(&frame, 0, sizeof(frame));

    frame.width  = width;
    frame.height = height;
    frame.format = context->pixel_format;
    frame.data[0] = (uint8 *)CFDataGetBytePtr(data);
    frame.linesize[0] = CGImageGetBytesPerRow(img);

    if (context->callback)
    {
        context->callback(&frame, context->userdata);
    }

    CGImageRelease(img);
    CFRelease(data);

    return TRUE;
}

BOOL avf_window_read(void * ctx)
{
    if (NULL == ctx)
    {
        return FALSE;
    }

    BOOL ret = FALSE;
    AVFWindowContext * context = (AVFWindowContext *)ctx;

    @autoreleasepool {
        ret = avf_window_capture(context);
    }

    return ret;
}