#include "ka_vcodec.h"

#include <stdio.h>
#include <string.h>
#include "ka_mem.h"
#include "ka_log.h"
#include "config.h"
#include "ka_error.h"

#define CHUNK_SIZE 0x40000 // 256KB

typedef struct
{
  uint8_t  y0, y1, y2, y3;
  int8_t   u, v;
  uint32_t rgb0, rgb1, rgb2, rgb3;
} codebook_t;

#define MAX_STRIPS      32

typedef struct
{
  uint32_t     id;
  uint32_t     x1, y1;
  uint32_t     x2, y2;
  codebook_t   v4_codebook[256];
  codebook_t   v1_codebook[256];
} strip_t;

typedef struct
{
  ka_vcodec_t hdr;
  ka_vcodec_event state;
  ka_vinfo_t   info;
  ka_vframe_t  frame[2];
  uint32_t     num_strips;
  uint32_t     frame_size;

  uint8_t*     chunk_ptr;
  ka_vframe_t* curr_frame;
  ka_vframe_t* prev_frame;
  int          frame_count;
  uint32_t     section;

  strip_t      strips[MAX_STRIPS];
  uint8_t      buffer[CHUNK_SIZE];
} data_t;

static uint32_t clip(int y)
{
  if (y > 255) return 255;
  if (y < 0) return 0;

  return (uint32_t) y;
}

static void cinepak_decode_codebook(codebook_t* codebook, int chunk_id
                                  , const uint8_t* sod, const uint8_t* eod)
{
  uint32_t flag, mask;
  int      i, n;
  int uvr, uvg, uvb;

  // check if this chunk contains 4- or 6-element vectors
  n    = (chunk_id & 0x0400) ? 4 : 6;
  flag = 0xffffffffu;
  mask = 0;

  for (i = 0; i < 256; i++)
  {
    if (!(mask >>= 1))
    {
      mask = 0x80000000;

      // inter coding uses 32-bit flag masks
      if (chunk_id & 0x0100)
      {
        if ((sod + 4) > eod)
          break;

        flag = (sod[0]<<24) + (sod[1]<<16) + (sod[2]<<8) + sod[3];
        sod += 4;
      }
    }

    if (flag & mask)
    {
      if ((sod + n) > eod)
        break;

      codebook->y0 = *sod++;
      codebook->y1 = *sod++;
      codebook->y2 = *sod++;
      codebook->y3 = *sod++;

      if (n == 6)
      {
        codebook->u = (int8_t) *sod++;
        codebook->v = (int8_t) *sod++;
      }
      else
      {
        codebook->u = 0;
        codebook->v = 0;
      }

      uvr = codebook->u << 1;
      uvg = ((codebook->v+1) >> 1) + codebook->u;
      uvb = codebook->v << 1;

      codebook->rgb0 = (clip(codebook->y0 + uvr) << KA_RGB_R_SHIFT)
                     | (clip(codebook->y0 - uvg) << KA_RGB_G_SHIFT)
                     | (clip(codebook->y0 + uvb) << KA_RGB_B_SHIFT);
      codebook->rgb1 = (clip(codebook->y1 + uvr) << KA_RGB_R_SHIFT)
                     | (clip(codebook->y1 - uvg) << KA_RGB_G_SHIFT)
                     | (clip(codebook->y1 + uvb) << KA_RGB_B_SHIFT);
      codebook->rgb2 = (clip(codebook->y2 + uvr) << KA_RGB_R_SHIFT)
                     | (clip(codebook->y2 - uvg) << KA_RGB_G_SHIFT)
                     | (clip(codebook->y2 + uvb) << KA_RGB_B_SHIFT);
      codebook->rgb3 = (clip(codebook->y3 + uvr) << KA_RGB_R_SHIFT)
                     | (clip(codebook->y3 - uvg) << KA_RGB_G_SHIFT)
                     | (clip(codebook->y3 + uvb) << KA_RGB_B_SHIFT);
    }
    codebook++;
  }
}

