#include "ka_scodec.h"

#include <string.h>
#include "ka_mem.h"
#include "ka_log.h"
#include "config.h"
#include "ka_error.h"
#include "dvd/dvd_commons.h"

#define BUF_SIZE 0x10000

typedef struct
{
  ka_scodec_t hdr;
  uint8_t  buf[BUF_SIZE];
  int      filled;
  uint32_t last_pts;
  int      parsed;
  uint32_t section;
  uint8_t  colmap[4];
  uint8_t  alpha[4];
  uint32_t offset0;
  uint32_t offset1;
  int      mode;
  uint8_t  forceDisplay;
  uint8_t  nrColors;
  uint32_t curDate;
  uint32_t x0;
  uint32_t y0;
  uint32_t x1; // exclusive
  uint32_t y1; // exclusive
  ka_subcolor palette[256];
} data_t;

typedef struct
{
  const uint8_t* start;
  const uint8_t* end;
  uint32_t bits;
  uint32_t buf;
} bitbuffer;

static void initbits(bitbuffer* gb, const uint8_t* start, const uint8_t* end)
{
  gb->start = start + 4;
  gb->end = end;
  gb->bits = 32;
  gb->buf = (start[0] << 24) + (start[1] << 16) + (start[2] << 8) + start[3];;
}

static uint32_t getbits(bitbuffer* gb, int count)
{
  uint32_t val;

  if ((gb->bits < 24) && (gb->start < gb->end))
  {
    gb->buf |= (*gb->start++) << (24 - gb->bits);
    gb->bits += 8;
  }

  if (gb->bits >= count)
  {
    val = gb->buf >> (32 - count);
    gb->buf <<= count;
    gb->bits -= count;
  }
  else val = 0;

  return val;
}

static ka_scodec_t* subs_dvd_new(const ka_sparams_t* params)
{
  data_t* data = ka_mem_calloc(sizeof(*data));

  if (!data)
  {
    ka_error_fill(params->pErrorBlock, ka_error_nomem);
    return NULL;
  }

  data->hdr.vptr = &kav_scodec_dvd;
  data->hdr.vars = params->vars;
  data->hdr.pErrorBlock = params->pErrorBlock;

  data->filled = 0;

  return &data->hdr;
}

static void subs_dvd_delete(ka_scodec_t** ppcodec)
{
  data_t* data = (data_t*) *ppcodec;
  *ppcodec = NULL;

  if (!data)
    return;

  ka_mem_free(data);
}

static int decode_run_2bit(bitbuffer* gb, int* color)
{
  unsigned int v, t;

  v = 0;
  for (t = 1; v < t && t <= 0x40; t <<= 2)
    v = (v << 4) | getbits(gb, 4);
  *color = v & 3;

  if (v < 4)
  { // Code for fill rest of line
    return 1 << 16;
  }

  return v >> 2;
}

static int decode_run_8bit(bitbuffer* gb, int* color)
{
  int len;
  int has_run = getbits(gb, 1);

  if (getbits(gb, 1))
    *color = getbits(gb, 8);
  else
    *color = getbits(gb, 2);

  if (has_run)
  {
    if (getbits(gb, 1))
    {
      len = getbits(gb, 7);
      if (len == 0)
        len = 1 << 16;
      else
        len += 9;
    }
    else
      len = getbits(gb, 3) + 2;
  }
  else
    len = 1;

  return len;
}

static int decode_rle(uint8_t* p, int linesize, int w, int h
               , const uint8_t* start, const uint8_t* end, int is_8bit)
{
  bitbuffer gb;
  int x, y, len, color;

  if (start >= end)
    return -1;

  if (w <= 0 || h <= 0)
    return -1;

  initbits(&gb, start, end);

  x = 0;
  y = 0;
  while(y < h)
  {
    if (is_8bit)
      len = decode_run_8bit(&gb, &color);
    else
      len = decode_run_2bit(&gb, &color);
    if (len > w - x)
      len = w - x;
    memset(p + x, color, len);
    x += len;
    if (x >= w)
    {
      y++;
      if (y >= h)
        break;
      p += linesize;
      x = 0;
      // byte align
      if (gb.bits & 7)
        getbits(&gb, gb.bits & 7);
    }
  }

  return 0;
}

