#include "ffmpeg_rga_decoder.h"

#include <opencv2/imgproc.hpp>
#include <opencv2/videoio.hpp>

#include "spdlog/spdlog.h"

// 结构体定义
struct RgaSrcInfo {
	int fd;
	int wstride;
	int hstride;
	void* vaddr;
	int format;
};

FFmpegRGADecoder::FFmpegRGADecoder() {
	pkt_ = av_packet_alloc();
	frame_ = av_frame_alloc();
}

FFmpegRGADecoder::~FFmpegRGADecoder() {
	release();
	if (pkt_)
		av_packet_free(&pkt_);
	if (frame_)
		av_frame_free(&frame_);
}

void FFmpegRGADecoder::cleanUp() {
	if (dec_ctx_) {
		avcodec_free_context(&dec_ctx_);
		dec_ctx_ = nullptr;
	}
	if (fmt_ctx_) {
		avformat_close_input(&fmt_ctx_);
		fmt_ctx_ = nullptr;
	}
	is_opened_ = false;
}

void FFmpegRGADecoder::release() {
	cleanUp();
	output_dma_buf_.reset();
}

bool FFmpegRGADecoder::isOpened() const {
	return is_opened_;
}

double FFmpegRGADecoder::get(int propId) {
	if (propId == cv::CAP_PROP_FRAME_WIDTH)
		return (double)width_;
	if (propId == cv::CAP_PROP_FRAME_HEIGHT)
		return (double)height_;
	if (propId == cv::CAP_PROP_FPS)
		return fps_;
	return 0.0;
}

