//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_encoder.h"
#include "aom/aomcx.h"
#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"

#define AV1_FOURCC 0x31305641

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;
static GlobalTypeLink *encodedFrameDataGT = NULL;
static GlobalTypeLink *encodedFrameDataArrayGT = 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;

typedef struct {
    AvxVideoInfo        info;
    aom_codec_iface_t   *intf;
    aom_codec_ctx_t     codec;
    aom_image_t         image;
	size_t width;
	size_t height;
	size_t frameCount;
	DanaEl* buffer[128];
} InstanceEncoder;

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);

			//printf("plane-%u::stride %u w %u h %u format %u AS %u\n", plane, stride, w, h, img -> fmt, stride * h);

			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;
	}

static int getOptionInt(DanaEl* options, char* name, int fallback, bool* set)
	{
	if (options != NULL)
		{
		int i;
		for (i = 0; i < api -> getArrayLength(options); i++)
			{
			DanaEl* cell = api -> getArrayCellEl(options, i);
			DanaEl* strKey = api -> getDataFieldEl(cell, 0);

			if ((api -> getArrayLength(strKey) == strlen(name)) && memcmp(api -> getArrayContent(strKey), name, strlen(name)) == 0)
				{
				char tmp[256];
				memset(tmp, 0, sizeof(tmp));
				DanaEl* strVal = api -> getDataFieldEl(cell, 1);
				if (strVal != NULL && api -> getArrayLength(strVal) != 0)
					{
					memcpy(tmp, api -> getArrayContent(strVal), api -> getArrayLength(strVal));
					*set = true;
					return atoi(tmp);
					}
				}
			}
		}
	
	return fallback;
	}

INSTRUCTION_DEF op_make_instance_encoder(FrameData *cframe)
	{
	#ifndef WASM
	InstanceEncoder* n = malloc(sizeof(InstanceEncoder));
	memset(n, 0, sizeof(InstanceEncoder));

	aom_codec_enc_cfg_t cfg;

	int usage = 0;
	int speed = 2;

	n -> intf = aom_codec_av1_cx();
	n -> info.codec_fourcc = AV1_FOURCC;
	n -> info.frame_width = api -> getParamInt(cframe, 0);
	n -> info.frame_height = api -> getParamInt(cframe, 1);
	n -> info.time_base.numerator = 1;
	n -> info.time_base.denominator = api -> getParamInt(cframe, 2);

	DanaEl* options = api -> getParamEl(cframe, 3);
	bool bitrate_set = false;
	int bitrate = getOptionInt(options, "bitrate", 4096, &bitrate_set);

	if (!aom_img_alloc(&n -> image, AOM_IMG_FMT_I420, n -> info.frame_width, n -> info.frame_height, 1))
		{
		api -> throwException(cframe, "failed to allocate encoder memory");
		return RETURN_OK;
		}

	//printf("encoder format: %u\n", n -> image.fmt);

	if (aom_codec_enc_config_default(n -> intf, &cfg, usage))
		{
		api -> throwException(cframe, "failed to configure encoder");
		return RETURN_OK;
		}

	cfg.g_w = n -> info.frame_width;
	cfg.g_h = n -> info.frame_height;
	cfg.g_timebase.num = n -> info.time_base.numerator;
	cfg.g_timebase.den = n -> info.time_base.denominator;
	cfg.rc_target_bitrate = bitrate;
	cfg.g_error_resilient = AOM_ERROR_RESILIENT_DEFAULT;

	if (aom_codec_enc_init(&n -> codec, n -> intf, &cfg, 0))
		{
		api -> throwException(cframe, "failed to start encoder");
		return RETURN_OK;
		}

	if (aom_codec_control(&n -> codec, AOME_SET_CPUUSED, speed))
		{
		api -> throwException(cframe, "failed to configure CPU for encoder");
		return RETURN_OK;
		}

	n -> width = api -> getParamInt(cframe, 0);
	n -> height = api -> getParamInt(cframe, 1);

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

	return RETURN_OK;
	}

#ifndef WASM
static DanaEl** encode_frame(InstanceEncoder* n, aom_codec_ctx_t *codec, aom_image_t *img, int frame_index, int flags, int *encodeCount)
	{
	int index = 0;
	aom_codec_iter_t iter = NULL;
	const aom_codec_cx_pkt_t *pkt = NULL;
	const aom_codec_err_t res =
	aom_codec_encode(codec, img, frame_index, 1, flags);
	if (res != AOM_CODEC_OK)
		{
		//printf("ENCODE FAILED\n");
		return NULL;//die_codec(codec, "Failed to encode frame");
		}

	while ((pkt = aom_codec_get_cx_data(codec, &iter)) != NULL)
		{
		DanaEl* ef = api -> makeData(encodedFrameDataGT);
		unsigned char* cnt = NULL;
		DanaEl* bytes = api -> makeArray(charArrayGT, pkt->data.frame.sz, &cnt);
		memcpy(cnt, pkt->data.frame.buf, pkt->data.frame.sz);
		api -> setDataFieldEl(ef, 0, bytes);
		api -> setDataFieldInt(ef, 1, pkt->data.frame.pts);

		n -> buffer[index] = ef;

		index ++;
		}
	
	DanaEl** result = NULL;
	if (index != 0)
		{
		result = malloc(sizeof(DanaEl*) * index);
		int i = 0;
		for (i = 0; i < index; i++)
			{
			result[i] = n -> buffer[i];
			}
		
		*encodeCount = index;
		}

	return result;
	}
