#include <SDL/SDL.h>
#include <SDL/SDL_thread.h>
#include <fstream>
#include <iostream>
#include <mutex>
#include <queue>
#include <thread>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswresample/swresample.h>
#include <libswscale/swscale.h>
}

#define MAX_AUDIO_FRAME_SIZE 192000
#define MAX_AUDIO_QUEUE_SIZE 250
#define MAX_VIDEO_QUEUE_SIZE 250
#define MAX_PICTURE_QUEUE_SIZE 1

class PacketQueue {
public:
  int size() {
    std::unique_lock<std::mutex> lock(m_mutex);
    return (int)m_queue.size();
  }

  void notify_all() { m_cond.notify_all(); }

  void push(AVPacket *pkt) {
    std::unique_lock<std::mutex> lock(m_mutex);
    m_queue.push(pkt);
    lock.unlock();
    m_cond.notify_all();
  }

  AVPacket *pop() {
    AVPacket *pkt = nullptr;
    std::unique_lock<std::mutex> lock(m_mutex);
    if (m_queue.empty()) {
      m_cond.wait(lock);
    }
    pkt = m_queue.front();
    m_queue.pop();
    return pkt;
  }

private:
  std::queue<AVPacket *> m_queue;
  std::condition_variable m_cond;
  std::mutex m_mutex;
};

struct VideoPicture {
  SDL_Overlay *bmp = nullptr;
  int width = 0;
  int height = 0;
};

class PictureQueue {
public:
  void notify_all() { m_cond.notify_all(); }

  void push(const VideoPicture &vp) {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (m_queue.size() >= MAX_PICTURE_QUEUE_SIZE) {
      m_cond.wait(lock);
    }

    m_queue.push(vp);
  }

  VideoPicture pop() {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (m_queue.empty()) {
      return VideoPicture();
    }

    VideoPicture vp = m_queue.front();
    m_queue.pop();
    lock.unlock();
    m_cond.notify_all();
    return vp;
  }

private:
  std::queue<VideoPicture> m_queue;
  std::condition_variable m_cond;
  std::mutex m_mutex;
};

struct VideoState {
  std::string filename;
  bool quit = false;

  AVFormatContext *format_ctx = nullptr;
  int video_stream_index = -1;
  int audio_stream_index = -1;

  AVStream *audio_stream = nullptr;
  AVCodecContext *audio_codec_ctx = nullptr;
  PacketQueue audio_pkt_queue;
  uint8_t audio_buf[MAX_AUDIO_FRAME_SIZE];
  int audio_buf_size = 0;
  int audio_buf_index = 0;

  AVStream *video_stream = nullptr;
  AVCodecContext *video_codec_ctx = nullptr;
  PacketQueue video_pkt_queue;

  SwsContext *sws_ctx = nullptr;
  SwrContext *swr_ctx = nullptr;

  PictureQueue picture_queue;

  std::shared_ptr<std::thread> demux_thread; // 解复用线程
  std::shared_ptr<std::thread> video_thread; // 视频解码线程
};

#define FF_REFRESH_EVENT (SDL_USEREVENT)
#define FF_QUIT_EVENT (SDL_USEREVENT + 1)

SDL_Surface *screen = nullptr;
std::mutex screen_mutex;

// 通过发送事件的方式让主线程更新视频画面
Uint32 sdl_refresh_timer_cb(Uint32 interval, void *param) {
  SDL_Event event;
  event.type = FF_REFRESH_EVENT;
  event.user.data1 = param;
  SDL_PushEvent(&event);
  return 0;
}

void schedule_refresh(VideoState *is, int delay) {
  SDL_AddTimer(delay, sdl_refresh_timer_cb, is);
}

