#include "rtsp_grabber.h"

#include <chrono>
#include <filesystem>
#include <iostream>
#include <sstream>

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

#define LOG_ERR(msg) std::cerr << "[ERR] " << msg << std::endl
#define LOG_INF(msg) std::cerr << "[INF] " << msg << std::endl

struct RtspFrameGrabberImpl {
  AVFormatContext *fmt_ctx = nullptr;
  AVCodecContext *dec_ctx = nullptr;
  int video_stream_index = -1;
  AVFrame *decoded_frame = nullptr;
  AVFrame *sws_frame = nullptr;
  SwsContext *sws_ctx = nullptr;
  AVCodecContext *jpeg_enc_ctx = nullptr;
  AVPacket *enc_pkt = nullptr;
  AVFrame *latest_frame_copy = nullptr;
  int width = 0;
  int height = 0;
  AVPixelFormat src_pix_fmt = AV_PIX_FMT_NONE;
  std::atomic<bool> stop_flag{false};
  std::atomic<bool> started{false};
  std::atomic<bool> ready{false};
  std::thread th;
  int reconnect_delay_ms = 1000;
};

static void free_impl(RtspFrameGrabberImpl *impl) {
  if (!impl)
    return;
  if (impl->th.joinable()) {
    impl->stop_flag = true;
    impl->th.join();
  }
  if (impl->enc_pkt)
    av_packet_free(&impl->enc_pkt);
  if (impl->jpeg_enc_ctx)
    avcodec_free_context(&impl->jpeg_enc_ctx);
  if (impl->latest_frame_copy)
    av_frame_free(&impl->latest_frame_copy);
  if (impl->sws_frame)
    av_frame_free(&impl->sws_frame);
  if (impl->decoded_frame)
    av_frame_free(&impl->decoded_frame);
  if (impl->sws_ctx)
    sws_freeContext(impl->sws_ctx);
  if (impl->dec_ctx)
    avcodec_free_context(&impl->dec_ctx);
  if (impl->fmt_ctx)
    avformat_close_input(&impl->fmt_ctx);
  delete impl;
}

static AVFrame *avframe_deep_copy(const AVFrame *src) {
  if (!src)
    return nullptr;
  AVFrame *dst = av_frame_alloc();
  dst->format = src->format;
  dst->width = src->width;
  dst->height = src->height;
  av_frame_get_buffer(dst, 0);
  av_frame_copy(dst, src);
  av_frame_copy_props(dst, src);
  return dst;
}
static bool init_jpeg_encoder(RtspFrameGrabberImpl *impl, int width, int height,
                              AVPixelFormat pix_fmt) {
  if (!impl || impl->jpeg_enc_ctx)
    return true;

  const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
  if (!codec)
    return false;

  AVCodecContext *c = avcodec_alloc_context3(codec);

  c->pix_fmt = pix_fmt;
  c->width = width;
  c->height = height;
  c->time_base = AVRational{1, 25};

  // ⭐ 最高质量
  c->qmin = 1;
  c->qmax = 2;

  c->compression_level = 0;

  if (avcodec_open2(c, codec, nullptr) < 0) {
    avcodec_free_context(&c);
    return false;
  }

  impl->jpeg_enc_ctx = c;
  impl->enc_pkt = av_packet_alloc();

  return true;
}

static bool encode_frame_to_jpeg(RtspFrameGrabberImpl *impl, AVFrame *frame,
                                 const std::string &filename) {
  if (!impl || !impl->jpeg_enc_ctx || !frame)
    return false;
  AVCodecContext *enc = impl->jpeg_enc_ctx;
  AVPacket *pkt = impl->enc_pkt;
  av_packet_unref(pkt);
  if (avcodec_send_frame(enc, frame) < 0)
    return false;
  if (avcodec_receive_packet(enc, pkt) < 0)
    return false;
  FILE *f = fopen(filename.c_str(), "wb");
  if (!f)
    return false;
  fwrite(pkt->data, 1, pkt->size, f);
  fclose(f);
  av_packet_unref(pkt);
  return true;
}

// -------- RtspFrameGrabber Methods --------
RtspFrameGrabber::RtspFrameGrabber(const std::string &rtsp_url)
    : rtsp_url_(rtsp_url) {
  avformat_network_init();
  opaque_impl_ = new RtspFrameGrabberImpl();
}

RtspFrameGrabber::~RtspFrameGrabber() {
  stop();
  free_impl(reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_));
  opaque_impl_ = nullptr;
}

bool RtspFrameGrabber::start() {
  RtspFrameGrabberImpl *impl =
      reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_);
  if (!impl || impl->started)
    return false;
  impl->stop_flag = false;
  impl->th = std::thread([this, impl]() { this->decode_loop(); });
  impl->started = true;
  return true;
}