bool FFmpegRGADecoder::open(const std::string& url) {
	cleanUp();
	int ret;
	AVDictionary* options = nullptr;

	// RTSP 优化参数
	av_dict_set(&options, "rtsp_transport", "tcp", 0);
	av_dict_set(&options, "buffer_size", "1024000", 0);
	av_dict_set(&options, "stimeout", "5000000", 0);
	av_dict_set(&options, "flags", "low_delay", 0);

	spdlog::info("FFmpeg: Opening stream: {}", url);
	if ((ret = avformat_open_input(&fmt_ctx_, url.c_str(), nullptr, &options)) < 0) {
		spdlog::error("FFmpeg: Could not open input URL: {}", url);
		return false;
	}

	if ((ret = avformat_find_stream_info(fmt_ctx_, nullptr)) < 0)
		return false;

	video_stream_idx_ = av_find_best_stream(fmt_ctx_, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
	if (video_stream_idx_ < 0)
		return false;

	AVStream* st = fmt_ctx_->streams[video_stream_idx_];
	AVCodecParameters* codecpar = st->codecpar;

	const char* decoder_name = nullptr;
	if (codecpar->codec_id == AV_CODEC_ID_HEVC)
		decoder_name = "hevc_rkmpp";
	else if (codecpar->codec_id == AV_CODEC_ID_H264)
		decoder_name = "h264_rkmpp";

	if (decoder_name) {
		decoder_ = avcodec_find_decoder_by_name(decoder_name);
		if (decoder_)
			spdlog::info("FFmpeg: Using Hardware Decoder: {}", decoder_name);
	}

	if (!decoder_) {
		spdlog::warn("FFmpeg: Hardware decoder not found, falling back.");
		decoder_ = avcodec_find_decoder(codecpar->codec_id);
	}
	if (!decoder_)
		return false;

	dec_ctx_ = avcodec_alloc_context3(decoder_);
	avcodec_parameters_to_context(dec_ctx_, codecpar);

	if ((ret = avcodec_open2(dec_ctx_, decoder_, nullptr)) < 0) {
		spdlog::error("FFmpeg: Failed to open codec");
		return false;
	}

	width_ = codecpar->width;
	height_ = codecpar->height;
	fps_ = (st->avg_frame_rate.den > 0) ? av_q2d(st->avg_frame_rate) : 25.0;

	is_opened_ = true;
	spdlog::info("FFmpeg: Decoder Ready. {}x{} @ {:.2f} fps", width_, height_, fps_);
	return true;
}

bool FFmpegRGADecoder::ensure_output_buffer(int width, int height) {
	size_t img_size = width * height * 3;
	if (!output_dma_buf_ || output_dma_buf_->size < img_size || width_ != width ||
		height_ != height) {
		width_ = width;
		height_ = height;
		output_dma_buf_ = std::make_unique<DmaBuffer>(img_size, "/dev/dma_heap/system-uncached");
		if (output_dma_buf_->fd < 0) {
			spdlog::error("FFmpeg: Failed to allocate output DMA buffer!");
			return false;
		}
		spdlog::info("FFmpeg: Allocated Output DMA Buffer (FD: {})", output_dma_buf_->fd);
	}
	return true;
}

// 获取源信息
bool get_src_info(const AVFrame* frame, RgaSrcInfo& info) {
	info.fd = 0;
	info.vaddr = nullptr;
	info.wstride = frame->width;
	info.hstride = frame->height;
	info.format = RK_FORMAT_YCbCr_420_SP;  // Default NV12

	if (frame->format == AV_PIX_FMT_DRM_PRIME && frame->data[0]) {
		AVDRMFrameDescriptor* desc = (AVDRMFrameDescriptor*)frame->data[0];
		info.fd = desc->objects[0].fd;

		// 1. 获取 Pitch (wstride)
		info.wstride = desc->layers[0].planes[0].pitch;

		// 2. 计算 hstride
		if (desc->layers[0].nb_planes >= 2) {
			int y_offset = desc->layers[0].planes[0].offset;
			int uv_offset = desc->layers[0].planes[1].offset;
			info.hstride = (uv_offset - y_offset) / info.wstride;
		} else {
			// [修改] 如果没有多平面信息，默认假设紧凑排列 (Compact)
			// 之前强制对齐到 16 可能导致读取位置错误（跳过了真正的 UV）
			// 如果画面变灰，通常是因为 hstride 算大了。
			info.hstride = frame->height;
		}
		return true;
	} else if (frame->format == AV_PIX_FMT_NV12) {
		info.wstride = frame->linesize[0];
		info.vaddr = (void*)frame->data[0];
		info.hstride = frame->height;
		return true;
	}
	// [新增] 检测 10bit 格式 (H.265 常见)
	else if (frame->format == AV_PIX_FMT_P010) {
		info.wstride = frame->linesize[0] / 2;	// P010 也是 2 bytes per pixel? 不，stride 是 bytes
		// RGA 对 P010 的处理需要特殊 flag，这里暂时回退到 NV12 处理尝试
		// 或者直接报错让软件处理
		info.wstride = frame->linesize[0];
		info.vaddr = (void*)frame->data[0];
		info.hstride = frame->height;
		info.format = RK_FORMAT_YCbCr_420_SP_10B;  // RGA 支持的 10bit 格式
		return true;
	}

	return false;
}

bool FFmpegRGADecoder::read(cv::Mat& output_mat) {
	if (!is_opened_)
		return false;

	static int frame_count = 0;
	frame_count++;

	int ret;
	while (true) {
		ret = av_read_frame(fmt_ctx_, pkt_);
		if (ret < 0)
			return false;

		if (pkt_->stream_index == video_stream_idx_) {
			ret = avcodec_send_packet(dec_ctx_, pkt_);
			if (ret < 0) {
				av_packet_unref(pkt_);
				continue;
			}

			ret = avcodec_receive_frame(dec_ctx_, frame_);
			if (ret == 0) {
				if (!ensure_output_buffer(frame_->width, frame_->height)) {
					av_packet_unref(pkt_);
					return false;
				}

				RgaSrcInfo src_info;
				bool has_info = get_src_info(frame_, src_info);

				// 调试日志：每 100 帧打印一次 Stride 信息，帮你确认计算是否正确
				if (frame_count % 100 == 0) {
					spdlog::info("Debug Frame: FD={} W={} H={} WStride={} HStride={} Fmt={}",
								 src_info.fd, frame_->width, frame_->height, src_info.wstride,
								 src_info.hstride, frame_->format);
				}

				if (has_info && src_info.fd > 0) {
					// --- 硬件路径 ---
					rga_buffer_t src_img, dst_img;
					rga_buffer_handle_t src_handle, dst_handle;

					// Import
					size_t src_size = src_info.wstride * src_info.hstride * 3 / 2;
					// 如果是 10bit，大小要翻倍? 不，P010 是 2 bytes per pixel
					if (src_info.format == RK_FORMAT_YCbCr_420_SP_10B)
						src_size *= 2;

					src_handle = importbuffer_fd(src_info.fd, src_size);
					dst_handle =
						importbuffer_fd(output_dma_buf_->fd, output_dma_buf_->actual_alloc_size);

					if (src_handle && dst_handle) {
						src_img = wrapbuffer_handle(src_handle, frame_->width, frame_->height,
													src_info.format);
						dst_img = wrapbuffer_handle(dst_handle, width_, height_, RK_FORMAT_BGR_888);

						// 设置 Stride
						src_img.wstride = src_info.wstride;
						src_img.hstride = src_info.hstride;
						dst_img.wstride = width_;
						dst_img.hstride = height_;

						imsetColorSpace(&src_img, IM_YUV_BT709_LIMIT_RANGE);
						imsetColorSpace(&dst_img, IM_RGB_FULL);

						IM_STATUS status =
							imcvtcolor(src_img, dst_img, src_info.format, RK_FORMAT_BGR_888);

						releasebuffer_handle(src_handle);
						releasebuffer_handle(dst_handle);

						if (status == IM_STATUS_SUCCESS) {
							output_mat = cv::Mat(height_, width_, CV_8UC3, output_dma_buf_->vaddr);
							av_packet_unref(pkt_);
							return true;
						} else {
							spdlog::error("RGA Fail: {}", imStrError(status));
						}
					}
				}

				// --- 软件回退 ---
				// 如果走到这里，说明硬件路径失败或无 FD
				// 简单的软件转换，确保颜色正常
				if (frame_->format == AV_PIX_FMT_NV12) {
					// 1. 将有效数据拷贝到连续内存 (去除 stride 影响)
					cv::Mat mYUV_Continous(frame_->height * 3 / 2, frame_->width, CV_8UC1);

					// 拷贝 Y
					uint8_t* src_y = frame_->data[0];
					uint8_t* dst_y = mYUV_Continous.data;
					for (int i = 0; i < frame_->height; i++) {
						memcpy(dst_y + i * frame_->width, src_y + i * frame_->linesize[0],
							   frame_->width);
					}

					// 拷贝 UV
					uint8_t* src_uv = frame_->data[1];
					// 如果 data[1] 为空，尝试根据 height 计算偏移
					if (!src_uv)
						src_uv = src_y + src_info.hstride * src_info.wstride;  // 尝试用 hstride

					uint8_t* dst_uv = dst_y + frame_->width * frame_->height;
					for (int i = 0; i < frame_->height / 2; i++) {
						memcpy(dst_uv + i * frame_->width, src_uv + i * frame_->linesize[0],
							   frame_->width);
					}

					cv::Mat mBGR(height_, width_, CV_8UC3, output_dma_buf_->vaddr);
					cv::cvtColor(mYUV_Continous, mBGR, cv::COLOR_YUV2BGR_NV12);

					if (width_ != frame_->width) {
						cv::resize(mBGR, mBGR, cv::Size(width_, height_));
					}

					output_mat = mBGR;
					av_packet_unref(pkt_);
					return true;
				}
			}
		}
		av_packet_unref(pkt_);
	}
}

bool FFmpegRGADecoder::read_raw(RgaFrameInfo& output_info) {
	if (!is_opened_)
		return false;
	int ret;
	while (true) {
		ret = av_read_frame(fmt_ctx_, pkt_);
		if (ret < 0)
			return false;

		if (pkt_->stream_index == video_stream_idx_) {
			ret = avcodec_send_packet(dec_ctx_, pkt_);
			if (ret < 0) {
				av_packet_unref(pkt_);
				continue;
			}

			ret = avcodec_receive_frame(dec_ctx_, frame_);
			if (ret == 0) {
				// 利用现有的逻辑获取信息
				RgaSrcInfo info;
				get_src_info(frame_, info);

				output_info.fd = info.fd;
				output_info.vaddr = info.vaddr;
				output_info.width = frame_->width;
				output_info.height = frame_->height;
				output_info.wstride = info.wstride;
				output_info.hstride = info.hstride;
				output_info.format = info.format;

				// 注意：这里我们返回了 info，但没有释放 AVPacket 或 AVFrame
				// 在单线程模型中，下一帧读取会覆盖 frame_。
				// 如果是 Zero-Copy，必须确保在 NPU 推理完成前，frame_ 内存不被释放/覆盖。
				// 但 rknn_run 是阻塞的，所以这里直接返回是安全的（前提是 infer 结束后才读下一帧）。

				av_packet_unref(pkt_);
				return true;
			}
		}
		av_packet_unref(pkt_);
	}
}