/*
g++ vaapi_jpeg_decoder.cpp -o vaapi_jpeg_decoder -ljpeg -lva -lva-drm
./vaapi_jpeg_decoder sample.jpg
*/

#include <iostream>
#include <stdexcept>

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

#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <jpeglib.h>
#include <va/va.h>
#include <va/va_drm.h>

static inline VAStatus vaSafeCall(VAStatus err)
{
  if (err != VA_STATUS_SUCCESS)
    throw std::runtime_error(vaErrorStr(err));
  return err;
}

class VaapiJpegRgbPacketProcessorImpl
{
public:
  int drm_fd;
  VADisplay display;
  VAConfigID config;
  VASurfaceID surface;
  VAContextID context;

  struct jpeg_decompress_struct dinfo;
  struct jpeg_error_mgr jerr;

  static const int WIDTH = 1920;
  static const int HEIGHT = 1080;

  VaapiJpegRgbPacketProcessorImpl()
  {
    dinfo.err = jpeg_std_error(&jerr);
    jpeg_create_decompress(&dinfo);

    initializeVaapi();
  }

  ~VaapiJpegRgbPacketProcessorImpl()
  {
    vaSafeCall(vaDestroyContext(display, context));
    vaSafeCall(vaDestroySurfaces(display, &surface, 1));
    vaSafeCall(vaDestroyConfig(display, config));
    vaSafeCall(vaTerminate(display));
    if (drm_fd >= 0)
      close(drm_fd);
    jpeg_destroy_decompress(&dinfo);
  }

  void initializeVaapi()
  {
    /* Open display */
    static const char *drm_devices[] = {
      "/dev/dri/renderD128",
      "/dev/dri/card0",
      NULL,
    };
    for (int i = 0; drm_devices[i]; i++) {
      drm_fd = open(drm_devices[i], O_RDWR);
      if (drm_fd < 0)
        continue;
      display = vaGetDisplayDRM(drm_fd);
      if (vaDisplayIsValid(display))
        break;
      close(drm_fd);
      drm_fd = -1;
      display = NULL;
    }
    if (!vaDisplayIsValid(display))
      throw std::runtime_error("display not found");

    /* Initialize and create config */
    int major_ver, minor_ver;
    vaSafeCall(vaInitialize(display, &major_ver, &minor_ver));

    const char *driver = vaQueryVendorString(display);
    std::cerr << "[VaapiJpegRgbPacketProcessor::initializeVaapi] driver: " << driver << std::endl;

    int max_entrypoints = vaMaxNumEntrypoints(display);
    if (max_entrypoints <= 0)
      throw std::runtime_error("invalid MaxNumEntrypoints");

    VAEntrypoint entrypoints[max_entrypoints];
    int num_entrypoints;
    vaSafeCall(vaQueryConfigEntrypoints(display, VAProfileJPEGBaseline, entrypoints, &num_entrypoints));
    if (num_entrypoints < 0 || num_entrypoints > max_entrypoints)
      throw std::runtime_error("invalid number of entrypoints");

    int vld_entrypoint;
    for  (vld_entrypoint = 0; vld_entrypoint < num_entrypoints; vld_entrypoint++) {
      if (entrypoints[vld_entrypoint] == VAEntrypointVLD)
        break;
    }
    if (vld_entrypoint == num_entrypoints)
      throw std::runtime_error("did not find VLD");

    VAConfigAttrib attr;
    attr.type = VAConfigAttribRTFormat;
    vaSafeCall(vaGetConfigAttributes(display, VAProfileJPEGBaseline, VAEntrypointVLD, &attr, 1));
	unsigned int rtformat = VA_RT_FORMAT_YUV444;
    if ((attr.value & rtformat) == 0) {
      std::cerr << "[VaapiJpegRgbPacketProcessor::initializeVaapi] warning: YUV444 not supported by libva, chroma will be halved" << std::endl;
      rtformat = VA_RT_FORMAT_YUV420;
    }
    if ((attr.value & rtformat) == 0)
      throw std::runtime_error("does not support YUV420");

    vaSafeCall(vaCreateConfig(display, VAProfileJPEGBaseline, VAEntrypointVLD, &attr, 1, &config));

    /* Create surface and context */
    vaSafeCall(vaCreateSurfaces(display, rtformat, WIDTH, HEIGHT, &surface, 1, NULL, 0));

    vaSafeCall(vaCreateContext(display, config, WIDTH, HEIGHT, 0, &surface, 1, &context));
  }

  VABufferID createBuffer(VABufferType type, unsigned int size, void *data)
  {
    VABufferID buffer;
    vaSafeCall(vaCreateBuffer(display, context, type, size, 1, data, &buffer));
    if (buffer == VA_INVALID_ID)
      throw std::runtime_error("created invalid buffer");
    return buffer;
  }