// 将视频帧刷新到SDL窗口
void video_refresh_timer(void *userdata) {
  VideoState *is = (VideoState *)userdata;

  VideoPicture vp = is->picture_queue.pop();
  if (is->video_stream) {
    if (vp.bmp == nullptr) {
      schedule_refresh(is, 1);
    } else {
      schedule_refresh(is, 10);

      // 对于非方形像素 sample_aspect_ratio 的值不为1
      SDL_Rect rect;
      float aspect_ratio;
      int w, h, x, y;
      if (is->video_codec_ctx->sample_aspect_ratio.num == 0) {
        aspect_ratio = 0;
      } else {
        aspect_ratio = (float)av_q2d(is->video_codec_ctx->sample_aspect_ratio) *
                       (float)is->video_codec_ctx->width /
                       (float)is->video_codec_ctx->height;
      }
      if (aspect_ratio <= 0.0) {
        aspect_ratio = (float)is->video_codec_ctx->width /
                       (float)is->video_codec_ctx->height;
      }
      h = screen->h;
      w = ((int)rint(h * aspect_ratio));
      if (w > screen->w) {
        w = screen->w;
        h = ((int)rint(w / aspect_ratio));
      }
      x = (screen->w - w) / 2;
      y = (screen->h - h) / 2;

      rect.x = x;
      rect.y = y;
      rect.w = w;
      rect.h = h;

      std::unique_lock lock(screen_mutex);
      SDL_DisplayYUVOverlay(vp.bmp, &rect);
    }
  } else {
    schedule_refresh(is, 100);
  }

  if (vp.bmp) {
    SDL_FreeYUVOverlay(vp.bmp);
  }
}

int audio_decode_frame(VideoState *is, uint8_t *audioBuf, int bufSize) {
  int dataSize(-1);

  AVPacket *packet = is->audio_pkt_queue.pop();
  if (packet == nullptr) {
    return -1;
  }

  // packet 发送到解码器
  avcodec_send_packet(is->audio_codec_ctx, packet);

  // 获取解码后的帧，一个音频帧包含多个音频采样点
  AVFrame *frame = av_frame_alloc();
  while (avcodec_receive_frame(is->audio_codec_ctx, frame) == 0) {
    // 音频重采样，将音频数据转换为所需的格式（采样频率、声道数、样本位数）
    int out_nb_samples =
        swr_convert(is->swr_ctx, &audioBuf, frame->nb_samples,
                    (const uint8_t **)frame->data, frame->nb_samples);

    dataSize = av_samples_get_buffer_size(nullptr, 2, out_nb_samples,
                                          AV_SAMPLE_FMT_S16, 1);
  }

  av_frame_free(&frame);
  av_packet_free(&packet);
  return dataSize;
}

void audio_callback(void *userdata, Uint8 *stream, int len) {
  // 存放解码出音频数据
  static uint8_t audioBuf[MAX_AUDIO_FRAME_SIZE];
  static uint32_t audioBufSize = 0;
  static uint32_t audioBufIndex = 0;

  VideoState *is = (VideoState *)userdata;
  AVCodecContext *codecCtx = is->audio_codec_ctx;

  // 需要存放的字节数剩余
  int lenRemain = len;
  while (lenRemain > 0) {
    if (audioBufIndex >= audioBufSize) {
      // 解码出的音频数据为空（或已消耗完毕），则执行音频解码
      int audioSize = audio_decode_frame(is, audioBuf, sizeof(audioBuf));
      if (audioSize < 0) {
        // 如果出错，填充数据 0，即无声音
        audioBufSize = 1024;
        std::memset(audioBuf, 0, audioBufSize);
      } else {
        audioBufSize = audioSize;
      }

      audioBufIndex = 0;
    }

    // 需要拷贝的音频数据大小
    int cpSize = audioBufSize - audioBufIndex;
    if (lenRemain < cpSize) {
      cpSize = lenRemain;
    }

    // 拷贝数据，并更新缓冲区相关信息
    std::memcpy(stream, (uint8_t *)audioBuf + audioBufIndex, cpSize);
    audioBufIndex += cpSize;
    stream += cpSize;
    lenRemain -= cpSize;
  }
}

// 视频解码线程
void video_thread(VideoState *is) {
  AVFrame *frame = av_frame_alloc();

  while (!is->quit) {
    AVPacket *pkt = is->video_pkt_queue.pop();
    if (pkt == nullptr) {
      break;
    }

    // 解码视频帧
    int ret = avcodec_send_packet(is->video_codec_ctx, pkt);
    if (ret < 0) {
      std::cout << "error: avcodec_send_packet() failed" << std::endl;
      av_packet_free(&pkt);
      goto fail;
    }

    ret = avcodec_receive_frame(is->video_codec_ctx, frame);
    if (ret == 0) {
      VideoPicture vp;
      vp.width = is->video_codec_ctx->width;
      vp.height = is->video_codec_ctx->height;

      std::unique_lock lock(screen_mutex);

      // YV12 SDL_Overlay
      vp.bmp = SDL_CreateYUVOverlay(is->video_codec_ctx->width,
                                    is->video_codec_ctx->height,
                                    SDL_YV12_OVERLAY, screen);

      lock.unlock();

      SDL_LockYUVOverlay(vp.bmp);

      // YV12 和 YUV420P 格式区别，UV分量顺序不同
      // https://blog.csdn.net/dss875914213/article/details/120836765
      AVFrame frameDst;
      frameDst.data[0] = vp.bmp->pixels[0];
      frameDst.data[1] = vp.bmp->pixels[2];
      frameDst.data[2] = vp.bmp->pixels[1];
      frameDst.linesize[0] = vp.bmp->pitches[0];
      frameDst.linesize[1] = vp.bmp->pitches[2];
      frameDst.linesize[2] = vp.bmp->pitches[1];

      // 执行像素格式转换
      sws_scale(is->sws_ctx, frame->data, frame->linesize, 0,
                is->video_codec_ctx->height, frameDst.data, frameDst.linesize);

      SDL_UnlockYUVOverlay(vp.bmp);

      is->picture_queue.push(vp);
    }

    av_packet_free(&pkt);
  }

fail:
  SDL_Event event;
  event.type = FF_QUIT_EVENT;
  event.user.data1 = is;
  SDL_PushEvent(&event);
  av_frame_free(&frame);
}

