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

#include <Argus/Argus.h>
#include <EGLStream/EGLStream.h>
#include <EGLStream/NV/ImageNativeBuffer.h>
#include <NvEglRenderer.h>

#include <cuda_runtime.h>
#include "cudaEGL.h"

#include "Error.h"

using namespace Argus;
using namespace EGLStream;

NvEglRenderer *g_renderer = NULL;
Size2D<uint32_t> STREAM_SIZE(1920, 1080);

bool do_exit = false;

static void *render_thread(void *arg) {
  OutputStream *stream = (OutputStream *)arg;

  UniqueObj<FrameConsumer> consumer;
  consumer.reset(FrameConsumer::create(stream));

  IEGLOutputStream *iEglOutputStream;
  IFrameConsumer *iFrameConsumer;

  iEglOutputStream = interface_cast<IEGLOutputStream>(stream);
  iFrameConsumer = interface_cast<IFrameConsumer>(consumer);
  if (!iFrameConsumer)
    fprintf(stderr, "Failed to get IFrameConsumer interface");

  /* Wait until the producer has connected to the stream */
  printf("Waiting until producer is connected...\n");
  if (iEglOutputStream->waitUntilConnected() != STATUS_OK) {
    fprintf(stderr, "Stream failed to connect.");
    return NULL;
  }
  printf("Producer has connected; continuing.\n");

  NvBufferCreateParams params = {
    .width=(int)STREAM_SIZE.width(),
    .height=(int)STREAM_SIZE.height(),
    .payloadType=NvBufferPayload_SurfArray,
    .layout=NvBufferLayout_Pitch,
    .colorFormat=NvBufferColorFormat_YUV420,
  };

  int dmabuf;
  int err = NvBufferCreateEx(&dmabuf, &params);
  assert(err == 0);

  NvBufferParams p;
  NvBufferGetParams(dmabuf, &p);
  printf("nvbuf pitches: %d %d %d\n", p.pitch[0], p.pitch[1], p.pitch[2]);

  int frame_id = 0;
  while (!do_exit) {
    UniqueObj<Frame> frame(iFrameConsumer->acquireFrame());
    IFrame *iFrame = interface_cast<IFrame>(frame);
    if (!iFrame) {
      fprintf(stderr, "failed to get frame");
      break;
    }

    NV::IImageNativeBuffer *iNativeBuffer = interface_cast<NV::IImageNativeBuffer>(iFrame->getImage());
    if (!iNativeBuffer)
      fprintf(stderr, "IImageNativeBuffer not supported by Image.");
    if (iNativeBuffer->copyToNvBuffer(dmabuf) != STATUS_OK)
      fprintf(stderr, "Failed to copy frame to NvBuffer.");

    printf("GOT FRAME %d\n", frame_id);

    EGLImageKHR egl_image = NULL;
    egl_image = NvEGLImageFromFd(NULL, dmabuf);
    if (egl_image == NULL)
      fprintf(stderr, "Error while mapping dmabuf fd (0x%X) to EGLImage\n", 0);

    CUresult status;
    CUeglFrame eglFrame;
    CUgraphicsResource pResource = NULL;

    cudaFree(0);
    if (cuGraphicsEGLRegisterImage(&pResource, egl_image, CU_GRAPHICS_MAP_RESOURCE_FLAGS_NONE) != CUDA_SUCCESS)
      fprintf(stderr, "cuGraphicsEGLRegisterImage failed: %d, cuda process stop\n", status);

    if (cuGraphicsResourceGetMappedEglFrame(&eglFrame, pResource, 0, 0) != CUDA_SUCCESS)
      fprintf(stderr, "cuGraphicsSubResourceGetMappedArray failed\n");

    if (cuCtxSynchronize() != CUDA_SUCCESS)
      fprintf(stderr, "cuCtxSynchronize failed");

    assert(eglFrame.frameType == CU_EGL_FRAME_TYPE_PITCH);

    if (frame_id == 30) {
       uint8_t *out = (uint8_t *)malloc(p.pitch[0]*STREAM_SIZE.height());
       cudaMemcpy((void *)out, eglFrame.frame.pPitch[0], p.pitch[0]*STREAM_SIZE.height(), cudaMemcpyDeviceToHost);
       FILE *f = fopen("out.mono", "wb");
       fwrite(out, sizeof(uint8_t), p.pitch[0]*STREAM_SIZE.height(), f);
       fclose(f);
       free(out);
       do_exit = true;
     }

    if (cuCtxSynchronize() != CUDA_SUCCESS)
      fprintf(stderr, "cuCtxSynchronize failed after memcpy");

    if (cuGraphicsUnregisterResource(pResource) != CUDA_SUCCESS)
      fprintf(stderr, "cuGraphicsEGLUnRegisterResource failed");

    NvDestroyEGLImage(NULL, egl_image);

    g_renderer->render(dmabuf);

    frame_id++;
  }

  NvBufferDestroy(dmabuf);

  return NULL;
}

int main() {
    g_renderer = NvEglRenderer::createEglRenderer("renderer", 0, 0, 0, 0); 
    if (!g_renderer)
        ORIGINATE_ERROR("failed to create renderer");
    g_renderer->enableProfiling();

    UniqueObj<CameraProvider> cameraProvider(CameraProvider::create());
    ICameraProvider *iCameraProvider = interface_cast<ICameraProvider>(cameraProvider);
    if (!iCameraProvider)
        ORIGINATE_ERROR("failed to create camera provider");

    std::vector<CameraDevice*> devices;
    iCameraProvider->getCameraDevices(&devices);
    printf("found %d camera(s)\n", devices.size());
    if (devices.size() == 0)
        ORIGINATE_ERROR("not enough cameras");

    UniqueObj<CaptureSession> captureSession(iCameraProvider->createCaptureSession(devices[0]));
    ICaptureSession *iSession = interface_cast<ICaptureSession>(captureSession);
    if (!iSession)
        ORIGINATE_ERROR("failed to create capture session");

    UniqueObj<OutputStreamSettings> streamSettings(iSession->createOutputStreamSettings(STREAM_TYPE_EGL));
    IEGLOutputStreamSettings *iEGLStreamSettings = interface_cast<IEGLOutputStreamSettings>(streamSettings);
    if (!iEGLStreamSettings)
        ORIGINATE_ERROR("failed to create output stream settings");
    iEGLStreamSettings->setPixelFormat(PIXEL_FMT_YCbCr_420_888);
    iEGLStreamSettings->setEGLDisplay(g_renderer->getEGLDisplay());
    iEGLStreamSettings->setResolution(STREAM_SIZE);

    UniqueObj<OutputStream> stream(iSession->createOutputStream(streamSettings.get()));
    if (!stream)
        ORIGINATE_ERROR("failed to create stream");

    UniqueObj<Request> request(iSession->createRequest());
    IRequest *iRequest = interface_cast<IRequest>(request);
    if (!iRequest) 
        ORIGINATE_ERROR("failed to create request");
    iRequest->enableOutputStream(stream.get());

    ISourceSettings *iSourceSettings = interface_cast<ISourceSettings>(request);
    iSourceSettings->setFrameDurationRange(Range<uint64_t>(1e9/30));

    pthread_t main_thread_handle;
    int err = pthread_create(&main_thread_handle, NULL, render_thread, (void *)stream.get());
    assert(err == 0);

    while (!do_exit) 
        iSession->capture(request.get());

    do_exit = true;

    pthread_join(main_thread_handle, NULL);

    g_renderer->printProfilingStats();

    iSession->waitForIdle();
    stream.reset();
    cameraProvider.reset();

    delete g_renderer;
}