#include "config.h"
#include "Imlib2_Loader.h"

#include <limits.h>
#include <stddef.h>
#include <stdint.h>

#include <libyuv.h>

#define DBG_PFX "LDR-y4m"

static const char  *const _formats[] = { "y4m" };

// Parser code taken from (commit 05ce1e4):
// https://codeberg.org/NRK/slashtmp/src/branch/master/parsers/y4m.c
//
// only parses the format, doesn't do any colorspace conversion.
// comments and frame parameters are ignored.
//
// This is free and unencumbered software released into the public domain.
// For more information, please refer to <https://unlicense.org/>
#define Y4M_PARSE_API static
#if IMLIB2_DEBUG
#define Y4M_PARSE_ASSERT(X)   do { if (!(X)) D("%d: %s\n", __LINE__, #X); } while (0)
#else
#define Y4M_PARSE_ASSERT(X)
#endif

enum Y4mParseStatus {
   Y4M_PARSE_OK,
   Y4M_PARSE_NOT_Y4M,
   Y4M_PARSE_EOF,
   Y4M_PARSE_CORRUPTED,
   Y4M_PARSE_UNSUPPORTED,
};

typedef struct {
   ptrdiff_t           w, h;
   ptrdiff_t           fps_num, fps_den;
   enum {
      Y4M_PARSE_CS_420,         /* default picked from ffmpeg */
      Y4M_PARSE_CS_420JPEG,
      Y4M_PARSE_CS_420MPEG2,
      Y4M_PARSE_CS_420PALDV,
      Y4M_PARSE_CS_422,
      Y4M_PARSE_CS_444,
      Y4M_PARSE_CS_MONO,
   } colour_space;
   enum {
      Y4M_PARSE_IL_PROGRESSIVE,
      Y4M_PARSE_IL_TOP,
      Y4M_PARSE_IL_BOTTOM,
      Y4M_PARSE_IL_MIXED,
   } interlacing;
   enum {
      Y4M_PARSE_ASPECT_DEFAULT,
      Y4M_PARSE_ASPECT_UNKNOWN,
      Y4M_PARSE_ASPECT_1_1,
      Y4M_PARSE_ASPECT_4_3,
      Y4M_PARSE_ASPECT_4_5,
      Y4M_PARSE_ASPECT_32_27,
      Y4M_PARSE_ASPECT_OTHER,
   } aspect;

   const void         *frame_data;
   ptrdiff_t           frame_data_len;
   const void         *y, *u, *v;
   ptrdiff_t           y_stride, u_stride, v_stride;

   // private
   const uint8_t      *p, *end;
} Y4mParse;

Y4M_PARSE_API enum Y4mParseStatus y4m_parse_init(Y4mParse * res,
                                                 const void *buffer,
                                                 ptrdiff_t size);
Y4M_PARSE_API enum Y4mParseStatus y4m_parse_frame(Y4mParse * res);

// implementation

static int
y4m__int(ptrdiff_t *out, const uint8_t ** p, const uint8_t * end)
{
   uint64_t            n = 0;
   const uint8_t      *s = *p;

   for (; s < end && *s >= '0' && *s <= '9'; ++s)
     {
        n = (n * 10) + (*s - '0');
        if (n > INT_MAX)
          {
             return 0;
          }
     }
   *out = n;
   *p = s;
   return 1;
}

static int
y4m__match(const char *match, ptrdiff_t mlen,
           const uint8_t ** p, const uint8_t * end)
{
   if (end - *p >= mlen && memcmp(match, *p, mlen) == 0)
     {
        *p += mlen;
        return 1;
     }
   return 0;
}