void RtspFrameGrabber::stop() {
  RtspFrameGrabberImpl *impl =
      reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_);
  if (!impl)
    return;
  impl->stop_flag = true;
  if (impl->th.joinable())
    impl->th.join();
  impl->started = false;
  impl->ready = false;
}

bool RtspFrameGrabber::is_ready() const {
  RtspFrameGrabberImpl *impl =
      reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_);
  return impl && impl->ready;
}

bool RtspFrameGrabber::capture_jpeg(const std::string &filename) {
  RtspFrameGrabberImpl *impl =
      reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_);
  if (!impl)
    return false;
  std::lock_guard<std::mutex> lk(frame_mutex_);
  if (!impl->latest_frame_copy)
    return false;
  int w = impl->latest_frame_copy->width;
  int h = impl->latest_frame_copy->height;
  AVPixelFormat pf = (AVPixelFormat)impl->latest_frame_copy->format;
  if (!init_jpeg_encoder(impl, w, h, pf))
    return false;
  return encode_frame_to_jpeg(impl, impl->latest_frame_copy, filename);
}

void RtspFrameGrabber::decode_loop() {
  RtspFrameGrabberImpl *impl =
      reinterpret_cast<RtspFrameGrabberImpl *>(opaque_impl_);
  if (!impl)
    return;

  AVDictionary *opts = nullptr;

  av_dict_set(&opts, "rtsp_transport", "tcp", 0);
  av_dict_set(&opts, "stimeout", "2000000", 0);

  av_dict_set(&opts, "fflags", "nobuffer", 0);
  av_dict_set(&opts, "flags", "low_delay", 0);
  av_dict_set(&opts, "max_delay", "0", 0);
  av_dict_set(&opts, "buffer_size", "102400", 0);

  AVPixelFormat dst_pix_fmt = AV_PIX_FMT_YUVJ420P;
#if LIBAVUTIL_VERSION_MAJOR >= 57
  dst_pix_fmt = AV_PIX_FMT_YUV420P;
#endif

  while (!impl->stop_flag) {
    AVFormatContext *fmt_ctx = nullptr;
    int ret = avformat_open_input(&fmt_ctx, rtsp_url_.c_str(), nullptr, &opts);
    if (ret != 0) {
      std::this_thread::sleep_for(
          std::chrono::milliseconds(impl->reconnect_delay_ms));
      continue;
    }
    impl->fmt_ctx = fmt_ctx;
    impl->reconnect_delay_ms = 1000;

    if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
      avformat_close_input(&impl->fmt_ctx);
      continue;
    }

    int video_index = -1;
    for (unsigned i = 0; i < fmt_ctx->nb_streams; ++i) {
      if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
        video_index = i;
        break;
      }
    }
    if (video_index < 0) {
      avformat_close_input(&impl->fmt_ctx);
      continue;
    }
    impl->video_stream_index = video_index;

    AVCodecParameters *codecpar = fmt_ctx->streams[video_index]->codecpar;
    const AVCodec *dec = avcodec_find_decoder(codecpar->codec_id);
    if (!dec) {
      avformat_close_input(&impl->fmt_ctx);
      continue;
    }

    AVCodecContext *dec_ctx = avcodec_alloc_context3(dec);
    avcodec_parameters_to_context(dec_ctx, codecpar);
    if (avcodec_open2(dec_ctx, dec, nullptr) < 0) {
      avcodec_free_context(&dec_ctx);
      avformat_close_input(&impl->fmt_ctx);
      continue;
    }
    impl->dec_ctx = dec_ctx;

    impl->decoded_frame = av_frame_alloc();
    impl->width = dec_ctx->width;
    impl->height = dec_ctx->height;
    impl->src_pix_fmt = dec_ctx->pix_fmt;
    impl->sws_frame = av_frame_alloc();
    impl->sws_frame->format = dst_pix_fmt;
    impl->sws_frame->width = impl->width;
    impl->sws_frame->height = impl->height;
    av_frame_get_buffer(impl->sws_frame, 32);
    impl->sws_ctx = sws_getContext(impl->width, impl->height, impl->src_pix_fmt,
                                   impl->width, impl->height,
                                   (AVPixelFormat)impl->sws_frame->format,
                                   SWS_BICUBIC, nullptr, nullptr, nullptr);
    init_jpeg_encoder(impl, impl->width, impl->height,
                      (AVPixelFormat)impl->sws_frame->format);
    impl->ready = true;
    // LOG_INF("connected and decoding: " << rtsp_url_);

    AVPacket *pkt = av_packet_alloc();
    while (!impl->stop_flag) {
      int r = av_read_frame(fmt_ctx, pkt);
      if (r < 0) {
        av_packet_unref(pkt);
        impl->ready = false;
        break;
      }
      if (pkt->stream_index != impl->video_stream_index) {
        av_packet_unref(pkt);
        continue;
      }
      if (avcodec_send_packet(dec_ctx, pkt) < 0) {
        av_packet_unref(pkt);
        continue;
      }

      while (!impl->stop_flag) {
        r = avcodec_receive_frame(dec_ctx, impl->decoded_frame);
        if (r == AVERROR(EAGAIN) || r == AVERROR_EOF)
          break;
        sws_scale(impl->sws_ctx, impl->decoded_frame->data,
                  impl->decoded_frame->linesize, 0, impl->height,
                  impl->sws_frame->data, impl->sws_frame->linesize);
        std::lock_guard<std::mutex> lk(frame_mutex_);
        if (impl->latest_frame_copy)
          av_frame_free(&impl->latest_frame_copy);
        impl->latest_frame_copy = avframe_deep_copy(impl->sws_frame);
      }
      av_packet_unref(pkt);
    }
    av_packet_free(&pkt);
    impl->ready = false;
    if (impl->sws_ctx) {
      sws_freeContext(impl->sws_ctx);
      impl->sws_ctx = nullptr;
    }
    if (impl->sws_frame) {
      av_frame_free(&impl->sws_frame);
      impl->sws_frame = nullptr;
    }
    if (impl->decoded_frame) {
      av_frame_free(&impl->decoded_frame);
      impl->decoded_frame = nullptr;
    }
    if (impl->dec_ctx) {
      avcodec_free_context(&impl->dec_ctx);
      impl->dec_ctx = nullptr;
    }
    if (impl->fmt_ctx) {
      avformat_close_input(&impl->fmt_ctx);
      impl->fmt_ctx = nullptr;
    }

    if (!impl->stop_flag)
      std::this_thread::sleep_for(
          std::chrono::milliseconds(impl->reconnect_delay_ms));
  }
  impl->ready = false;
  LOG_INF("decode thread exit for: " << rtsp_url_);
}
RtspManager::~RtspManager() { shutdown(); }