#endif

INSTRUCTION_DEF op_encode(FrameData *cframe)
	{
	#ifndef WASM
	InstanceEncoder* handle = NULL;
	memcpy(&handle, api -> getParamRaw(cframe, 0), sizeof(void*));

	DanaEl* img = api -> getParamEl(cframe, 1);
	unsigned char keyFrame = api -> getParamRaw(cframe, 2)[0];

	for (size_t plane = 0; plane < 3; ++plane) {
		unsigned char *buf = handle -> image.planes[plane];
		//calculate our copy-volume as the lower of the two stride values (we already know the w/h are the same between input and output)
		int inputStride = api -> getDataFieldInt(img, plane + 4);
		int outputStride = handle -> image.stride[plane];
		int stride = inputStride;
		if (outputStride < stride) stride = outputStride;
		int w = aom_img_plane_width(&handle -> image, plane) * ((handle -> image.fmt & AOM_IMG_FMT_HIGHBITDEPTH) ? 2 : 1);
		int h = aom_img_plane_height(&handle -> image, plane);
		
		DanaEl* pixels = api -> getDataFieldEl(img, plane + 1);
		unsigned char* rawPix = api -> getArrayContent(pixels);

		for (int i = 0; i < h; i++)
			{
			memcpy(buf, rawPix, stride);
			buf += outputStride;
			rawPix += inputStride;
			}
		}
	
	int flags = 0;
	if (keyFrame) flags |= AOM_EFLAG_FORCE_KF;

	int ec = 0;
	DanaEl** result = encode_frame(handle, &handle -> codec, &handle -> image, handle -> frameCount, flags, &ec);
	if (result != NULL)
		{
		DanaEl* array = api -> makeArray(encodedFrameDataArrayGT, ec, NULL);
		int i = 0;
		for (i = 0; i < ec; i++)
			{
			api -> setArrayCellEl(array, i, result[i]);
			}
		free(result);

		api -> returnEl(cframe, array);
		}
	
	handle -> frameCount ++;

	#endif

	return RETURN_OK;
	}

INSTRUCTION_DEF op_encode_finish(FrameData *cframe)
	{
	#ifndef WASM
	
	InstanceEncoder* handle = NULL;
	memcpy(&handle, api -> getParamRaw(cframe, 0), sizeof(void*));

	// flush the encoder
	DanaEl** buffer[128];
	int countBuffer[128];
	DanaEl** result = NULL;
	int index = 0;
	int totalLength = 0;
	int ec = 0;
	while ((result = encode_frame(handle, &handle -> codec, NULL, -1, 0, &ec)) != NULL)
		{
		buffer[index] = result;
		countBuffer[index] = ec;
		totalLength += ec;
		index ++;
		}
	
	if (totalLength != 0)
		{
		DanaEl* value = api -> makeArray(encodedFrameDataArrayGT, totalLength, NULL);
		int i = 0;
		int ndx = 0;
		for (i = 0; i < index; i++)
			{
			int j = 0;
			for (j = 0; j < countBuffer[i]; j++)
				{
				api -> setArrayCellEl(value, ndx, buffer[i][j]);
				ndx ++;
				}
			
			free(buffer[i]);
			}
		free(result);

		api -> returnEl(cframe, value);
		}
	
		#endif

	return RETURN_OK;
	}

INSTRUCTION_DEF op_destroy_encoder(FrameData *cframe)
	{
	#ifndef WASM
	InstanceEncoder* handle = NULL;
	memcpy(&handle, api -> getParamRaw(cframe, 0), sizeof(void*));

	aom_img_free(&handle -> image);
	aom_codec_destroy(&handle -> codec);

	free(handle);

	#endif

	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);

	AV1Lib_setInterfaceFunction("newEncoder", op_make_instance_encoder);
	AV1Lib_setInterfaceFunction("encode", op_encode);
	AV1Lib_setInterfaceFunction("finishEncode", op_encode_finish);
	AV1Lib_setInterfaceFunction("destroyEncoder", op_destroy_encoder);
	
	charArrayGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("byte[]"));
	frameDataGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("PixelMapYUV"));
	frameDataArrayGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("PixelMapYUV[]"));
	encodedFrameDataGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("EncodedFrame"));
	encodedFrameDataArrayGT = api -> resolveGlobalTypeMapping(AV1Lib_getTypeDefinition("EncodedFrame[]"));
	
	return AV1Lib_getPublicInterface();
	}

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