#include "h264_reader.h"

#include <string.h>

typedef struct {
  size_t start;
  size_t payload;
  size_t end;
  guint8 type;
  gboolean first_slice;
} H264Nal;

typedef struct {
  size_t start;
  size_t end;
  gboolean idr;
  gboolean has_sps;
  gboolean has_pps;
} H264Frame;

struct H264Reader {
  guint8 *file_data;
  size_t file_size;
  H264Frame *frames;
  size_t frame_count;
  size_t cursor;
  guint8 *sps;
  size_t sps_size;
  guint8 *pps;
  size_t pps_size;
  guint8 *send_buffer;
  size_t send_capacity;
};

static gboolean find_start_code(const guint8 *data, size_t size, size_t from,
                                size_t *offset, size_t *length)
{
  size_t i;

  for (i = from; i + 3 <= size; i++) {
    if (i + 4 <= size && data[i] == 0 && data[i + 1] == 0 &&
        data[i + 2] == 0 && data[i + 3] == 1) {
      *offset = i;
      *length = 4;
      return TRUE;
    }
    if (data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1) {
      *offset = i;
      *length = 3;
      return TRUE;
    }
  }
  return FALSE;
}

static gboolean first_mb_is_zero(const guint8 *data, size_t size)
{
  guint8 rbsp[128];
  size_t in = 0;
  size_t out = 0;
  unsigned zero_run = 0;
  size_t bit = 0;
  unsigned leading_zero_bits = 0;

  while (in < size && out < sizeof(rbsp)) {
    guint8 b = data[in++];
    if (zero_run >= 2 && b == 0x03) {
      zero_run = 0;
      continue;
    }
    rbsp[out++] = b;
    if (b == 0)
      zero_run++;
    else
      zero_run = 0;
  }

  while (bit < out * 8) {
    guint8 value = (rbsp[bit / 8] >> (7 - (bit % 8))) & 1;
    bit++;
    if (value)
      return leading_zero_bits == 0; /* ue(v) == 0 */
    leading_zero_bits++;
    if (leading_zero_bits > 31)
      break;
  }
  return FALSE;
}

static gboolean append_frame(H264Reader *reader, const H264Nal *nals,
                             size_t first, size_t past_last)
{
  H264Frame *frame;
  size_t i;

  if (past_last <= first)
    return TRUE;

  reader->frames = g_realloc_n(reader->frames, reader->frame_count + 1,
                               sizeof(*reader->frames));
  frame = &reader->frames[reader->frame_count++];
  memset(frame, 0, sizeof(*frame));
  frame->start = nals[first].start;
  frame->end = nals[past_last - 1].end;
  for (i = first; i < past_last; i++) {
    frame->idr |= nals[i].type == 5;
    frame->has_sps |= nals[i].type == 7;
    frame->has_pps |= nals[i].type == 8;
  }
  return TRUE;
}

static void cache_parameter_set(H264Reader *reader, const H264Nal *nal)
{
  guint8 **dst = nal->type == 7 ? &reader->sps : &reader->pps;
  size_t *dst_size = nal->type == 7 ? &reader->sps_size : &reader->pps_size;
  size_t size = nal->end - nal->start;

  if (*dst || size == 0)
    return;
  *dst = g_memdup2(reader->file_data + nal->start, size);
  *dst_size = size;
}