void RtspManager::configure(
    const std::unordered_map<int, std::string> &side_sources) {

  std::lock_guard<std::mutex> lk(mgr_mutex_);

  // 删除不再需要的grabber
  for (auto it = side_grabbers_.begin(); it != side_grabbers_.end();) {
    auto f = side_sources.find(it->first);

    if (f == side_sources.end() || f->second != it->second->get_rtsp_url()) {
      try {
        it->second->stop();
      } catch (...) {
        // 忽略异常，避免影响主程序
      }
      it = side_grabbers_.erase(it);
    } else {
      ++it;
    }
  }

  // 新增grabber
  for (const auto &kv : side_sources) {

    int sid = kv.first;
    const std::string &url = kv.second;

    if (side_grabbers_.find(sid) != side_grabbers_.end())
      continue;

    try {

      auto g = std::make_unique<RtspFrameGrabber>(url);

      if (!g) {
        continue;
      }

      // start失败直接跳过
      if (!g->start()) {
        continue;
      }

      side_grabbers_.emplace(sid, std::move(g));

    } catch (...) {
      // 捕获所有异常，避免configure影响系统稳定
      continue;
    }
  }
}

void RtspManager::shutdown() {
  std::lock_guard<std::mutex> lk(mgr_mutex_);

  for (auto &kv : side_grabbers_) {
    kv.second->stop();
  }
  side_grabbers_.clear();
}
std::string RtspManager::capture_side_frame(int source_id) {
  if (!is_ready(source_id))
    return "";
  std::lock_guard<std::mutex> lk(mgr_mutex_);

  auto it = side_grabbers_.find(source_id);

  if (it == side_grabbers_.end())
    return "";

  std::filesystem::create_directories("./temp");

  auto now = std::chrono::system_clock::now();
  std::time_t t = std::chrono::system_clock::to_time_t(now);

  std::ostringstream oss;

  oss << "side_" << source_id << "_" << t << ".jpg";

  std::string name = oss.str();
  std::string path = "./temp/" + name;

  if (!it->second->capture_jpeg(path))
    return "";

  return name;
}
bool RtspManager::is_ready(int source_id) const {
  std::lock_guard<std::mutex> lk(mgr_mutex_);

  auto it = side_grabbers_.find(source_id);

  return it != side_grabbers_.end() && it->second->is_ready();
}