/* ****************************************************************************** *\

INTEL CORPORATION PROPRIETARY INFORMATION
This software is supplied under the terms of a license agreement or nondisclosure
agreement with Intel Corporation and may not be copied or disclosed except in
accordance with the terms of that agreement
Copyright(c) 2011-2015 Intel Corporation. All Rights Reserved.

\* ****************************************************************************** */

#include <stdio.h>
#include <assert.h>

#include "vaapiAllocator.hpp"

enum {
  MFX_FOURCC_VP8_NV12 = MFX_MAKEFOURCC('V', 'P', '8', 'N'),
  MFX_FOURCC_VP8_MBDATA = MFX_MAKEFOURCC('V', 'P', '8', 'M'),
  MFX_FOURCC_VP8_SEGMAP = MFX_MAKEFOURCC('V', 'P', '8', 'S'),
};

unsigned int ConvertMfxFourccToVAFormat(mfxU32 fourcc) {
  switch (fourcc) {
    case MFX_FOURCC_NV12:
      return VA_FOURCC_NV12;
    case MFX_FOURCC_YUY2:
      return VA_FOURCC_YUY2;
    case MFX_FOURCC_YV12:
      return VA_FOURCC_YV12;
    case MFX_FOURCC_RGB4:
      return VA_FOURCC_ARGB;
    case MFX_FOURCC_P8:
      return VA_FOURCC_P208;

    default:
      assert(!"unsupported fourcc");
      return 0;
  }
}

unsigned int ConvertVP8FourccToMfxFourcc(mfxU32 fourcc) {
  switch (fourcc) {
    case MFX_FOURCC_VP8_NV12:
    case MFX_FOURCC_VP8_MBDATA:
      return MFX_FOURCC_NV12;
    case MFX_FOURCC_VP8_SEGMAP:
      return MFX_FOURCC_P8;

    default:
      return fourcc;
  }
}
vaapiFrameAllocator::vaapiFrameAllocator() : m_dpy(0) {}

vaapiFrameAllocator::~vaapiFrameAllocator() { Close(); }

mfxStatus vaapiFrameAllocator::Init(mfxAllocatorParams *pParams) {
  vaapiAllocatorParams *p_vaapiParams = dynamic_cast<vaapiAllocatorParams *>(pParams);

  if ((NULL == p_vaapiParams) || (NULL == p_vaapiParams->m_dpy)) return MFX_ERR_NOT_INITIALIZED;

  m_dpy = p_vaapiParams->m_dpy;

  return MFX_ERR_NONE;
}

mfxStatus vaapiFrameAllocator::CheckRequestType(mfxFrameAllocRequest *request) {
  mfxStatus sts = BaseFrameAllocator::CheckRequestType(request);
  if (MFX_ERR_NONE != sts) return sts;

  if ((request->Type & (MFX_MEMTYPE_VIDEO_MEMORY_DECODER_TARGET | MFX_MEMTYPE_VIDEO_MEMORY_PROCESSOR_TARGET)) != 0)
    return MFX_ERR_NONE;
  else
    return MFX_ERR_UNSUPPORTED;
}

mfxStatus vaapiFrameAllocator::Close() { return BaseFrameAllocator::Close(); }