bool stream_component_open(VideoState *is, int stream_index) {
  if (stream_index < 0 || stream_index >= (int)is->format_ctx->nb_streams) {
    return false;
  }

  // 获取解码器
  const AVCodec *codec = avcodec_find_decoder(
      is->format_ctx->streams[stream_index]->codecpar->codec_id);
  if (!codec) {
    std::cout << "error: Unsupported codec" << std::endl;
    goto fail;
  }

  // 创建解码器上下文
  AVCodecContext *codec_ctx = avcodec_alloc_context3(codec);
  if (!codec_ctx) {
    std::cout << "error: avcodec_alloc_context3() failed" << std::endl;
    goto fail;
  }

  int ret = avcodec_parameters_to_context(
      codec_ctx, is->format_ctx->streams[stream_index]->codecpar);
  if (ret < 0) {
    std::cout << "error: avcodec_parameters_to_context() failed" << std::endl;
    goto fail;
  }

  ret = avcodec_open2(codec_ctx, codec, nullptr);
  if (ret < 0) {
    std::cout << "error: avcodec_open2() failed" << std::endl;
    goto fail;
  }

  switch (codec_ctx->codec_type) {
  case AVMEDIA_TYPE_AUDIO:
    is->audio_stream = is->format_ctx->streams[stream_index];
    is->audio_codec_ctx = codec_ctx;

    AVChannelLayout outChLayout;
    AVChannelLayout inChLayout;
    av_channel_layout_default(&outChLayout, 2);
    av_channel_layout_default(&inChLayout, codec_ctx->ch_layout.nb_channels);

    // 创建音频重采样上下文
    ret = swr_alloc_set_opts2(
        &is->swr_ctx, &outChLayout, AV_SAMPLE_FMT_S16, codec_ctx->sample_rate,
        &inChLayout, codec_ctx->sample_fmt, codec_ctx->sample_rate, 0, nullptr);
    if (ret < 0) {
      goto fail;
    }

    swr_init(is->swr_ctx);

    SDL_AudioSpec wanted_spec, spec;
    wanted_spec.freq = codec_ctx->sample_rate;
    wanted_spec.format = AUDIO_S16SYS;
    wanted_spec.channels = 2;
    wanted_spec.silence = 0;
    wanted_spec.samples = 1024;
    wanted_spec.callback = audio_callback;
    wanted_spec.userdata = is;
    if (SDL_OpenAudio(&wanted_spec, &spec) < 0) {
      std::cout << "error: SDL_OpenAudio() failed, " << SDL_GetError()
                << std::endl;
      goto fail;
    }
    SDL_PauseAudio(0);
    break;

  case AVMEDIA_TYPE_VIDEO:
    is->video_stream = is->format_ctx->streams[stream_index];
    is->video_codec_ctx = codec_ctx;
    is->video_thread = std::make_shared<std::thread>(video_thread, is);
    is->sws_ctx =
        sws_getContext(is->video_codec_ctx->width, is->video_codec_ctx->height,
                       is->video_codec_ctx->pix_fmt, is->video_codec_ctx->width,
                       is->video_codec_ctx->height, AV_PIX_FMT_YUV420P,
                       SWS_BILINEAR, NULL, NULL, NULL);
    break;

  default:
    break;
  }

  return true;

fail:
  avcodec_free_context(&codec_ctx);
  return false;
}