static enum Y4mParseStatus
y4m__parse_params(Y4mParse * res, const uint8_t ** start, const uint8_t * end)
{
   const uint8_t      *p = *start, *pp;

   for (;;)
     {
        if (end - p <= 0)
           return Y4M_PARSE_CORRUPTED;
        switch (*p++)
          {
          case ' ':
             break;
          case '\n':
             *start = p;
             if (res->w <= 0 || res->h <= 0 ||
                 res->fps_num <= 0 || res->fps_den <= 0)
                return Y4M_PARSE_CORRUPTED;
             return Y4M_PARSE_OK;
          case 'W':
             if (!y4m__int(&res->w, &p, end))
                return Y4M_PARSE_CORRUPTED;
             break;
          case 'H':
             if (!y4m__int(&res->h, &p, end))
                return Y4M_PARSE_CORRUPTED;
             break;
          case 'F':
             pp = p;
             if (!y4m__int(&res->fps_num, &p, end) ||
                 !y4m__match(":", 1, &p, end) ||
                 !y4m__int(&res->fps_den, &p, end))
             {
#if IMLIB2_DEBUG
               char str[1024];
               sscanf((const char *)(pp-1), "%s", str);
               D("%s: unknown frame rate: '%s'\n", __func__, str);
#endif
                return Y4M_PARSE_CORRUPTED;
             }
             break;
          case 'I':
             if (y4m__match("p", 1, &p, end))
                res->interlacing = Y4M_PARSE_IL_PROGRESSIVE;
             else if (y4m__match("t", 1, &p, end))
                res->interlacing = Y4M_PARSE_IL_TOP;
             else if (y4m__match("b", 1, &p, end))
                res->interlacing = Y4M_PARSE_IL_BOTTOM;
             else if (y4m__match("m", 1, &p, end))
                res->interlacing = Y4M_PARSE_IL_MIXED;
             else {
#if IMLIB2_DEBUG
               char str[1024];
               sscanf((const char *)(p-1), "%s", str);
               D("%s: unknown interlace type: '%s'\n", __func__, str);
#endif
                return Y4M_PARSE_CORRUPTED;
             }
             break;
          case 'C':
             if (y4m__match("mono", 4, &p, end))
                res->colour_space = Y4M_PARSE_CS_MONO;
             else if (y4m__match("420jpeg", 7, &p, end))
                res->colour_space = Y4M_PARSE_CS_420JPEG;
             else if (y4m__match("420mpeg2", 8, &p, end))
                res->colour_space = Y4M_PARSE_CS_420MPEG2;
             else if (y4m__match("420paldv", 8, &p, end))
                res->colour_space = Y4M_PARSE_CS_420PALDV;
             else if (y4m__match("420 ", 4, &p, end))
                res->colour_space = Y4M_PARSE_CS_420;
             else if (y4m__match("422 ", 4, &p, end))
                res->colour_space = Y4M_PARSE_CS_422;
             else if (y4m__match("444 ", 4, &p, end))
                res->colour_space = Y4M_PARSE_CS_444;
             else {
#if IMLIB2_DEBUG
                char str[1024];
                sscanf((const char *)(p-1), "%s", str);
                D("%s: unknown color type: '%s'\n", __func__, str);
#endif
                return Y4M_PARSE_CORRUPTED;
             }
             break;
          case 'A':
             if (y4m__match("0:0", 3, &p, end))
                res->aspect = Y4M_PARSE_ASPECT_UNKNOWN;
             else if (y4m__match("1:1", 3, &p, end))
                res->aspect = Y4M_PARSE_ASPECT_1_1;
             else if (y4m__match("4:3", 3, &p, end))
                res->aspect = Y4M_PARSE_ASPECT_4_3;
             else if (y4m__match("4:5", 3, &p, end))
                res->aspect = Y4M_PARSE_ASPECT_4_5;
             else if (y4m__match("32:27", 5, &p, end))
                res->aspect = Y4M_PARSE_ASPECT_32_27;
             else
               {
                  res->aspect = Y4M_PARSE_ASPECT_OTHER;
                  for (; p < end && *p != ' ' && *p != '\n'; ++p)
                     ;
               }
             break;
          case 'X':            /* comments ignored */
             for (; p < end && *p != ' ' && *p != '\n'; ++p)
                ;
             break;
          default:
             return Y4M_PARSE_CORRUPTED;
             break;
          }
     }
   Y4M_PARSE_ASSERT(!"unreachable");
   return -1;                   // silence warning
}

Y4M_PARSE_API enum Y4mParseStatus
y4m_parse_init(Y4mParse * res, const void *buffer, ptrdiff_t size)
{
   const char          magic[10] = "YUV4MPEG2 ";

   memset(res, 0x0, sizeof(*res));
   res->w = res->h = -1;
   res->p = buffer;
   res->end = res->p + size;

   if (!y4m__match(magic, sizeof(magic), &res->p, res->end))
     {
        return Y4M_PARSE_NOT_Y4M;
     }
   return y4m__parse_params(res, &res->p, res->end);
}