mfxStatus vaapiFrameAllocator::AllocImpl(mfxFrameAllocRequest *request, mfxFrameAllocResponse *response) {
  mfxStatus mfx_res = MFX_ERR_NONE;
  VAStatus va_res = VA_STATUS_SUCCESS;
  unsigned int va_fourcc = 0;
  VASurfaceID *surfaces = NULL;
  VASurfaceAttrib attrib;
  vaapiMemId *vaapi_mids = NULL, *vaapi_mid = NULL;
  mfxMemId *mids = NULL;
  mfxU32 fourcc = request->Info.FourCC;
  mfxU16 surfaces_num = request->NumFrameSuggested, numAllocated = 0, i = 0;
  bool bCreateSrfSucceeded = false;

  memset(response, 0, sizeof(mfxFrameAllocResponse));

  // VP8 hybrid driver has weird requirements for allocation of surfaces/buffers for VP8 encoding
  // to comply with them additional logic is required to support regular and VP8 hybrid allocation pathes
  mfxU32 mfx_fourcc = ConvertVP8FourccToMfxFourcc(fourcc);
  va_fourcc = ConvertMfxFourccToVAFormat(mfx_fourcc);
  if (!va_fourcc || ((VA_FOURCC_NV12 != va_fourcc) && (VA_FOURCC_YV12 != va_fourcc) && (VA_FOURCC_YUY2 != va_fourcc) &&
                     (VA_FOURCC_ARGB != va_fourcc) && (VA_FOURCC_P208 != va_fourcc))) {
    return MFX_ERR_MEMORY_ALLOC;
  }
  if (!surfaces_num) {
    return MFX_ERR_MEMORY_ALLOC;
  }

  if (MFX_ERR_NONE == mfx_res) {
    surfaces = (VASurfaceID *)calloc(surfaces_num, sizeof(VASurfaceID));
    vaapi_mids = (vaapiMemId *)calloc(surfaces_num, sizeof(vaapiMemId));
    mids = (mfxMemId *)calloc(surfaces_num, sizeof(mfxMemId));
    if ((NULL == surfaces) || (NULL == vaapi_mids) || (NULL == mids)) mfx_res = MFX_ERR_MEMORY_ALLOC;
  }
  if (MFX_ERR_NONE == mfx_res) {
    if (VA_FOURCC_P208 != va_fourcc) {
      unsigned int format;

      attrib.type = VASurfaceAttribPixelFormat;
      attrib.flags = VA_SURFACE_ATTRIB_SETTABLE;
      attrib.value.type = VAGenericValueTypeInteger;
      attrib.value.value.i = va_fourcc;
      format = va_fourcc;

      if (fourcc == MFX_FOURCC_VP8_NV12) {
        // special configuration for NV12 surf allocation for VP8 hybrid encoder is required
        attrib.type = (VASurfaceAttribType)VASurfaceAttribUsageHint;
        attrib.value.value.i = VA_SURFACE_ATTRIB_USAGE_HINT_ENCODER;
      } else if (fourcc == MFX_FOURCC_VP8_MBDATA) {
        // special configuration for MB data surf allocation for VP8 hybrid encoder is required
        attrib.value.value.i = VA_FOURCC_P208;
        format = VA_FOURCC_P208;
      } else if (va_fourcc == VA_FOURCC_NV12) {
        format = VA_RT_FORMAT_YUV420;
      }

      va_res = vaCreateSurfaces(m_dpy, format, request->Info.Width, request->Info.Height, surfaces, surfaces_num,
                                &attrib, 1);

      mfx_res = va_to_mfx_status(va_res);
      bCreateSrfSucceeded = (MFX_ERR_NONE == mfx_res);
    } else {
      VAContextID context_id = request->reserved[0];
      int codedbuf_size;

      int width32 = 32 * ((request->Info.Width + 31) >> 5);
      int height32 = 32 * ((request->Info.Height + 31) >> 5);

      VABufferType codedbuf_type;
      if (fourcc == MFX_FOURCC_VP8_SEGMAP) {
        codedbuf_size = request->Info.Width * request->Info.Height;
        codedbuf_type = (VABufferType)VAEncMacroblockMapBufferType;
      } else {
        codedbuf_size = static_cast<int>((width32 * height32) * 400LL / (16 * 16));
        codedbuf_type = VAEncCodedBufferType;
      }

      for (numAllocated = 0; numAllocated < surfaces_num; numAllocated++) {
        VABufferID coded_buf;

        va_res = vaCreateBuffer(m_dpy, context_id, codedbuf_type, codedbuf_size, 1, NULL, &coded_buf);
        mfx_res = va_to_mfx_status(va_res);
        if (MFX_ERR_NONE != mfx_res) break;
        surfaces[numAllocated] = coded_buf;
      }
    }
  }
  if (MFX_ERR_NONE == mfx_res) {
    for (i = 0; i < surfaces_num; ++i) {
      vaapi_mid = &(vaapi_mids[i]);
      vaapi_mid->m_fourcc = fourcc;
      vaapi_mid->m_surface = &(surfaces[i]);
      mids[i] = vaapi_mid;
    }

    response->mids = mids;
    response->NumFrameActual = surfaces_num;
  } else  // i.e. MFX_ERR_NONE != mfx_res
  {
    response->mids = NULL;
    response->NumFrameActual = 0;
    if (VA_FOURCC_P208 != va_fourcc || fourcc == MFX_FOURCC_VP8_MBDATA) {
      if (bCreateSrfSucceeded) vaDestroySurfaces(m_dpy, surfaces, surfaces_num);
    } else {
      for (i = 0; i < numAllocated; i++) vaDestroyBuffer(m_dpy, surfaces[i]);
    }
    if (mids) {
      free(mids);
      mids = NULL;
    }
    if (vaapi_mids) {
      free(vaapi_mids);
      vaapi_mids = NULL;
    }
    if (surfaces) {
      free(surfaces);
      surfaces = NULL;
    }
  }
  return mfx_res;
}