static gboolean parse_annexb(H264Reader *reader, char **error_message)
{
  H264Nal *nals = NULL;
  size_t nal_count = 0;
  size_t pos = 0;
  size_t start = 0;
  size_t sc_len = 0;
  size_t i;
  size_t au_first = 0;
  size_t pending_prefix = G_MAXSIZE;
  gboolean have_vcl = FALSE;
  gboolean have_idr = FALSE;

  while (find_start_code(reader->file_data, reader->file_size, pos,
                         &start, &sc_len)) {
    size_t next = 0;
    size_t next_sc_len = 0;
    size_t end;
    H264Nal *nal;

    if (start + sc_len >= reader->file_size)
      break;
    if (find_start_code(reader->file_data, reader->file_size,
                        start + sc_len, &next, &next_sc_len))
      end = next;
    else
      end = reader->file_size;

    nals = g_realloc_n(nals, nal_count + 1, sizeof(*nals));
    nal = &nals[nal_count++];
    nal->start = start;
    nal->payload = start + sc_len;
    nal->end = end;
    nal->type = reader->file_data[nal->payload] & 0x1f;
    nal->first_slice = FALSE;
    if (nal->type >= 1 && nal->type <= 5 && nal->payload + 1 < nal->end)
      nal->first_slice = first_mb_is_zero(reader->file_data + nal->payload + 1,
                                          nal->end - nal->payload - 1);
    if (nal->type == 7 || nal->type == 8)
      cache_parameter_set(reader, nal);

    pos = end;
    if (end >= reader->file_size)
      break;
  }

  if (nal_count == 0) {
    *error_message = g_strdup("not an Annex-B H.264 file (no start code)");
    g_free(nals);
    return FALSE;
  }

  for (i = 0; i < nal_count; i++) {
    gboolean vcl = nals[i].type >= 1 && nals[i].type <= 5;
    gboolean boundary = have_vcl &&
        (nals[i].type == 9 || (vcl && nals[i].first_slice));

    if (boundary) {
      size_t next_first = pending_prefix != G_MAXSIZE ? pending_prefix : i;
      append_frame(reader, nals, au_first, next_first);
      au_first = next_first;
      have_vcl = FALSE;
      pending_prefix = G_MAXSIZE;
    }

    if (vcl) {
      have_vcl = TRUE;
      pending_prefix = G_MAXSIZE;
    } else if (have_vcl && pending_prefix == G_MAXSIZE &&
               (nals[i].type == 6 || nals[i].type == 7 ||
                nals[i].type == 8 || nals[i].type == 9)) {
      pending_prefix = i;
    }
  }
  if (have_vcl)
    append_frame(reader, nals, au_first, nal_count);

  g_free(nals);

  for (i = 0; i < reader->frame_count; i++) {
    if (reader->frames[i].idr) {
      reader->cursor = i;
      have_idr = TRUE;
      break;
    }
  }

  if (!reader->sps || !reader->pps || !have_idr || reader->frame_count == 0) {
    *error_message = g_strdup_printf(
        "invalid H.264 stream: frames=%zu sps=%d pps=%d idr=%d",
        reader->frame_count, reader->sps != NULL, reader->pps != NULL, have_idr);
    return FALSE;
  }
  return TRUE;
}

H264Reader *h264_reader_open(const char *path, char **error_message)
{
  H264Reader *reader;
  gchar *contents = NULL;
  gsize length = 0;
  GError *error = NULL;

  if (error_message)
    *error_message = NULL;
  if (!path || !path[0]) {
    if (error_message)
      *error_message = g_strdup("empty H.264 file path");
    return NULL;
  }
  if (!g_file_get_contents(path, &contents, &length, &error)) {
    if (error_message)
      *error_message = g_strdup(error ? error->message : "read failed");
    g_clear_error(&error);
    return NULL;
  }

  reader = g_new0(H264Reader, 1);
  reader->file_data = (guint8 *)contents;
  reader->file_size = length;
  if (!parse_annexb(reader, error_message)) {
    h264_reader_close(reader);
    return NULL;
  }
  return reader;
}

void h264_reader_close(H264Reader *reader)
{
  if (!reader)
    return;
  g_free(reader->file_data);
  g_free(reader->frames);
  g_free(reader->sps);
  g_free(reader->pps);
  g_free(reader->send_buffer);
  g_free(reader);
}

static size_t find_next_idr(const H264Reader *reader, size_t from)
{
  size_t i;
  for (i = from; i < reader->frame_count; i++)
    if (reader->frames[i].idr)
      return i;
  for (i = 0; i < from && i < reader->frame_count; i++)
    if (reader->frames[i].idr)
      return i;
  return G_MAXSIZE;
}

gboolean h264_reader_next(H264Reader *reader, gboolean force_idr,
                          const uint8_t **data, size_t *size,
                          gboolean *is_idr)
{
  H264Frame *frame;
  size_t frame_size;

  if (!reader || !data || !size || reader->frame_count == 0)
    return FALSE;

  if (force_idr || reader->cursor >= reader->frame_count) {
    size_t next = find_next_idr(reader,
        reader->cursor < reader->frame_count ? reader->cursor : 0);
    if (next == G_MAXSIZE)
      return FALSE;
    reader->cursor = next;
  }

  frame = &reader->frames[reader->cursor++];
  frame_size = frame->end - frame->start;
  if (frame->idr && (!frame->has_sps || !frame->has_pps)) {
    size_t needed = reader->sps_size + reader->pps_size + frame_size;
    if (reader->send_capacity < needed) {
      reader->send_buffer = g_realloc(reader->send_buffer, needed);
      reader->send_capacity = needed;
    }
    memcpy(reader->send_buffer, reader->sps, reader->sps_size);
    memcpy(reader->send_buffer + reader->sps_size, reader->pps, reader->pps_size);
    memcpy(reader->send_buffer + reader->sps_size + reader->pps_size,
           reader->file_data + frame->start, frame_size);
    *data = reader->send_buffer;
    *size = needed;
  } else {
    *data = reader->file_data + frame->start;
    *size = frame_size;
  }
  if (is_idr)
    *is_idr = frame->idr;
  return TRUE;
}

size_t h264_reader_frame_count(const H264Reader *reader)
{
  return reader ? reader->frame_count : 0;
}