  void createParameterBuffers(struct jpeg_decompress_struct &dinfo, VABufferID va_bufs[3], VABufferID slice_bufs[2])
  {
    /* Picture Parameter */
    VAPictureParameterBufferJPEGBaseline pic = {0};
    pic.picture_width = dinfo.image_width;
    pic.picture_height = dinfo.image_height;
    for (int i = 0; i< dinfo.num_components; i++) {
      pic.components[i].component_id = dinfo.comp_info[i].component_id;
      pic.components[i].h_sampling_factor = dinfo.comp_info[i].h_samp_factor;
      pic.components[i].v_sampling_factor = dinfo.comp_info[i].v_samp_factor;
      pic.components[i].quantiser_table_selector = dinfo.comp_info[i].quant_tbl_no;
    }
    pic.num_components = dinfo.num_components;

    va_bufs[0] = createBuffer(VAPictureParameterBufferType, sizeof(pic), &pic);

    /* IQ Matrix */
    VAIQMatrixBufferJPEGBaseline iq = {0};
    for (int i = 0; i < NUM_QUANT_TBLS; i++) {
      if (!dinfo.quant_tbl_ptrs[i])
        continue;
      iq.load_quantiser_table[i] = 1;
      /* Assuming dinfo.data_precision == 8 */
      const int natural_order[DCTSIZE2] = {
        0,  1,  8, 16,  9,  2,  3, 10,
        17, 24, 32, 25, 18, 11,  4,  5,
        12, 19, 26, 33, 40, 48, 41, 34,
        27, 20, 13,  6,  7, 14, 21, 28,
        35, 42, 49, 56, 57, 50, 43, 36,
        29, 22, 15, 23, 30, 37, 44, 51,
        58, 59, 52, 45, 38, 31, 39, 46,
        53, 60, 61, 54, 47, 55, 62, 63,
      };

      for (int j = 0; j < DCTSIZE2; j++)
        iq.quantiser_table[i][j] = dinfo.quant_tbl_ptrs[i]->quantval[natural_order[j]];
    }

    va_bufs[1] = createBuffer(VAIQMatrixBufferType, sizeof(iq), &iq);

    /* Huffman Table */
    VAHuffmanTableBufferJPEGBaseline huff = {0};
    const int num_huffman_tables = 2;
    for (int i = 0; i < num_huffman_tables; i++) {
      if (!dinfo.dc_huff_tbl_ptrs[i] || !dinfo.ac_huff_tbl_ptrs[i])
        continue;
      huff.load_huffman_table[i] = 1;
      memcpy(huff.huffman_table[i].num_dc_codes, &dinfo.dc_huff_tbl_ptrs[i]->bits[1],
           sizeof(huff.huffman_table[i].num_dc_codes));
      memcpy(huff.huffman_table[i].dc_values, dinfo.dc_huff_tbl_ptrs[i]->huffval,
           sizeof(huff.huffman_table[i].dc_values));
      memcpy(huff.huffman_table[i].num_ac_codes, &dinfo.ac_huff_tbl_ptrs[i]->bits[1],
           sizeof(huff.huffman_table[i].num_ac_codes));
      memcpy(huff.huffman_table[i].ac_values, dinfo.ac_huff_tbl_ptrs[i]->huffval,
           sizeof(huff.huffman_table[i].ac_values));
    }

    va_bufs[2] = createBuffer(VAHuffmanTableBufferType, sizeof(huff), &huff);

    /* Slice Parameter */
    VASliceParameterBufferJPEGBaseline slice = {0};
    slice.slice_data_size = dinfo.src->bytes_in_buffer;
    slice.slice_data_flag = VA_SLICE_DATA_FLAG_ALL;
    for (int i = 0; i < dinfo.comps_in_scan; i++) {
      slice.components[i].component_selector = dinfo.cur_comp_info[i]->component_id;
      slice.components[i].dc_table_selector = dinfo.cur_comp_info[i]->dc_tbl_no;
      slice.components[i].ac_table_selector = dinfo.cur_comp_info[i]->ac_tbl_no;
    }
    slice.num_components = dinfo.comps_in_scan;
    slice.restart_interval = dinfo.restart_interval;
    unsigned int mcu_h_size = dinfo.max_h_samp_factor * DCTSIZE;
    unsigned int mcu_v_size = dinfo.max_v_samp_factor * DCTSIZE;
    unsigned int mcus_per_row = (WIDTH + mcu_h_size - 1) / mcu_h_size;
    unsigned int mcu_rows_in_scan = (HEIGHT + mcu_v_size - 1) / mcu_v_size;
    slice.num_mcus = mcus_per_row * mcu_rows_in_scan;

    slice_bufs[0] = createBuffer(VASliceParameterBufferType, sizeof(slice), &slice);
    slice_bufs[1] = createBuffer(VASliceDataBufferType, dinfo.src->bytes_in_buffer, (void*)dinfo.src->next_input_byte);
  }

  void decompress(unsigned char *buf, size_t len)
  {
    jpeg_mem_src(&dinfo, buf, len);
    int header_status = jpeg_read_header(&dinfo, true);
    if (header_status != JPEG_HEADER_OK)
      throw std::runtime_error("jpeg not ready for decompression");

    if (dinfo.image_width != WIDTH || dinfo.image_height != HEIGHT)
      throw std::runtime_error("image dimensions do not match preset");

    VABufferID va_bufs[3];
    VABufferID slice_bufs[2];
    createParameterBuffers(dinfo, va_bufs, slice_bufs);

    jpeg_abort_decompress(&dinfo);

    /* Commit buffers */
    vaSafeCall(vaBeginPicture(display, context, surface));
    vaSafeCall(vaRenderPicture(display, context, va_bufs, 3));
    vaSafeCall(vaRenderPicture(display, context, slice_bufs, 2));
    vaSafeCall(vaEndPicture(display, context));

    /* Sync surface */
    vaSafeCall(vaSyncSurface(display, surface));
  }
};

int main(int argc, char **argv)
{
  if (argc < 2)
    return 1;

  struct stat statbuf;
  if (stat(argv[1], &statbuf) < 0 )
    return 1;
  size_t size = statbuf.st_size;

  FILE *f = fopen(argv[1], "r");
  if (!f)
    return 1;

  unsigned char *mem = new unsigned char[size];
  if (!mem)
    return 1;

  fread(mem, size, 1, f);
  fclose(f);

  VaapiJpegRgbPacketProcessorImpl *jdec = new VaapiJpegRgbPacketProcessorImpl();

  for (;;)
  {
    jdec->decompress(mem, size);
  }

  delete[] mem;
  delete jdec;
  return 0;
}