static inline uint32_t read_offset(const uint8_t* buf, int bigoffsets)
{
  uint32_t a = *buf;
  a = buf[1] + (a << 8);
  if (bigoffsets)
  {
    a = buf[2] + (a << 8);
    a = buf[3] + (a << 8);
  }

  return a;
}

static int defineRect(data_t* data, ka_sub_rect* rect, int typ, const uint8_t* buf, int size)
{
  rect->typ = typ;
  rect->section = data->section;
  rect->startDate = rect->endDate = data->curDate;

  if ((data->offset0 > 0)
  &&  (data->offset0 < size)
  &&  (data->offset1 > 0)
  &&  (data->offset1 < size))
  {
    int w = data->x1 - data->x0;
    int h = data->y1 - data->y0;

    if ((w > 0) && (h > 0))
    {
      rect->data = ka_mem_alloc(w * h);
      if (rect->data == NULL)
      {
        ka_log(ka_log_error | ka_log_subtitle, "Cannot allocate rect of (%d x %d)", w, h);
        goto fail;
      }

      if (decode_rle(rect->data, w * 2, w, (h + 1) / 2
                   , buf + data->offset0, buf + size, data->nrColors == 256) < 0)
      {
        ka_log(ka_log_error | ka_log_subtitle, "Cannot decode rle rect");
        goto fail;
      }
      if (decode_rle(rect->data + w, w * 2, w, h / 2
                   , buf + data->offset1, buf + size, data->nrColors == 256) < 0)
      {
        ka_log(ka_log_error | ka_log_subtitle, "Cannot decode rle rect");
        goto fail;
      }

      rect->bpr = w;

      return 0;
    }
    else
    {
      ka_log(ka_log_error | ka_log_subtitle, "Incorrect subtitle rect (%d x %d)", w, h);
    }
  }
  else
  {
    ka_log(ka_log_error | ka_log_subtitle
         , "Incorrect subtitle. [%010d], offset1 %d, offset2 %d , size %d"
         , data->curDate, data->offset0, data->offset1, size);
  }

fail:
  return -1;
}

static void updateRect(data_t* data, ka_sub_rect* rect)
{
  rect->startDate = rect->endDate = data->curDate;
  rect->forceDisplay = data->forceDisplay;
  rect->x0 = data->x0;
  rect->y0 = data->y0;
  rect->x1 = data->x1;
  rect->y1 = data->y1;
  rect->nrColors = data->nrColors;

  // define a default palette ?
  if (data->nrColors == 4)
  {
    const dvd_color_t* palette = ka_vars_get(data->hdr.vars, "DVD_PALETTE");
    if (palette != NULL)
    {
      rect->nrColors = 0;
      for (int i = 0; i < 4; i++)
      {
        dvd_color_t color = palette[data->colmap[i]];
        rect->palette[i].y = color.y;
        rect->palette[i].u = color.u;
        rect->palette[i].v = color.v;
        rect->palette[i].a = data->alpha[i];
        if (data->alpha[i])
          rect->nrColors++;
      }
    }
    else
    {
      const uint8_t y[5][4] = {{0x00}, {0xff}, {0xff, 0x00}, {0xff, 0x80, 0x00}, {0x00, 0x55, 0xaa, 0xff}};
      uint8_t used[16] = {0};
      rect->nrColors = 0;

      for (int i = 0; i < 4; i++)
      {
        rect->palette[i].y = 0;
        rect->palette[i].u = 0x80;
        rect->palette[i].v = 0x80;
        rect->palette[i].a = data->alpha[i];
        if (!used[data->colmap[i]] && data->alpha[i])
        {
          used[data->colmap[i]] = 1;
          rect->nrColors++;
        }
      }

      memset(used, 0xff, 16);
      for (int i = 0, j = 0; i < 4; i++)
      {
        if (data->alpha[i])
        {
          if (used[data->colmap[i]] == 0xff)
          {
            used[data->colmap[i]] = i;
            rect->palette[i].y = y[rect->nrColors][j];
            j++;
          }
          else
            rect->palette[i].y = rect->palette[used[data->colmap[i]]].y;
        }
      }
    }
  }
  else
  {
    memcpy(rect->palette, data->palette, sizeof(data->palette));
  }
}

