/*
 * Copyright 2025 NXP
 * All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <stdlib.h>
#include <stdio.h>

#include "image_process.h"
#include "fsl_jpegdec.h"
#include "jpeg.h"
#include "app.h"
#include "vglite_window.h"

#if defined(FSL_SDK_ENABLE_DRIVER_CACHE_CONTROL) && FSL_SDK_ENABLE_DRIVER_CACHE_CONTROL
#include "fsl_cache.h"
#endif

#include "eiq_app.h"
#include "fsl_debug_console.h"
#include "profile_pins.h"

/*******************************************************************************
 * Definitions
 ******************************************************************************/
#define VG_LITE_ALIGN(number, alignment)    \
        (((number) + ((alignment) - 1)) & ~((alignment) - 1))

#define APP_JPEGDEC            (&g_jpegdec)
#define APP_JPEGDEC_IRQHandler JPEGDEC_IRQHandler
#define APP_JPEGDEC_IRQn       JPEGDEC_IRQn

#ifdef MATCH_RESOLUTION_FIXED
  #define RGB_BUFFER_SIZE (MATCH_RES_W * MATCH_RES_H * 3)
#else
  #define RGB_BUFFER_SIZE (320 * 240 * 3)
#endif

/*******************************************************************************
 * Prototypes
 ******************************************************************************/
extern int APP_Decode_JPEG(uint8_t *imagebuffer, uint32_t imagelen, uint32_t *decoded_width, uint32_t *decoded_height, uint8_t *aImgYUVOut, uint8_t *aImgRGBOut);
extern void redraw_img(vg_lite_buffer_t * source_buffer1, vg_lite_buffer_t * source_buffer2, vg_lite_buffer_t * source_buffer3, vg_lite_buffer_t * source_buffer4, box_data *boxes, bool  clear_boxes);

/*******************************************************************************
 * Variables
 ******************************************************************************/
AT_NONCACHEABLE_SECTION(static jpegdec_descpt_t s_decoderDespt);
extern JPEG_DECODER_Type g_jpegdec;
uint8_t *s_frameBufferAddr;
static vg_lite_buffer_t jpegdecblit_source;

AT_NONCACHEABLE_SECTION_ALIGN(uint8_t  usbbuff[2][256*1024], 64);
AT_NONCACHEABLE_SECTION_ALIGN(uint8_t bufML[MODEL_WIDTH * MODEL_HEIGHT * 3], 64);

#ifdef USE_UNCACHED_JPG_BUFFERS
AT_NONCACHEABLE_SECTION_ALIGN(uint8_t  aImgYUV[RGB_BUFFER_SIZE], 64);
AT_NONCACHEABLE_SECTION_ALIGN(uint8_t  aImgRGB[RGB_BUFFER_SIZE], 64);
#elif defined(USE_PSRAM_JPG_BUFFERS)
 uint8_t * aImgYUV = (uint8_t *) 0x60400000U;
 uint8_t * aImgRGB = (uint8_t *) 0x60600000U;

#else
__ALIGNED(64) uint8_t  aImgYUV[RGB_BUFFER_SIZE];
__ALIGNED(64) uint8_t  aImgRGB[RGB_BUFFER_SIZE];
#endif

/*******************************************************************************
 * Code
 ******************************************************************************/
/*
 * @note Address and size should be aligned to XCACHE_LINESIZE_BYTE due to the cache operation unit
 * FSL_FEATURE_XCACHE_LINESIZE_BYTE. On RT700 that is 16 bytes.
 */

int CacheFlushByAddr(uint32_t addr, uint32_t len, uint8_t invalidate)
{

#if defined(FSL_SDK_ENABLE_DRIVER_CACHE_CONTROL) && FSL_SDK_ENABLE_DRIVER_CACHE_CONTROL
    if (!invalidate) // This is for sending data, just flush the cache range to RAM
    {
        /* clear the DCACHE */
        DCACHE_CleanByRange(addr, len);
    }
    else /* (Send/Receive) Flush and invalidate so any new access uses the new data*/
    {
        /* clear the DCACHE */
        DCACHE_CleanInvalidateByRange(addr, len);
    }
#endif
    return 1;
}

