//Written by Barry Porter, 2025

#include "dana_lib_defs.h"
#include "nli_util.h"
#include "vmi_util.h"

#include 
#include 
#include 

#ifdef WINDOWS
#include 
#endif

#ifdef LINUX
#include 
#include 
#include 
#include 
#endif

#include 

#include "aom/aom_decoder.h"
#include "common/tools_common.h"
#include "common/video_reader.h"
#include "common/video_common.h"
#include "aom/aomdx.h"
#include "aom_ports/mem_ops.h"

void AV1Lib_setInterfaceFunction(char* name, void* ptr);
Interface* AV1Lib_getPublicInterface();
const DanaType* AV1Lib_getTypeDefinition(char* name);

static CoreAPI *api;

static GlobalTypeLink *charArrayGT = NULL;
static GlobalTypeLink *frameDataGT = NULL;
static GlobalTypeLink *frameDataArrayGT = NULL;

typedef struct {
    AvxVideoInfo        info;
    aom_codec_iface_t   *intf;
    aom_codec_ctx_t     codec;
    aom_codec_iter_t    iterator;
    aom_image_t         *image;
	DanaEl* buffer[128];
} Instance;

INSTRUCTION_DEF op_make_instance(FrameData *cframe)
	{
	Instance* n = malloc(sizeof(Instance));
	memset(n, 0, sizeof(Instance));

	// get the av1 decoder interface
    n->intf = aom_codec_av1_dx();
    // init the codec from our chosen interface
    if (aom_codec_dec_init(&n->codec, n->intf, NULL, 0)) {
		api -> throwException(cframe, "decoder init failed");
        return RETURN_OK;
    }

	api -> returnRaw(cframe, (unsigned char*) &n, sizeof(void*));

	return RETURN_OK;
	}

INSTRUCTION_DEF op_decode(FrameData *cframe)
	{
	Instance* handle = NULL;
	memcpy(&handle, api -> getParamRaw(cframe, 0), sizeof(void*));

	//printf("::decode\n");

	DanaEl* array = api -> getParamEl(cframe, 1);
	unsigned char* content = api -> getArrayContent(array);
	size_t contentLen = api -> getArrayLength(array);

	//size_t width = api -> getParamInt(cframe, 2);
	//size_t height = api -> getParamInt(cframe, 3);

	aom_codec_err_t err = 0;

	// decode the frame
	if ((err = aom_codec_decode(&handle->codec, content, contentLen, NULL)) != 0)
		{
		api -> throwException(cframe, aom_codec_err_to_string(err));
		return RETURN_OK;
		}
	
	//we should keep calling get_frame until it returns "null", since a single decode operation can return multiple frames
	int frameCount = 0;
	memset(&handle->iterator, 0, sizeof(handle->iterator));
	while ((handle -> image = aom_codec_get_frame(&handle->codec, &handle->iterator)) != NULL)
		{
		if (handle -> image == NULL)
			{
			api -> throwException(cframe, "image frame get failed");
			return RETURN_OK;
			}
		
		DanaEl* cell = api -> makeData(frameDataGT);
		api -> setDataFieldEl(cell, 0, api -> getParamEl(cframe, 2));

		unsigned char* rawPix = NULL;
		
		aom_image_t* img = handle->image;
		// copy the decoded data into our buffer
		for (size_t plane = 0; plane < 3; ++plane) {
			const unsigned char *buf = img->planes[plane];
			const int stride = img->stride[plane];
			const int w = aom_img_plane_width(img, plane) * ((img->fmt & AOM_IMG_FMT_HIGHBITDEPTH) ? 2 : 1);
			const int h = aom_img_plane_height(img, plane);

			DanaEl* pixels = api -> makeArray(charArrayGT, stride * h, &rawPix);
			memcpy(rawPix, buf, stride * h);

			api -> setDataFieldEl(cell, plane + 1, pixels);
			api -> setDataFieldInt(cell, plane + 4, stride);
			}
		
		handle -> buffer[frameCount] = cell;
		frameCount ++;
		}

	if (frameCount != 0)
		{
		DanaEl* result = api -> makeArray(frameDataArrayGT, frameCount, NULL);

		for (int i = 0; i < frameCount; i++)
			{
			api -> setArrayCellEl(result, i, handle -> buffer[i]);
			}

		api -> returnEl(cframe, result);
		}
	
	return RETURN_OK;
	}

INSTRUCTION_DEF op_destroy(FrameData *cframe)
	{
	Instance* handle = NULL;
	memcpy(&handle, api -> getParamRaw(cframe, 0), sizeof(void*));

	aom_codec_destroy(&handle -> codec);

	free(handle);

	return RETURN_OK;
	}

#ifdef STATIC_NATIVE_LIBRARIES
Interface* AV1Lib_load(CoreAPI *capi)
#else
Interface* load(CoreAPI *capi)
#endif
	{
	api = capi;
	
	AV1Lib_setInterfaceFunction("newDecoder", op_make_instance);
	AV1Lib_setInterfaceFunction("decode", op_decode);
	AV1Lib_setInterfaceFunction("destroyDecoder", op_destroy);
	
	charArrayGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("byte[]"));
	frameDataGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("PixelMapYUV"));
	frameDataArrayGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("PixelMapYUV[]"));
	
	return AV1Lib_getPublicInterface();
	}

#ifdef STATIC_NATIVE_LIBRARIES
void AV1Lib_unload()
#else
void unload()
#endif
	{
	api -> decrementGTRefCount(charArrayGT);
	api -> decrementGTRefCount(frameDataGT);
	api -> decrementGTRefCount(frameDataArrayGT);
	}