static int subs_dvd_decode(ka_scodec_t* pcodec, ka_subtitles_t* subs, ka_block_t* b)
{
  data_t* data = (data_t*) pcodec;
  int size = b->eod - b->sod;
  ka_sub_rect* rect = NULL;

  // if pts, start of new data
  if (b->pts)
  {
    data->filled = 0;
    data->last_pts = b->pts;
    data->parsed = 0;
    data->section = b->section;
  }

  // merge in buffer
  if (BUF_SIZE < data->filled + size)
  {
    ka_log(ka_log_error | ka_log_subtitle, "DVD subtile too large for buffer");
    // extend size as if read
    data->filled += size;
    return 1; // consumed
  }

  memcpy(data->buf + data->filled, b->sod, size);
  data->filled += size;

  // decode
  if ((data->filled < 10) || data->parsed)
    return 1;

  const uint8_t* buf = data->buf;

  int bigoffsets = (!buf[0] && !buf[1]);
  size = read_offset(buf + (bigoffsets ? 2 : 0), bigoffsets);

  if (size > BUF_SIZE)
  {
    // should not happen, discard data
    data->buf[0] = data->buf[1] = data->buf[2] = data->buf[3] = 0;
    data->filled = 10;
    ka_log(ka_log_error | ka_log_subtitle, "Subtitle too large, discarded");
    return 1;
  }

  // wait till the buffer is filled to expected size
  if (size > data->filled)
    return 1;

//ka_log_dump(data->buf, size);

  uint32_t cmd_pos = read_offset(buf + (bigoffsets ? 6 : 2), bigoffsets);
  uint32_t pos = cmd_pos + (bigoffsets ? 6 : 4);
  uint8_t cmd;
  data->curDate = data->last_pts;
  data->colmap[0] = 0;
  data->colmap[1] = 1;
  data->colmap[2] = 2;
  data->colmap[3] = 3;
  data->alpha[0] = 0;
  data->alpha[1] = data->alpha[2] = data->alpha[3] = 0xff;
  data->offset0 = 0;
  data->offset1 = 0;
  data->nrColors = 4;
  data->x0 = 0;
  data->y0 = 0;
  data->x1 = 0;
  data->y1 = 0;

  data->parsed = 1;

  while (pos < size)
  {
    uint32_t date = read_offset(buf + cmd_pos, 0); // 2 bytes
    uint32_t next_cmd_pos = read_offset(buf + cmd_pos + 2, bigoffsets);
    if (next_cmd_pos == cmd_pos)
    {
      // last command
      next_cmd_pos = size;
    }
    else if ((next_cmd_pos > size) || (next_cmd_pos < pos))
    {
      ka_log(ka_log_error | ka_log_subtitle
           , "Subtitle command chain corrupted. Cmd pos %x -> %x, data size %x"
           , cmd_pos, next_cmd_pos, size);
      next_cmd_pos = size;
    }
    else if (data->curDate > data->last_pts + (date << 10))
    {
      ka_log(ka_log_error | ka_log_subtitle
           , "Subtitle command delay lower than previous one. Cmd pos %x", cmd_pos);
      next_cmd_pos = size;
    }
    data->curDate = data->last_pts + (date << 10);

    cmd = 0;
    data->mode = -1;

    // loop on commands
    while (pos < next_cmd_pos)
    {
      cmd = buf[pos++];

      switch(cmd)
      {
        case 0x00: // forced mode (like menus)
          data->mode = 0;
          data->forceDisplay = 1;
        break;
        case 0x01: // start
        {
          data->mode = 1;
          data->forceDisplay = 0;
        }
        break;
        case 0x02: // end
        {
          data->mode = 2;
        }
        break;
        case 0x03: // colormap
        {
          if (pos + 2 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode colormap");
            goto fail;
          }

          data->colmap[3] = buf[pos] >> 4;
          data->colmap[2] = buf[pos] & 0xf;
          data->colmap[1] = buf[pos+1] >> 4;
          data->colmap[0] = buf[pos+1] & 0xf;

          pos += 2;
        }
        break;
        case 0x04: // alpha
        {
          if (pos + 2 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode alpha");
            goto fail;
          }

          int a;
          a = (buf[pos] >> 4)   ;data->alpha[3] = a + (a << 4);
          a = (buf[pos] & 0xf)  ;data->alpha[2] = a + (a << 4);
          a = (buf[pos+1] >> 4) ;data->alpha[1] = a + (a << 4);
          a = (buf[pos+1] & 0xf);data->alpha[0] = a + (a << 4);
          pos += 2;
        }
        break;
        case 0x05: // rect
        case 0x85: // HD rect
        {
          if (pos + 6 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode rect");
            goto fail;
          }

          data->x0 = (buf[pos] << 4) + (buf[pos+1] >> 4);
          data->x1 = 1 + ((buf[pos+1] & 0xf) << 8) + buf[pos+2];
          data->y0 = (buf[pos+3] << 4) + (buf[pos+4] >> 4);
          data->y1 = 1 + ((buf[pos+4] & 0xf) << 8) + buf[pos+5];

          pos += 6;
        }
        break;
        case 0x06: // offsets
        {
          data->nrColors = 4;
          if (pos + 4 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode offsets");
            goto fail;
          }

          data->offset0 = (buf[pos] << 8) + buf[pos+1];
          data->offset1 = (buf[pos+2] << 8) + buf[pos+3];

          pos += 4;
        }
        break;
        case 0x83: // HD palette
        {
          if (pos + 768 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode palette");
            goto fail;
          }

          for (int i = 0; i < 256; i++)
          {
            data->palette[i].y = buf[pos + 3*i];
            data->palette[i].u = buf[pos + 3*i + 1];
            data->palette[i].v = buf[pos + 3*i + 2];
          }

          pos += 768;
        }
        break;
        case 0x84: // HD contrast
        {
          if (pos + 256 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode contrast");
            goto fail;
          }

          for (int i = 0; i < 256; i++)
            data->palette[i].a = 255 - buf[pos + i];

          pos += 256;
        }
        break;
        case 0x86: // HD offsets
        {
          data->nrColors = 256;
          if (pos + 8 > next_cmd_pos)
          {
            ka_log(ka_log_error | ka_log_subtitle, "Cannot decode HD offsets");
            goto fail;
          }

          data->offset0 = read_offset(buf + pos, 1);
          data->offset1 = read_offset(buf + pos + 4, 1);

          pos += 8;
        }
        break;
        case 0xff: // end
          pos = next_cmd_pos;
        break;
        default:
          ka_log(ka_log_error | ka_log_subtitle, "Unexpected DVD subtitle command %02x", cmd);
      }
    }

    if (cmd != 0xff)
      ka_log(ka_log_error | ka_log_subtitle, "DVD subtitle last command should be FF");

    if (data->mode >= 0)
    {
      if (rect != NULL)
      {
        // Existing rect ?
        if ((data->mode == 2) || (rect->startDate < data->curDate))
        {
          // Complete rect
          rect->endDate = data->curDate;
          ka_subtitles_rectStatus(subs, rect, 1);
          // Clear rect
          rect = NULL;
        }
        else
        {
          // Update definition
          updateRect(data, rect);
        }
      }

      // New rect needed ?
      if ((rect == NULL) && (data->mode < 2))
      {
        rect = ka_subtitles_getNewRect(subs);

        if (rect == NULL)
        {
          ka_log(ka_log_error | ka_log_subtitle, "No free subtitle rect");
          return 1;
        }

        if (defineRect(data, rect, b->typ, buf, size))
          goto fail;

        updateRect(data, rect);
      }
    }

    // next
    cmd_pos = next_cmd_pos;
    pos = cmd_pos + (bigoffsets ? 6 : 4);
  }

  if (rect != NULL)
  {
    // Complete rect
    rect->endDate = data->curDate;
    ka_subtitles_rectStatus(subs, rect, 1);
    return 1; // consumed
  }

fail:
  if (rect != NULL)
    ka_subtitles_rectStatus(subs, rect, 0);

  return 1; // consumed
}

const kav_scodec_t kav_scodec_dvd =
{
  .name = "DVD"
, .FNnew             = subs_dvd_new
, .FNdelete          = subs_dvd_delete
, .FNdecode          = subs_dvd_decode
};