int init_buffer(vg_lite_buffer_t * buffer, vg_lite_buffer_format_t fmt, int width,int height, int stride, const uint8_t *mem_source)
{
	memset(buffer, 0, sizeof(buffer));

	 /* Fill in blit information */
	buffer->width	 = width;
	buffer->height	 = height;

	buffer->format = fmt;
	buffer->stride = stride;

	buffer->tiled	 = VG_LITE_LINEAR;
	buffer->compress_mode = VG_LITE_DEC_DISABLE;//VG_LITE_DEC_DISABLE; // DC8000 fbConfig.decMode

	buffer->memory = (uint32_t *) mem_source;
	buffer->address	 = (uint32_t) mem_source;

	buffer->image_mode = VG_LITE_NORMAL_IMAGE_MODE;
	buffer->transparency_mode = VG_LITE_IMAGE_OPAQUE;

	return 1;
}


int jpg_decode(uint8_t *jpg_src_buff, uint32_t jpg_src_buff_size, uint8_t *main_plane_buffer, uint8_t *secondary_plane_buffer)
{
    static jpegdec_config_t config;
    uint32_t status;
    static int frame = 0;
    uint8_t uc1_firstrun = 0;


    if (1)
    {
        memset(&s_decoderDespt, 0U, sizeof(jpegdec_descpt_t));

        /* Step 1: Init JPEG decoder module. */
        JPEGDEC_GetDefaultConfig(&config);
        config.slots = kJPEGDEC_Slot0; /* Enable only one slot. */
        JPEGDEC_Init(APP_JPEGDEC, &config);
        JPEGDEC_SetJpegBuffer(&s_decoderDespt.config, (uint8_t *) jpg_src_buff,  jpg_src_buff_size/*jpegimg_10_29Len*/);

        /* Step 2: Set buffer of generated image for JPEG decoder. */
        JPEGDEC_SetOutputBuffer(&s_decoderDespt.config, (uint8_t *)main_plane_buffer, (uint8_t *)secondary_plane_buffer);

        /* Step 3: Parse header. */
        status = JPEGDEC_ParseHeader(&s_decoderDespt.config);

        if (status != kStatus_Success)
        {
            PRINTF("Unsupported JPEG status:0x%x\r\n", status);
            PRINTF("JPEG info w:%u h:%u pixdepth:%u pixfmt:0x%X\r\n",
            		s_decoderDespt.config.width, s_decoderDespt.config.height,
					s_decoderDespt.config.pixelDepth, s_decoderDespt.config.pixelFormat);
            s_decoderDespt.config.width = 640;
            s_decoderDespt.config.height = 480;
            s_decoderDespt.config.pixelDepth = kJPEGDEC_PixelDepth8Bit;
            s_decoderDespt.config.pixelFormat = kJPEGDEC_PixelFormatYUV422;

        }

        /* Fill in blit information */
        jpegdecblit_source.width	 = s_decoderDespt.config.width;
        jpegdecblit_source.height	 = s_decoderDespt.config.height;

#if USE_SW_YUV420_TO_RGB
        jpegdecblit_source.format	 = VG_LITE_BGR565;
#else
        // Assume YUV420/NV12 format
        if (s_decoderDespt.config.pixelFormat == kJPEGDEC_PixelFormatYUV420)
        {
            jpegdecblit_source.format	 = VG_LITE_NV12;
            jpegdecblit_source.stride  = VG_LITE_ALIGN(jpegdecblit_source.width, 64);

        }
        else if (s_decoderDespt.config.pixelFormat == kJPEGDEC_PixelFormatYUV422)
        {
            jpegdecblit_source.format	 = VG_LITE_YUYV;
            jpegdecblit_source.stride  = VG_LITE_ALIGN(jpegdecblit_source.width*2, 64);
        }
        else if (s_decoderDespt.config.pixelFormat == kJPEGDEC_PixelFormatYUV444)
        {
#if CONVERT_YUV444TORGB
            jpegdecblit_source.format	 = VG_LITE_RGB888;
            //jpegdecblit_source.format	 = VG_LITE_BGR888;

            jpegdecblit_source.stride  = VG_LITE_ALIGN(jpegdecblit_source.width*3, 64);
#else
            jpegdecblit_source.format	 = VG_LITE_YV24;
            jpegdecblit_source.stride  = VG_LITE_ALIGN(jpegdecblit_source.width*3, 64);
#endif
        }
        else
        {
             PRINTF("Unsupported JPEG format :0x%x\r\n", s_decoderDespt.config.pixelFormat);
             assert(false);
        }

#endif

        jpegdecblit_source.tiled	 = VG_LITE_LINEAR;
        jpegdecblit_source.compress_mode = VG_LITE_DEC_DISABLE;//VG_LITE_DEC_DISABLE; // DC8000 fbConfig.decMode

        if((jpegdecblit_source.format==VG_LITE_RGB565) || (jpegdecblit_source.format==VG_LITE_BGR565))
        {
          jpegdecblit_source.memory = (uint32_t *) s_frameBufferAddr;
          jpegdecblit_source.address	 = (uint32_t) s_frameBufferAddr;
          jpegdecblit_source.stride	 = jpegdecblit_source.width *2;
        }
        else if (jpegdecblit_source.format == VG_LITE_RGB888)
        {
          jpegdecblit_source.memory  = (uint32_t *) main_plane_buffer;
          jpegdecblit_source.address = (uint32_t) main_plane_buffer;
        }
        else if (jpegdecblit_source.format == VG_LITE_BGR888)
        {
          jpegdecblit_source.memory  = (uint32_t *) main_plane_buffer;
          jpegdecblit_source.address = (uint32_t) main_plane_buffer;
        }
        else
        { /*YUV420*/
          jpegdecblit_source.memory  = (uint32_t *) main_plane_buffer;
          jpegdecblit_source.address = (uint32_t) main_plane_buffer;
          jpegdecblit_source.width   = VG_LITE_ALIGN(jpegdecblit_source.width, 8);
          jpegdecblit_source.height  = VG_LITE_ALIGN(jpegdecblit_source.height, 8);

          jpegdecblit_source.image_mode = VG_LITE_NORMAL_IMAGE_MODE;
          jpegdecblit_source.transparency_mode = VG_LITE_IMAGE_OPAQUE;
          // Set secondary plane
          if (jpegdecblit_source.format	== VG_LITE_NV12)
          {
              jpegdecblit_source.yuv.uv_stride = jpegdecblit_source.stride;
              jpegdecblit_source.yuv.uv_height = jpegdecblit_source.height / 2;

              jpegdecblit_source.yuv.uv_memory = (uint32_t *) secondary_plane_buffer;
              jpegdecblit_source.yuv.uv_planar = (uint32_t) secondary_plane_buffer;
          }
          else if (jpegdecblit_source.format == VG_LITE_YV24)
          {
              jpegdecblit_source.yuv.uv_stride = jpegdecblit_source.yuv.v_stride = VG_LITE_ALIGN(jpegdecblit_source.width, 64);
              jpegdecblit_source.yuv.uv_height = jpegdecblit_source.yuv.v_height = jpegdecblit_source.height;

              jpegdecblit_source.yuv.uv_planar = (uint32_t)  (main_plane_buffer + jpegdecblit_source.stride * jpegdecblit_source.height);
              jpegdecblit_source.yuv.uv_memory = (uint32_t *) jpegdecblit_source.yuv.uv_planar;

              jpegdecblit_source.yuv.v_planar = (uint32_t) secondary_plane_buffer;
              jpegdecblit_source.yuv.v_memory = (uint32_t *) jpegdecblit_source.yuv.v_planar;
          }
        }

        /* Step 4: Set output pitch, auto start decode when descriptor is loaded. */
        if((jpegdecblit_source.format==VG_LITE_RGB565) || (jpegdecblit_source.format==VG_LITE_BGR565))
        {
            JPEGDEC_SetDecodeOption(&s_decoderDespt.config,  s_decoderDespt.config.width, false, true);
        }
        else
        {
            /* Change the pitch/stride to be a multiple of 64 bytes required for NV12 format */
            JPEGDEC_SetDecodeOption(&s_decoderDespt.config,  jpegdecblit_source.stride /*VG_LITE_ALIGN(s_decoderDespt.config.width, 64)*/, false, true);
        }

        /* Step 5: Set slot descriptor. */
        JPEGDEC_SetSlotNextDescpt(APP_JPEGDEC, 0U, &s_decoderDespt);

        /* Step 6: Enable the descriptor to start the decoding. */
        JPEGDEC_EnableSlotNextDescpt(APP_JPEGDEC, 0U);

    }

    frame++;

    /* Step 7: Wait for decoding complete. */
    status = JPEGDEC_GetStatusFlags(APP_JPEGDEC, 0U);

    while ((status & (kJPEGDEC_DecodeCompleteFlag | kJPEGDEC_ErrorFlags)) == 0U)
    {
        if (status & kJPEGDEC_StreamBufferDoneFlag)
        {
          PRINTF("JPEG stream buffer done status:0x%x\r\n", status);
        }
        status = JPEGDEC_GetStatusFlags(APP_JPEGDEC, 0U);
    }

    if (uc1_firstrun)
    {
       PRINTF("JPEG info w:%u h:%u pixdepth:%u pixfmt:0x%X\r\n",
              s_decoderDespt.config.width, s_decoderDespt.config.height,
              s_decoderDespt.config.pixelDepth, s_decoderDespt.config.pixelFormat);
    }

    if ((status & kJPEGDEC_DecodeCompleteFlag) == 0U)
    {
        JPEGDEC_ClearStatusFlags(APP_JPEGDEC, 0U, status);
        PRINTF("Error occured during JPEG decoding, status:0x%x\r\n", status);
        PRINTF("JPEG decode input size:%i\r\n", jpg_src_buff_size);

        PRINTF("JPEG info w:%u h:%u pixdepth:%u pixfmt:0x%X\r\n",
         s_decoderDespt.config.width, s_decoderDespt.config.height,
              s_decoderDespt.config.pixelDepth, s_decoderDespt.config.pixelFormat);
        assert(false);
    }
    else
    {
#if JPEG_VERBOSE_LOG
        PRINTF("JPEG info w:%u h:%u pixdepth:%u pixfmt:0x%X\r\n",
              s_decoderDespt.config.width, s_decoderDespt.config.height,
              s_decoderDespt.config.pixelDepth, s_decoderDespt.config.pixelFormat);
#endif
    }
    if((jpegdecblit_source.format==VG_LITE_RGB565) || (jpegdecblit_source.format==VG_LITE_BGR565))
    {   /* Step 8: Convert the YUV420 format pixel to RGB565 to display. */
    }
    else if (s_decoderDespt.config.pixelFormat == kJPEGDEC_PixelFormatYUV444)
    {
#if CONVERT_YUV444TORGB
      Convert_yuv444_to_rgb888(s_decoderDespt.config.width, s_decoderDespt.config.height, main_plane_buffer,
                             main_plane_buffer, jpegdecblit_source.stride);
#else
            jpegdecblit_source.stride = jpegdecblit_source.yuv.uv_stride;
#endif
    }

    return 0;
}