static int cinepak_decode_vectors(data_t* data, strip_t* strip, int chunk_id
                                , const uint8_t* sod, const uint8_t* eod)
{
  uint32_t         flag, mask;
  codebook_t*      codebook;
  uint32_t         x, y;
  uint32_t*        prgb;
  uint32_t dy  = data->info.luminance.width;

  flag = 0;
  mask = 0;

  for (y = strip->y1; y < strip->y2; y += 4)
  {
    prgb = (uint32_t*) data->curr_frame.planes.y;
    prgb += strip->x1 + (y * dy);

    for (x = strip->x1; x < strip->x2; x += 4)
    {
      if ((chunk_id & 0x0100) && !(mask >>= 1))
      {
        if ((sod + 4) > eod)
          return -1;

        flag = (sod[0]<<24) + (sod[1]<<16) + (sod[2]<<8) + sod[3];
        sod += 4;
        mask = 0x80000000;
      }

      if (!(chunk_id & 0x0100) || (flag & mask))
      {
        if (!(chunk_id & 0x0200) && !(mask >>= 1))
        {
          if ((sod + 4) > eod)
            return -1;

          flag = (sod[0]<<24) + (sod[1]<<16) + (sod[2]<<8) + sod[3];
          sod += 4;
          mask = 0x80000000;
        }

        if (!(chunk_id & 0x0200) && (flag & mask))
        {
          if ((sod + 4) > eod)
            return -1;

          codebook = &strip->v4_codebook[*sod++];
          prgb[0] = codebook->rgb0;
          prgb[1] = codebook->rgb1;
          prgb[dy] = codebook->rgb2;
          prgb[dy+1] = codebook->rgb3;

          codebook = &strip->v4_codebook[*sod++];
          prgb[2] = codebook->rgb0;
          prgb[3] = codebook->rgb1;
          prgb[dy+2] = codebook->rgb2;
          prgb[dy+3] = codebook->rgb3;

          codebook = &strip->v4_codebook[*sod++];
          prgb[2*dy] = codebook->rgb0;
          prgb[2*dy+1] = codebook->rgb1;
          prgb[3*dy] = codebook->rgb2;
          prgb[3*dy+1] = codebook->rgb3;

          codebook = &strip->v4_codebook[*sod++];
          prgb[2*dy+2] = codebook->rgb0;
          prgb[2*dy+3] = codebook->rgb1;
          prgb[3*dy+2] = codebook->rgb2;
          prgb[3*dy+3] = codebook->rgb3;
        }
        else
        {
          if (sod >= eod)
            return -1;

          codebook = &strip->v1_codebook[*sod++];
          prgb[0] = codebook->rgb0;
          prgb[1] = codebook->rgb0;
          prgb[dy] = codebook->rgb0;
          prgb[dy+1] = codebook->rgb0;

          prgb[2] = codebook->rgb1;
          prgb[3] = codebook->rgb1;
          prgb[dy+2] = codebook->rgb1;
          prgb[dy+3] = codebook->rgb1;

          prgb[2*dy] = codebook->rgb2;
          prgb[2*dy+1] = codebook->rgb2;
          prgb[3*dy] = codebook->rgb2;
          prgb[3*dy+1] = codebook->rgb2;

          prgb[2*dy+2] = codebook->rgb3;
          prgb[2*dy+3] = codebook->rgb3;
          prgb[3*dy+2] = codebook->rgb3;
          prgb[3*dy+3] = codebook->rgb3;
        }
      }

      prgb += 4;
    }
  }

  return 0;
}