mfxStatus vaapiFrameAllocator::ReleaseResponse(mfxFrameAllocResponse *response) {
  vaapiMemId *vaapi_mids = NULL;
  mfxU32 i = 0;

  if (!response) return MFX_ERR_NULL_PTR;

  if (response->mids) {
    VASurfaceID *surfaces = NULL;
    vaapi_mids = (vaapiMemId *)(response->mids[0]);
    mfxU32 mfx_fourcc = ConvertVP8FourccToMfxFourcc(vaapi_mids->m_fourcc);
    bool isBitstreamMemory = (MFX_FOURCC_P8 == mfx_fourcc) ? true : false;
    surfaces = vaapi_mids->m_surface;
    for (i = 0; i < response->NumFrameActual; ++i) {
      if (MFX_FOURCC_P8 == vaapi_mids[i].m_fourcc)
        vaDestroyBuffer(m_dpy, surfaces[i]);
      else if (vaapi_mids[i].m_sys_buffer)
        free(vaapi_mids[i].m_sys_buffer);
    }
    free(vaapi_mids);
    free(response->mids);
    response->mids = NULL;

    if (!isBitstreamMemory) {
      vaDestroySurfaces(m_dpy, surfaces, response->NumFrameActual);
    }
    free(surfaces);
  }
  response->NumFrameActual = 0;
  return MFX_ERR_NONE;
}