/*
 * Process an image from the imagebuffer buffer with a size of image size and encoding as given in format
 * the format is based on the IMG_FORMAT_* defines
 * Currently only IMG_FORMAT_JPG is implemented
 */
int Process_Image_Frame(uint8_t *imagebuffer, uint32_t imagesize, uint32_t width, uint32_t height, uint8_t format)
{
	static uint64_t counter = 0;
	static vg_lite_buffer_t vg_buffer;
	uint8_t inf_res;
	uint8_t *inf_out_img;
	static vg_lite_buffer_t inf_out_vg_buffer;
	static vg_lite_buffer_t logo_vg_buffer;
	static box_data obj_boxes[MAX_BOXES];
	char fileName[32];
	uint32_t decoded_width;
	uint32_t decoded_height;

	SET_CORE0_PROFILE_PIN();

	// Only support MJPEG images
	if (format != IMG_FORMAT_JPG)
	{
		PRINTF("PI unsupported format fmt:%d\r\n", format);
		return -1;
	}

	// Decode the JPEG image into the aImgRGB buffer
	jpg_decode(imagebuffer, imagesize, aImgYUV, NULL);
	decoded_width =  jpegdecblit_source.width;
	decoded_height = jpegdecblit_source.height;

	//Resize the RGB image from the aImgRGB image buffer according to the NN model
	// requirements and convert to input format (if needed) and perform inference.
	// The image in the aImgRGB buffer is saved as  decoded_height lines where each line has decoded_width pixels
	// and each pixel has uint8_t RGB components stored one after the other in packed format R0G0B0 R1G1B1 ...
	init_buffer(&vg_buffer, VG_LITE_RGB888, MODEL_WIDTH, MODEL_HEIGHT, MODEL_WIDTH*3, aImgRGB);

	/*
	 * Workaround to make it work at 325MHz.
	 * Adding a 2ms delay between the hwjepg decoder "done" flag and before we start using the decoded output.
	 * Otherwise we might be trying to use the data before JPEG decoding is done.
	 */
	SDK_DelayAtLeastUs(2000U, SDK_DEVICE_MAXIMUM_CPU_CLOCK_FREQUENCY);

	vg_lite_matrix_t matrix;
	vg_lite_identity(&matrix);
	vg_lite_rectangle_t rect;

	// For input 640x480 and output 320x320 we can consider getting a centered square image at x offset 640-480/2
	// and full height that gets reduced to 320x320
	rect.x = (decoded_width-decoded_height)/2;
	rect.y = 0;
	rect.width =  decoded_height;
	rect.height = decoded_height;
	float scale_factor = MODEL_HEIGHT/(float)decoded_height;
	vg_lite_scale(scale_factor, scale_factor, &matrix);
	vg_lite_blit_rect(&vg_buffer, &jpegdecblit_source, &rect, &matrix,
										VG_LITE_BLEND_NONE,
										0,
										VG_LITE_FILTER_BI_LINEAR); // Possible interpolation methods, according to the vg_lite interface:
																   // VG_LITE_FILTER_POINT - Used for linear paint
																   // VG_LITE_FILTER_LINEAR - Use a 2x2 box around the image pixel and perform an interpolation
																   // VG_LITE_FILTER_BI_LINEAR - Use a 2x2 box around the image pixel and perform an interpolation
																   // VG_LITE_FILTER_GAUSSIAN - Does not seem to work (shows black image)
	vg_lite_finish();

	if (counter % DISPLAY_INFERENCE_FREQ == 0) {
		eIQApp_Inference_Run(aImgRGB, MODEL_WIDTH, MODEL_HEIGHT, 3, &inf_res, obj_boxes);
	}

	redraw_img(&jpegdecblit_source, &vg_buffer, NULL, NULL, obj_boxes, (counter % DISPLAY_INFERENCE_FREQ) == (DISPLAY_INFERENCE_FREQ-1));
	counter++;

	CLR_CORE0_PROFILE_PIN();
	return 0;
}