static int cinepak_decode_strip(data_t* data, strip_t* strip
                              , const uint8_t* sod, const uint8_t* eod)
{
  int chunk_id, chunk_size;

  // coordinate sanity checks
  if (strip->x1 >= data->info.luminance.width
  ||  strip->x2 >  data->info.luminance.width
  ||  strip->y1 >= data->info.luminance.height
  ||  strip->y2 >  data->info.luminance.height
  ||  strip->x1 >= strip->x2
  ||  strip->y1 >= strip->y2)
  {
    if (config.debug & cfg_printwarnings)
      ka_log(ka_log_error, "invalid strip: [%d,%d] to ]%d,%d["
                         , strip->x1, strip->y1, strip->x2, strip->y2);
    return -1;
  }

  while ((sod + 4) <= eod)
  {
    chunk_id   = (sod[0]<<8) + sod[1];
    chunk_size = (sod[2]<<8) + sod[3] - 4;
    sod += 4;
    chunk_size = ((sod + chunk_size) > eod) ? (eod - sod) : chunk_size;

    switch (chunk_id)
    {
      case 0x2000:
      case 0x2100:
      case 0x2400:
      case 0x2500:
        cinepak_decode_codebook(strip->v4_codebook, chunk_id, sod, sod + chunk_size);
      break;
      case 0x2200:
      case 0x2300:
      case 0x2600:
      case 0x2700:
        cinepak_decode_codebook(strip->v1_codebook, chunk_id, sod, sod + chunk_size);
      break;
      case 0x3000:
      case 0x3100:
      case 0x3200:
      {
        if (cinepak_decode_vectors(data, strip, chunk_id, sod, sod + chunk_size))
          if (config.debug & cfg_printwarnings)
            ka_log(ka_log_error, "vector truncated");
        return 0;
      }
      break;
    }

    sod += chunk_size;
  }

  return -1;
}

static ka_vframe_type cinepak_frametype(const uint8_t* sod, const uint8_t* eod)
{
  if ((sod + 22) > eod)
    return ka_vframe_type_P;

  // check ID of first strip
  if (sod[10] != 0x10)
    return ka_vframe_type_P;

  return ka_vframe_type_I;
}

static int cinepak_decode_frame(data_t* data, const uint8_t* sod, const uint8_t* eod)
{
  int i, result, strip_size, frame_flags;
  int y0 = 0;

  frame_flags = sod[0];
  sod += 10;

  for (i=0; i < data->num_strips; i++)
  {
    if ((sod + 12) > eod)
      return -1;

    data->strips[i].id = (sod[0]<<8) + sod[1];
    data->strips[i].y1 = y0;
    data->strips[i].x1 = 0;
    data->strips[i].y2 = y0 + (sod[8]<<8) + sod[9];
    data->strips[i].x2 = data->info.luminance.width;

    strip_size = (sod[2]<<8) + sod[3] - 12;
    sod += 12;
    if ((sod + strip_size) > eod)
      if (config.debug & cfg_printwarnings)
        ka_log(ka_log_error, "Strip truncated");
    strip_size = ((sod + strip_size) > eod) ? (eod - sod) : strip_size;

    if (!(frame_flags & 0x01))
    {
      strip_t* ostrip = &data->strips[i ? i-1 : data->num_strips-1];
      memcpy(data->strips[i].v4_codebook, ostrip->v4_codebook,
             sizeof(ostrip->v4_codebook));
      memcpy(data->strips[i].v1_codebook, ostrip->v1_codebook,
             sizeof(ostrip->v1_codebook));
    }

    result = cinepak_decode_strip(data, &data->strips[i], sod, sod + strip_size);

    if (result != 0)
      return result;

    sod += strip_size;
    y0    = data->strips[i].y2;
  }

  return 0;
}

static void cvid_delete(ka_vcodec_t** ppcodec)
{
  data_t* data = (data_t*) *ppcodec;
  *ppcodec = NULL;

  if (!data)
    return;

  if (data->frame[0].planes.y) ka_mem_free(data->frame[0].planes.y);
  if (data->frame[1].planes.y) ka_mem_free(data->frame[1].planes.y);
  ka_mem_free(data);
}