Y4M_PARSE_API enum Y4mParseStatus
y4m_parse_frame(Y4mParse * res)
{
   ptrdiff_t           npixels, sdiv, voff;
   const uint8_t      *p = res->p, *end = res->end;

   Y4M_PARSE_ASSERT(p <= end);
   if (p == end)
     {
        return Y4M_PARSE_EOF;
     }
   if (!y4m__match("FRAME", 5, &p, end))
     {
        return Y4M_PARSE_CORRUPTED;
     }
   // NOTE: skip frame params, ffmpeg seems to do the same...
   for (; p < end && *p != '\n'; ++p)
      ;
   if (p == end)
      return Y4M_PARSE_CORRUPTED;
   ++p;                         /* skip '\n' */

   res->frame_data = p;
   npixels = res->w * res->h;
   switch (res->colour_space)
     {
     case Y4M_PARSE_CS_420JPEG:
     case Y4M_PARSE_CS_420MPEG2:
     case Y4M_PARSE_CS_420PALDV:
     case Y4M_PARSE_CS_420:
        res->frame_data_len = npixels * 3 / 2;
        sdiv = 2;
        voff = (npixels * 5) / 4;
        break;
     case Y4M_PARSE_CS_422:
        res->frame_data_len = npixels * 2;
        sdiv = 2;
        voff = (npixels * 3) / 2;
        break;
     case Y4M_PARSE_CS_444:
        res->frame_data_len = npixels * 3;
        sdiv = 1;
        voff = npixels * 2;
        break;
     case Y4M_PARSE_CS_MONO:
        res->frame_data_len = npixels;
        sdiv = voff = 0;        // silence bogus compiler warning
        break;
     default:
        return Y4M_PARSE_UNSUPPORTED;
        break;
     }
   if (end - p < res->frame_data_len)
     {
        return Y4M_PARSE_CORRUPTED;
     }

   res->y = p;
   res->y_stride = res->w;
   if (res->colour_space == Y4M_PARSE_CS_MONO)
     {
        res->u = res->v = NULL;
        res->u_stride = res->v_stride = 0;
     }
   else
     {
        res->u = p + npixels;
        res->v = p + voff;
        res->u_stride = res->v_stride = res->w / sdiv;
     }

   res->p = p + res->frame_data_len;    /* advance to next potential frame */
   Y4M_PARSE_ASSERT(res->p <= end);

   return Y4M_PARSE_OK;
}

// END Y4mParse library

/* wrapper for mono colour space to match the signature of other yuv conversion
 * routines. */
static int
conv_mono(const uint8_t * y, int y_stride, const uint8_t * u, int u_stride,
          const uint8_t * v, int v_stride, uint8_t * dst, int dst_stride,
          int width, int height)
{
   return I400ToARGB(y, y_stride, dst, dst_stride, width, height);
}

static int
_load(ImlibImage * im, int load_data)
{
   Y4mParse            y4m;
   int                 res, fcount, frame;
   ImlibImageFrame    *pf = NULL;
   int                 (*conv)(const uint8_t *, int, const uint8_t *, int,
                               const uint8_t *, int, uint8_t *, int, int, int);

   if (y4m_parse_init(&y4m, im->fi->fdata, im->fi->fsize) != Y4M_PARSE_OK)
      return LOAD_FAIL;

   frame = im->frame;
   if (frame > 0)
     {
        fcount = 0;
        // NOTE: this is fairly cheap since nothing is being decoded.
        for (Y4mParse tmp = y4m; y4m_parse_frame(&tmp) == Y4M_PARSE_OK;)
           ++fcount;
        if (fcount == 0)
           return LOAD_BADIMAGE;
        if (frame > fcount)
           return LOAD_BADFRAME;

        pf = __imlib_GetFrame(im);
        if (!pf)
           return LOAD_OOM;
        pf->frame_count = fcount;
        pf->loop_count = 0;     /* Loop forever */
        if (pf->frame_count > 1)
           pf->frame_flags |= FF_IMAGE_ANIMATED;
     }
   else
     {
        frame = 1;
     }

   for (int i = 0; i < frame; ++i)
     {
        if (y4m_parse_frame(&y4m) != Y4M_PARSE_OK)
           return LOAD_BADIMAGE;
     }

   if (!IMAGE_DIMENSIONS_OK(y4m.w, y4m.h))
      return LOAD_BADIMAGE;
   // no interlacing support
   if (y4m.interlacing != Y4M_PARSE_IL_PROGRESSIVE)
      return LOAD_BADIMAGE;

   im->w = y4m.w;
   im->h = y4m.h;
   im->has_alpha = 0;

   switch (y4m.colour_space)
     {
     case Y4M_PARSE_CS_MONO:
        conv = conv_mono;
        break;
     case Y4M_PARSE_CS_422:
        conv = I422ToARGB;
        break;
     case Y4M_PARSE_CS_444:
        conv = I444ToARGB;
        break;
     case Y4M_PARSE_CS_420JPEG:
     case Y4M_PARSE_CS_420MPEG2:
     case Y4M_PARSE_CS_420PALDV:
     case Y4M_PARSE_CS_420:
        conv = I420ToARGB;
        break;
     default:
        DL("colour_space: %d\n", y4m.colour_space);
        return LOAD_BADIMAGE;
        break;
     }

   if (pf)
     {
        pf->canvas_w = im->w;
        pf->canvas_h = im->h;
        pf->frame_delay = (1000 * y4m.fps_den) / y4m.fps_num;
     }

   if (!load_data)
      return LOAD_SUCCESS;

   if (!__imlib_AllocateData(im))
      return LOAD_OOM;

   res = conv(y4m.y, y4m.y_stride, y4m.u, y4m.u_stride, y4m.v, y4m.v_stride,
              (uint8_t *) im->data, im->w * 4, im->w, im->h);
   if (res != 0)
      return LOAD_BADIMAGE;

   if (im->lc)
      __imlib_LoadProgressRows(im, 0, im->h);

   return LOAD_SUCCESS;
}

IMLIB_LOADER(_formats, _load, NULL);