mfxStatus vaapiFrameAllocator::LockFrame(mfxMemId mid, mfxFrameData *ptr) {
  mfxStatus mfx_res = MFX_ERR_NONE;
  VAStatus va_res = VA_STATUS_SUCCESS;
  vaapiMemId *vaapi_mid = (vaapiMemId *)mid;
  mfxU8 *pBuffer = 0;

  if (!vaapi_mid || !(vaapi_mid->m_surface)) return MFX_ERR_INVALID_HANDLE;

  mfxU32 mfx_fourcc = ConvertVP8FourccToMfxFourcc(vaapi_mid->m_fourcc);

  if (MFX_FOURCC_P8 == mfx_fourcc)  // bitstream processing
  {
    VACodedBufferSegment *coded_buffer_segment;
    if (vaapi_mid->m_fourcc == MFX_FOURCC_VP8_SEGMAP)
      va_res = vaMapBuffer(m_dpy, *(vaapi_mid->m_surface), (void **)(&pBuffer));
    else
      va_res = vaMapBuffer(m_dpy, *(vaapi_mid->m_surface), (void **)(&coded_buffer_segment));
    mfx_res = va_to_mfx_status(va_res);
    if (MFX_ERR_NONE == mfx_res) {
      if (vaapi_mid->m_fourcc == MFX_FOURCC_VP8_SEGMAP)
        ptr->Y = pBuffer;
      else
        ptr->Y = (mfxU8 *)coded_buffer_segment->buf;
    }
  } else  // Image processing
  {
    va_res = vaSyncSurface(m_dpy, *(vaapi_mid->m_surface));
    mfx_res = va_to_mfx_status(va_res);

    if (MFX_ERR_NONE == mfx_res) {
      va_res = vaDeriveImage(m_dpy, *(vaapi_mid->m_surface), &(vaapi_mid->m_image));
      mfx_res = va_to_mfx_status(va_res);
    }
    if (MFX_ERR_NONE == mfx_res) {
      va_res = vaMapBuffer(m_dpy, vaapi_mid->m_image.buf, (void **)&pBuffer);
      mfx_res = va_to_mfx_status(va_res);
    }
    if (MFX_ERR_NONE == mfx_res) {
      switch (vaapi_mid->m_image.format.fourcc) {
        case VA_FOURCC_NV12:
          if (mfx_fourcc == MFX_FOURCC_NV12) {
            ptr->Pitch = (mfxU16)vaapi_mid->m_image.pitches[0];
            ptr->Y = pBuffer + vaapi_mid->m_image.offsets[0];
            ptr->U = pBuffer + vaapi_mid->m_image.offsets[1];
            ptr->V = ptr->U + 1;
          } else
            mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
        case VA_FOURCC_YV12:
          if (mfx_fourcc == MFX_FOURCC_YV12) {
            ptr->Pitch = (mfxU16)vaapi_mid->m_image.pitches[0];
            ptr->Y = pBuffer + vaapi_mid->m_image.offsets[0];
            ptr->V = pBuffer + vaapi_mid->m_image.offsets[1];
            ptr->U = pBuffer + vaapi_mid->m_image.offsets[2];
          } else
            mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
        case VA_FOURCC_YUY2:
          if (mfx_fourcc == MFX_FOURCC_YUY2) {
            ptr->Pitch = (mfxU16)vaapi_mid->m_image.pitches[0];
            ptr->Y = pBuffer + vaapi_mid->m_image.offsets[0];
            ptr->U = ptr->Y + 1;
            ptr->V = ptr->Y + 3;
          } else
            mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
        case VA_FOURCC_ARGB:
          if (mfx_fourcc == MFX_FOURCC_RGB4) {
            ptr->Pitch = (mfxU16)vaapi_mid->m_image.pitches[0];
            ptr->B = pBuffer + vaapi_mid->m_image.offsets[0];
            ptr->G = ptr->B + 1;
            ptr->R = ptr->B + 2;
            ptr->A = ptr->B + 3;
          } else
            mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
        case VA_FOURCC_P208:
          if (mfx_fourcc == MFX_FOURCC_NV12) {
            ptr->Pitch = (mfxU16)vaapi_mid->m_image.pitches[0];
            ptr->Y = pBuffer + vaapi_mid->m_image.offsets[0];
          } else
            mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
        default:
          mfx_res = MFX_ERR_LOCK_MEMORY;
          break;
      }
    }
  }
  return mfx_res;
}

mfxStatus vaapiFrameAllocator::UnlockFrame(mfxMemId mid, mfxFrameData *ptr) {
  vaapiMemId *vaapi_mid = (vaapiMemId *)mid;

  if (!vaapi_mid || !(vaapi_mid->m_surface)) return MFX_ERR_INVALID_HANDLE;

  mfxU32 mfx_fourcc = ConvertVP8FourccToMfxFourcc(vaapi_mid->m_fourcc);

  if (MFX_FOURCC_P8 == mfx_fourcc)  // bitstream processing
  {
    vaUnmapBuffer(m_dpy, *(vaapi_mid->m_surface));
  } else  // Image processing
  {
    vaUnmapBuffer(m_dpy, vaapi_mid->m_image.buf);
    vaDestroyImage(m_dpy, vaapi_mid->m_image.image_id);

    if (NULL != ptr) {
      ptr->Pitch = 0;
      ptr->Y = NULL;
      ptr->U = NULL;
      ptr->V = NULL;
      ptr->A = NULL;
    }
  }
  return MFX_ERR_NONE;
}

mfxStatus vaapiFrameAllocator::GetFrameHDL(mfxMemId mid, mfxHDL *handle) {
  vaapiMemId *vaapi_mid = (vaapiMemId *)mid;

  if (!handle || !vaapi_mid || !(vaapi_mid->m_surface)) return MFX_ERR_INVALID_HANDLE;

  *handle = vaapi_mid->m_surface;  // VASurfaceID* <-> mfxHDL
  return MFX_ERR_NONE;
}