static ka_vcodec_t* cvid_new(ka_error_t* pErrorBlock, const ka_vparams_t* params, uint32_t hardware)
{
  hardware = hardware;
  ka_vinfo_t* info;
  ka_vframe_t* frame;

  data_t* data = ka_mem_calloc(sizeof(*data));
  if (!data)
    goto error;

  data->hdr.vptr = &kav_vcodec_cvid;
  data->hdr.pErrorBlock = pErrorBlock;

  info = &data->info;
  info->videotype = "Cinepak";
  info->luminance.width = (params->width + 7) & ~7;
  info->luminance.height = (params->height + 7) & ~7;
  info->chroma.wshift = 1;
  info->chroma.hshift = 1;
  info->chroma.type = ka_chroma_RGB;
  info->picture.width = params->width;
  info->picture.height = params->height;
  info->frame_period = params->period;
  info->aspect_ratio.width = 1;
  info->aspect_ratio.height = 1;
  info->bitrate = 0;
  info->progressive = 1;

  frame = data->curr_frame = &data->frame[0];
  frame->type = ka_vframe_type_I;
  frame->pts = 0;
  frame->duration = info->frame_period;
  frame->display_type = 2;
  frame->valid = 1;
  frame->sequence_ref = 1;
  frame->planes.y = ka_mem_alloc(4*info->luminance.width
                                  *info->luminance.height);
  if (!frame->planes.y)
    goto error;

  frame = data->prev_frame = &data->frame[1];
  frame->type = ka_vframe_type_I;
  frame->pts = 0;
  frame->duration = info->frame_period;
  frame->display_type = 2;
  frame->valid = 1;
  frame->sequence_ref = 1;
  frame->planes.y = ka_mem_alloc(4*info->luminance.width
                                  *info->luminance.height);
  if (!frame->planes.y)
    goto error;

  data->state = ka_vcodec_event_buffering;
  data->chunk_ptr = data->buffer;
  data->frame_size = 0;
  data->num_strips = 0;

  return &data->hdr;

error:
  cvid_delete((ka_vcodec_t**) &data);

  ka_error_fill(pErrorBlock, ka_error_nomem);
  return NULL;
}

static void cvid_copyframe(ka_vframe_t* pto, const ka_vframe_t* pfrom, const ka_vinfo_t* info)
{
  // cf. all buffers allocated in one go
  memcpy(pto->planes.y, pfrom->planes.y, 4*info->luminance.width
                                          *info->luminance.height);
}

static ka_vcodec_event cvid_decode(ka_vcodec_t* pcodec, ka_block_t* b, uint32_t exit_time)
{
  data_t* data = (data_t*) pcodec;
  ka_vcodec_event event = data->state;

  exit_time = exit_time;     // unused

  if (!b) return ka_vcodec_event_buffering;

  if (data->frame_count > 1)
    return ka_vcodec_event_buffer_full;

  data->section = b->section;

  switch (event)
  {
    case ka_vcodec_event_buffering:
    {
      uint32_t size, width, height, strips;

      // Frame already started in previous block?
      if (data->chunk_ptr != data->buffer)
      {
        // Copy required data to end of previous frame
        size = data->frame_size - (data->chunk_ptr - data->buffer);
        if (size > (b->eod - b->sod))
          size = b->eod - b->sod;
        memcpy(data->chunk_ptr, b->sod, size);
        data->chunk_ptr += size;
        b->sod += size;

        // If enough data for previous frame, process it
        if (data->frame_size == (data->chunk_ptr - data->buffer))
        {
          ka_vframe_t* frame = data->prev_frame;

          data->state = ka_vcodec_event_decode_frame;
          data->prev_frame = data->curr_frame;
          data->curr_frame = frame;

          frame->valid = 1;
          frame->type = cinepak_frametype(data->buffer, data->chunk_ptr);
          if (frame->type != ka_vframe_type_I)
            cvid_copyframe(frame, data->prev_frame, &data->info);
          return ka_vcodec_event_decode_frame;
        }

        return ka_vcodec_event_buffering;
      }

      // Start of new frame?
      if ((b->eod - b->sod) < 10)
      {
        if (config.debug & cfg_printwarnings)
          ka_log(ka_log_error, "Short buffer");
        b->sod = b->eod;
        return ka_vcodec_event_buffering;
      }

      size   = (b->sod[1]<<16) + (b->sod[2]<<8) + b->sod[3];
      width  = (b->sod[4]<<8) + b->sod[5];
      height = (b->sod[6]<<8) + b->sod[7];
      strips = (b->sod[8]<<8) + b->sod[9];

      if ((strips > MAX_STRIPS)
      ||  (data->num_strips && (strips != data->num_strips))
      ||  (width != data->info.luminance.width)
      ||  (height != data->info.luminance.height))
      {
        if (config.debug & cfg_printwarnings)
        {
          ka_log(ka_log_error, "Invalid start of buffer, ignoring.");
          ka_log(ka_log_error, " Params: %dx%d, %d strips", width, height, strips);
          ka_log(ka_log_error, " Expected: %dx%d, %d strips", data->info.luminance.width, data->info.luminance.height, data->num_strips);
        }
        // Probably continuation of as provious frame, skipping
        b->sod = b->eod;
        return ka_vcodec_event_buffering;
      }

      // Save frame parameters
      if (!data->num_strips) data->num_strips = strips;
      data->frame_size = size;

      if (data->frame_size <= (b->eod - b->sod))
      {
        ka_vframe_t* frame = data->prev_frame;

        data->state = ka_vcodec_event_decode_frame;
        data->prev_frame = data->curr_frame;
        data->curr_frame = frame;

        frame->valid = 1;
        frame->type = cinepak_frametype(b->sod, b->sod + data->frame_size);
        if (frame->type != ka_vframe_type_I)
          cvid_copyframe(frame, data->prev_frame, &data->info);
        return ka_vcodec_event_decode_frame;
      }

      // Save end of block as start of frame
      memcpy(data->buffer, b->sod, b->eod - b->sod);
      data->chunk_ptr = data->buffer + (b->eod - b->sod);
      b->sod = b->eod;
      return ka_vcodec_event_buffering;
    }
    break;
    case ka_vcodec_event_decode_frame:
    {
      if (data->frame_size == (data->chunk_ptr - data->buffer))
      {
        if (data->curr_frame->valid > 0)
          cinepak_decode_frame(data, data->buffer, data->chunk_ptr);
        data->chunk_ptr = data->buffer;
      }
      else
      {
        if (data->curr_frame->valid > 0)
          cinepak_decode_frame(data, b->sod, b->sod + data->frame_size);
        b->sod += data->frame_size;
      }
      data->state = ka_vcodec_event_end_frame;
      data->frame_count++;
      return data->state;
    }
    break;
    case ka_vcodec_event_end_frame:
    {
      data->state = ka_vcodec_event_buffering;
      return data->state;
    }
    break;
  }

  return event;
}