void demux_thread(VideoState *is) {
  // 打开文件
  int ret = avformat_open_input(&is->format_ctx, is->filename.c_str(), nullptr,
                                nullptr);
  if (ret < 0) {
    std::cout << "error: avformat_open_input() failed" << std::endl;
    goto fail;
  }

  ret = avformat_find_stream_info(is->format_ctx, nullptr);
  if (ret < 0) {
    std::cout << "error: avformat_find_stream_info() failed" << std::endl;
    goto fail;
  }

  av_dump_format(is->format_ctx, 0, is->filename.c_str(), 0);

  // 查找视频流、音频流
  for (uint32_t i = 0; i < is->format_ctx->nb_streams; ++i) {
    if (is->format_ctx->streams[i]->codecpar->codec_type ==
            AVMEDIA_TYPE_VIDEO &&
        is->video_stream_index < 0) {
      is->video_stream_index = i;
    } else if (is->format_ctx->streams[i]->codecpar->codec_type ==
                   AVMEDIA_TYPE_AUDIO &&
               is->audio_stream_index < 0) {
      is->audio_stream_index = i;
    }
  }

  if (is->video_stream_index == -1) {
    std::cout << "error: Didn't find a video stream";
    goto fail;
  }

  if (is->audio_stream_index == -1) {
    std::cout << "error: Didn't find a audio stream";
    goto fail;
  }

  if (is->video_stream_index >= 0) {
    if (!stream_component_open(is, is->video_stream_index)) {
      std::cout << "error: video stream_component_open() failed";
      goto fail;
    }
  }

  if (is->audio_stream_index >= 0) {
    if (!stream_component_open(is, is->audio_stream_index)) {
      std::cout << "error: audio stream_component_open() failed";
      goto fail;
    }
  }

  AVPacket *packet = av_packet_alloc();

  // 循环读取 packet 放到音视频对应的队列中
  while (true) {
    if (is->quit) {
      break;
    }

    if (is->audio_pkt_queue.size() >= MAX_AUDIO_QUEUE_SIZE ||
        is->video_pkt_queue.size() >= MAX_VIDEO_QUEUE_SIZE) {
      SDL_Delay(10);
      continue;
    }

    if (av_read_frame(is->format_ctx, packet) < 0) {
      break;
    }

    if (packet->stream_index == is->video_stream_index) {
      is->video_pkt_queue.push(av_packet_clone(packet));
    } else if (packet->stream_index == is->audio_stream_index) {
      is->audio_pkt_queue.push(av_packet_clone(packet));
    }

    av_packet_unref(packet);
  }

fail:
  SDL_Event event;
  event.type = FF_QUIT_EVENT;
  event.user.data1 = is;
  SDL_PushEvent(&event);
  av_packet_free(&packet);
}

// Spawning Threads
int main(int argc, char *argv[]) {
  VideoState is;

  // 取命令行参数获取文件名
  std::string filename;
  if (argc > 1) {
    filename = argv[1];
  } else {
    filename = "C:/Users/zhou/Videos/bigbuckbunny.mp4";
  }

  // 初始化SDL
  if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_AUDIO | SDL_INIT_TIMER) < 0) {
    std::cout << "error: Could not initialize SDL - " << SDL_GetError()
              << std::endl;
    goto fail;
  }

  // 创建 SDL_Surface
  screen = SDL_SetVideoMode(800, 600, 0, 0);
  if (!screen) {
    std::cout << "error: SDL: could not set video mode" << std::endl;
    goto fail;
  }

  is.filename = filename;

  // 启动解复用线程
  is.demux_thread = std::make_shared<std::thread>(demux_thread, &is);

  schedule_refresh(&is, 40);

  while (true) {
    SDL_Event event;
    SDL_WaitEvent(&event);

    switch (event.type) {
    case FF_QUIT_EVENT:
    case SDL_QUIT:
      is.quit = true;
      break;

    case FF_REFRESH_EVENT:
      // 将视频帧刷新到SDL窗口
      video_refresh_timer(event.user.data1);
      break;

    default:
      break;
    }

    if (is.quit) {
      break;
    }
  }

fail:
  SDL_FreeSurface(screen);
  SDL_Quit();

  avcodec_free_context(&is.audio_codec_ctx);
  avcodec_free_context(&is.video_codec_ctx);
  sws_freeContext(is.sws_ctx);
  swr_free(&is.swr_ctx);
  avformat_free_context(is.format_ctx);

  is.audio_pkt_queue.notify_all();
  is.video_pkt_queue.notify_all();
  is.picture_queue.notify_all();
  is.video_thread->join();
  is.demux_thread->join();

  return 0;
}