#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>

static int g_video_size = 0;
static int g_audio_size = 0;
static uint8_t* g_video_buf = NULL;
static uint8_t* g_audio_buf = NULL;
static uint8_t* g_pps_buf = NULL;
static uint8_t* g_sps_buf = NULL;
static const uint32_t nalu_start_4bytecode = 0x01000000;
static const uint32_t nalu_start_3bytecode = 0x010000;

typedef enum H264_NALU_TYPE {
  NALU_TYPE_SPS = 7,
  NALU_TYPE_PPS = 8,
  NALU_TYPE_IDR = 5,
  NALU_TYPE_NON_IDR = 1,
} H264_NALU_TYPE;

// 添加循环控制变量
static int g_loop_count = 0;
static int g_max_loops = -1; // -1表示无限循环
static int g_audio_pos = 0;  // 音频位置

// 添加读取状态结构
typedef struct {
  uint8_t* video_start;
  uint8_t* video_current;
  uint8_t* video_end;
  int audio_pos;
  int loop_count;
  int sps_size;
  int pps_size;
} ReaderState;

static ReaderState g_reader_state = {0};

// 添加函数设置循环次数
void reader_set_loop_count(int max_loops) {
  g_max_loops = max_loops;
}

int reader_get_max_loops() {
  return g_max_loops;
}

int reader_init(char *video_url) {
  FILE* video_fp = NULL;
  FILE* audio_fp = NULL;
  char audiofile[] = "alaw08m.wav";
  printf("open file %s \n", video_url);

  video_fp = fopen(video_url, "rb");
  if (video_fp == NULL) {
    printf("open file %s failed\n", video_url);
    return -1;
  }

  fseek(video_fp, 0, SEEK_END);
  g_video_size = ftell(video_fp);
  fseek(video_fp, 0, SEEK_SET);
  g_video_buf = (uint8_t*)calloc(1, g_video_size);
  fread(g_video_buf, 1, g_video_size, video_fp);
  fclose(video_fp);

  audio_fp = fopen(audiofile, "rb");
  if (audio_fp == NULL) {
    printf("open file %s failed\n", audiofile);
    return -1;
  }

  fseek(audio_fp, 0, SEEK_END);
  g_audio_size = ftell(audio_fp);
  fseek(audio_fp, 0, SEEK_SET);
  g_audio_buf = (uint8_t*)calloc(1, g_audio_size);
  fread(g_audio_buf, 1, g_audio_size, audio_fp);
  fclose(audio_fp);

  // 初始化读取状态
  g_reader_state.video_start = g_video_buf;
  g_reader_state.video_current = g_video_buf;
  g_reader_state.video_end = g_video_buf + g_video_size;
  g_reader_state.audio_pos = 0;
  g_reader_state.loop_count = 0;
  g_reader_state.sps_size = 0;
  g_reader_state.pps_size = 0;

  // 重置全局循环计数
  g_loop_count = 0;
  g_audio_pos = 0;

  return 0;
}

uint8_t* reader_h264_find_nalu(uint8_t* buf_start, uint8_t* buf_end) {
  uint8_t* p = buf_start;

  while ((p + 3) < buf_end) {
    if (memcmp(p, &nalu_start_4bytecode, 4) == 0) {
      return p;
    } else if (memcmp(p, &nalu_start_3bytecode, 3) == 0) {
      return p;
    }
    p++;
  }

  return buf_end;
}

uint8_t* reader_get_video_frame(int* size) {
  uint8_t* buf = NULL;
  static int pps_size = 0;
  static int sps_size = 0;
  uint8_t* buf_end = g_video_buf + g_video_size;

  static uint8_t* pstart = NULL;
  static uint8_t* pend = NULL;
  size_t nalu_size;

  if (!pstart)
    pstart = g_video_buf;

  pend = reader_h264_find_nalu(pstart + 2, buf_end);

  if (pend == buf_end) {
    pstart = NULL;
    return NULL;
  }

  nalu_size = pend - pstart;
  int start_code_offset = memcmp(pstart, &nalu_start_3bytecode, 3) == 0 ? 3 : 4;
  H264_NALU_TYPE nalu_type = (H264_NALU_TYPE)(pstart[start_code_offset] & 0x1f);

  switch (nalu_type) {
    case NALU_TYPE_SPS:
      sps_size = nalu_size;
      if (g_sps_buf != NULL) {
        free(g_sps_buf);
        g_sps_buf = NULL;
      }
      g_sps_buf = (uint8_t*)calloc(1, sps_size);
      memcpy(g_sps_buf, pstart, sps_size);
      break;
    case NALU_TYPE_PPS:
      pps_size = nalu_size;
      if (g_pps_buf != NULL) {
        free(g_pps_buf);
        g_pps_buf = NULL;
      }
      g_pps_buf = (uint8_t*)calloc(1, pps_size);
      memcpy(g_pps_buf, pstart, pps_size);

      break;
    case NALU_TYPE_IDR:
      *size = sps_size + pps_size + nalu_size;
      buf = (uint8_t*)calloc(1, *size);
      memcpy(buf, g_sps_buf, sps_size);
      memcpy(buf + sps_size, g_pps_buf, pps_size);
      memcpy(buf + sps_size + pps_size, pstart, nalu_size);

      break;
    case NALU_TYPE_NON_IDR:
    default:
      *size = nalu_size;
      buf = (uint8_t*)calloc(1, *size);
      memcpy(buf, pstart, nalu_size);

      break;
  }

  pstart = pend;

  return buf;
}

uint8_t* reader_get_audio_frame(int* size) {
  // sample-rate=8000 channels=1 format=S16LE duration=20ms alaw-size=160
  uint8_t* buf = NULL;
  static int pos = 0;
  *size = 160;
  if ((pos + *size) > g_audio_size) {
    pos = 0;
  }

  buf = g_audio_buf + pos;
  pos += *size;

  return buf;
}

void reader_deinit() {
  if (g_sps_buf != NULL) {
    free(g_sps_buf);
    g_sps_buf = NULL;
  }

  if (g_pps_buf != NULL) {
    free(g_pps_buf);
    g_pps_buf = NULL;
  }

  if (g_video_buf != NULL) {
    free(g_video_buf);
    g_video_buf = NULL;
  }

  if (g_audio_buf != NULL) {
    free(g_audio_buf);
    g_audio_buf = NULL;
  }

  // 重置读取状态
  memset(&g_reader_state, 0, sizeof(g_reader_state));
  g_audio_pos = 0;
  g_loop_count = 0;
}

// 添加时间戳计算相关函数
uint32_t reader_get_video_timestamp(int frame_rate) {
  static uint32_t base_timestamp = 0;
  static uint32_t frame_count = 0;

  // 90000 Hz clock rate for video
  uint32_t timestamp_increment = 90000 / frame_rate;
  uint32_t timestamp = base_timestamp + frame_count * timestamp_increment;

  frame_count++;
  return timestamp;
}

uint32_t reader_get_audio_timestamp() {
  static uint32_t base_timestamp = 0;
  static uint32_t sample_count = 0;

  // 8000 Hz sample rate for audio
  uint32_t timestamp_increment = 160; // 20ms packets at 8000 Hz = 160 samples
  uint32_t timestamp = base_timestamp + sample_count;

  sample_count += timestamp_increment;
  return timestamp;
}