static void cvid_action(ka_vcodec_t* pcodec, ka_vcodec_action action)
{
  data_t* data = (data_t*) pcodec;

  switch(action)
  {
    case ka_vcodec_action_reset:
    {
      data->state = ka_vcodec_event_buffering;
      data->chunk_ptr = data->buffer;
      data->frame_size = 0;
      data->frame_count = 0;
    }
    break;
    case ka_vcodec_action_skipframe:
    {
      data->curr_frame->valid = 0;
    }
    break;
  }
}

static void cvid_info(ka_vcodec_t* pcodec, ka_vinfo_t* info)
{
  data_t* data = (data_t*) pcodec;

  *info = data->info;
}

static int cvid_getDecodeFrame(ka_vcodec_t* pcodec, ka_vframe_t* frame)
{
  data_t* data = (data_t*) pcodec;

  *frame = *data->curr_frame;
  frame->section = data->section;

  return 1;
}

static int cvid_getDisplayFrame(ka_vcodec_t* pcodec, ka_vframe_t* frame)
{
  data_t* data = (data_t*) pcodec;

  *frame = *data->curr_frame;

  return 1;
}

static void cvid_releaseFrame(ka_vcodec_t* pcodec, ka_vframe_t* frame)
{
  data_t* data = (data_t*) pcodec;
  if (data->frame_count > 0)
    data->frame_count--;
  else
   ka_log(ka_log_error, "release already released frame %d", frame->index);
  frame = frame; // suppress not used warning
}

const kav_vcodec_t kav_vcodec_cvid =
{
  .name = "Radius Cinepak"
, .FNnew             = cvid_new
, .FNdelete          = cvid_delete
, .FNdecode          = cvid_decode
, .FNact             = cvid_action
, .FNgetInfo         = cvid_info
, .FNgetDecodeFrame  = cvid_getDecodeFrame
, .FNgetDisplayFrame = cvid_getDisplayFrame
, .FNreleaseFrame    = cvid_releaseFrame
};
