#include "videoplayer.hpp"

#ifdef ZHAVE_VIDEO

#define EGL_EGLEXT_PROTOTYPES
#define GL_GLEXT_PROTOTYPES

#include <EGL/egl.h>
#include <GL/gl.h>
#include <GL/glext.h>
#include <cuda.h>

#include <algorithm>
#include <chrono>
#include <climits>

#include <QtLogging>

extern "C" CUresult cuGraphicsGLRegisterImage(
	CUgraphicsResource* resource,
	GLuint image,
	GLenum target,
	unsigned int flags);

Q_LOGGING_CATEGORY(videoLog, "ZShell.Internal.wallpaper")

namespace {

enum AVPixelFormat chooseFormat(
	AVCodecContext* ctx, const AVPixelFormat* pix_fmts) {
	(void)ctx;
	for (const AVPixelFormat* p = pix_fmts; *p != AV_PIX_FMT_NONE; ++p)
		if (*p == AV_PIX_FMT_CUDA) return AV_PIX_FMT_CUDA;
	return pix_fmts[0];
}

} // namespace

VideoPlayer::VideoPlayer(EGLDisplay display) : m_eglDisplay(display) {}

VideoPlayer::~VideoPlayer() {
	stop();

	std::lock_guard<std::mutex> lock(m_mtx);
	if (m_pending) {
		av_frame_unref(m_pending);
		av_frame_free(&m_pending);
	}
	if (m_displayed) {
		av_frame_unref(m_displayed);
		av_frame_free(&m_displayed);
	}

	if (m_cuCtx) {
		CUcontext prev = nullptr;
		cuCtxPushCurrent(m_cuCtx);
		if (m_resY) {
			cuGraphicsUnregisterResource(m_resY);
			m_resY = nullptr;
		}
		if (m_resUV) {
			cuGraphicsUnregisterResource(m_resUV);
			m_resUV = nullptr;
		}
		if (m_texY) {
			glDeleteTextures(1, &m_texY);
			m_texY = 0;
		}
		if (m_texUV) {
			glDeleteTextures(1, &m_texUV);
			m_texUV = 0;
		}
		cuCtxPopCurrent(&prev);
		cuCtxDestroy(m_cuCtx);
		m_cuCtx = nullptr;
	}
	m_texSize = QSize();
}

void VideoPlayer::setReadyCallback(std::function<void()> callback) {
	m_readyCb = std::move(callback);
}

void VideoPlayer::setFrameCallback(std::function<void()> callback) {
	m_frameCb = std::move(callback);
}

bool VideoPlayer::openDecoder(const QString& path) {
	if (avformat_open_input(
			&m_fmt, path.toUtf8().constData(), nullptr, nullptr) < 0) {
		qCWarning(
			videoLog,
			"VideoPlayer: could not open %s",
			path.toUtf8().constData());
		return false;
	}
	if (avformat_find_stream_info(m_fmt, nullptr) < 0) {
		qCWarning(videoLog, "VideoPlayer: could not read stream info");
		return false;
	}

	m_streamIndex = -1;
	for (unsigned i = 0; i < m_fmt->nb_streams; ++i)
		if (m_fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
			m_streamIndex = int(i);
			break;
		}
	if (m_streamIndex < 0) {
		qCWarning(videoLog, "VideoPlayer: no video stream");
		return false;
	}

	AVCodecParameters* par = m_fmt->streams[m_streamIndex]->codecpar;
	const AVCodec* codec = avcodec_find_decoder(par->codec_id);
	const AVCodec* hwcodec = nullptr;
	void* it = nullptr;
	while (const AVCodec* c = av_codec_iterate(&it)) {
		if (c->id != par->codec_id || !av_codec_is_decoder(c)) continue;
		for (int i = 0;
			 const AVCodecHWConfig* cfg = avcodec_get_hw_config(c, i);
			 ++i)
			if (cfg->device_type == AV_HWDEVICE_TYPE_CUDA &&
				(cfg->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)) {
				hwcodec = c;
				break;
			}
		if (hwcodec) break;
	}
	if (hwcodec) codec = hwcodec;
	if (!codec) {
		qCWarning(
			videoLog,
			"VideoPlayer: no decoder for codec id %d",
			int(par->codec_id));
		return false;
	}

	m_codec = avcodec_alloc_context3(codec);
	if (!m_codec) return false;
	if (avcodec_parameters_to_context(m_codec, par) < 0) {
		qCWarning(videoLog, "VideoPlayer: could not copy codec params");
		return false;
	}

	m_hwDeviceCtx = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_CUDA);
	if (!m_hwDeviceCtx) {
		qCWarning(
			videoLog, "VideoPlayer: could not allocate cuda hwdevice context");
		return false;
	}
	AVCUDADeviceContext* cudaHw = reinterpret_cast<AVCUDADeviceContext*>(
		reinterpret_cast<AVHWDeviceContext*>(m_hwDeviceCtx->data)->hwctx);
	if (m_cuCtx) cudaHw->cuda_ctx = m_cuCtx;
	if (av_hwdevice_ctx_init(m_hwDeviceCtx) < 0) {
		qCWarning(
			videoLog, "VideoPlayer: could not create cuda hwdevice context");
		av_buffer_unref(&m_hwDeviceCtx);
		return false;
	}
	m_codec->hw_device_ctx = av_buffer_ref(m_hwDeviceCtx);
	m_codec->get_format = chooseFormat;

	if (avcodec_open2(m_codec, codec, nullptr) < 0) {
		qCWarning(videoLog, "VideoPlayer: could not open codec");
		return false;
	}

	return true;
}

void VideoPlayer::closeDecoder() {
	if (m_codec) {
		avcodec_free_context(&m_codec);
		m_codec = nullptr;
	}
	if (m_hwDeviceCtx) {
		av_buffer_unref(&m_hwDeviceCtx);
		m_hwDeviceCtx = nullptr;
	}
	if (m_fmt) {
		avformat_close_input(&m_fmt);
		m_fmt = nullptr;
	}
	m_streamIndex = -1;
}

void VideoPlayer::decodeLoop(const QString& path) {
	if (!openDecoder(path)) {
		m_threadRunning = false;
		return;
	}

	AVRational fr = m_fmt->streams[m_streamIndex]->avg_frame_rate;
	if (fr.num <= 0 || fr.den <= 0) fr = m_codec->framerate;
	if (fr.num <= 0 || fr.den <= 0) fr = AVRational{30, 1};
	const int64_t intervalUs =
		av_rescale_q(1, AVRational{fr.den, fr.num}, AVRational{1, 1000000});

	AVFrame* frame = av_frame_alloc();
	if (!frame) {
		m_threadRunning = false;
		return;
	}

	AVPacket pkt{};

	while (m_threadRunning) {
		int ret = av_read_frame(m_fmt, &pkt);
		if (ret == AVERROR_EOF) {
			if (avformat_seek_file(
					m_fmt, m_streamIndex, INT64_MIN, 0, INT64_MAX, 0) < 0) {
				qCWarning(videoLog, "VideoPlayer: loop seek failed; stopping");
				break;
			}
			avcodec_flush_buffers(m_codec);
			continue;
		}
		if (ret < 0) {
			qCWarning(videoLog, "VideoPlayer: read error; stopping");
			break;
		}
		if (pkt.stream_index != m_streamIndex) {
			av_packet_unref(&pkt);
			continue;
		}

		if (avcodec_send_packet(m_codec, &pkt) < 0) {
			av_packet_unref(&pkt);
			continue;
		}
		av_packet_unref(&pkt);

		while (m_threadRunning && avcodec_receive_frame(m_codec, frame) == 0) {
			bool wasEmpty = false;
			{
				std::lock_guard<std::mutex> lock(m_mtx);
				if (!m_pending) m_pending = av_frame_alloc();
				av_frame_unref(m_pending);
				if (av_frame_ref(m_pending, frame) >= 0) {
					wasEmpty = !m_haveFrame;
					m_haveFrame = true;
				}
			}
			if (wasEmpty && m_readyCb) m_readyCb();
			if (m_frameCb) m_frameCb();
			av_frame_unref(frame);
		}

		std::this_thread::sleep_for(std::chrono::microseconds(intervalUs));
	}

	av_frame_free(&frame);
	closeDecoder();
	m_threadRunning = false;
}

bool VideoPlayer::start(const QString& path) {
	stop();

	if (!ensureCudaContext()) return false;

	m_threadRunning = true;
	m_thread = std::thread([this, path]() { decodeLoop(path); });
	return true;
}

void VideoPlayer::stop() {
	m_threadRunning = false;
	if (m_thread.joinable()) m_thread.join();

	closeDecoder();

	std::lock_guard<std::mutex> lock(m_mtx);
	if (m_pending) {
		av_frame_unref(m_pending);
		m_pending = nullptr;
	}
	if (m_displayed) {
		av_frame_unref(m_displayed);
		m_displayed = nullptr;
	}
	m_haveFrame = false;
	m_size = QSize();
}

QSize VideoPlayer::frameSize() const {
	std::lock_guard<std::mutex> lock(m_mtx);
	return m_size;
}

void VideoPlayer::deleteTexture() {
	std::lock_guard<std::mutex> lock(m_mtx);
	CUcontext prev = nullptr;
	if (m_cuCtx) cuCtxPushCurrent(m_cuCtx);
	if (m_resY) {
		cuGraphicsUnregisterResource(m_resY);
		m_resY = nullptr;
	}
	if (m_resUV) {
		cuGraphicsUnregisterResource(m_resUV);
		m_resUV = nullptr;
	}
	if (m_texY) {
		glDeleteTextures(1, &m_texY);
		m_texY = 0;
	}
	if (m_texUV) {
		glDeleteTextures(1, &m_texUV);
		m_texUV = 0;
	}
	if (m_cuCtx) cuCtxPopCurrent(&prev);
	m_texSize = QSize();
}

bool VideoPlayer::ensureCudaContext() {
	if (m_cuCtx) return true;

	if (cuInit(0) != CUDA_SUCCESS) {
		qCWarning(videoLog, "VideoPlayer: cuInit failed");
		return false;
	}
	CUdevice dev = 0;
	if (cuDeviceGet(&dev, 0) != CUDA_SUCCESS) {
		qCWarning(videoLog, "VideoPlayer: cuDeviceGet failed");
		return false;
	}
	if (cuCtxCreate_v4(&m_cuCtx, nullptr, 0, dev) != CUDA_SUCCESS) {
		m_cuCtx = nullptr;
		qCWarning(videoLog, "VideoPlayer: cuCtxCreate failed");
		return false;
	}
	CUcontext popped = nullptr;
	cuCtxPopCurrent(&popped);
	return true;
}

bool VideoPlayer::ensureTextures(unsigned w, unsigned h, int compBytes) {
	if (!ensureCudaContext()) return false;
	CUcontext prev = nullptr;
	cuCtxPushCurrent(m_cuCtx);

	const QSize sz = QSize(int(w), int(h));
	if (m_texSize == sz && m_texComp == compBytes && m_texY != 0) {
		cuCtxPopCurrent(&prev);
		return true;
	}

	if (m_resY) cuGraphicsUnregisterResource(m_resY);
	if (m_resUV) cuGraphicsUnregisterResource(m_resUV);
	m_resY = m_resUV = nullptr;
	if (m_texY) glDeleteTextures(1, &m_texY);
	if (m_texUV) glDeleteTextures(1, &m_texUV);
	m_texY = m_texUV = 0;

	const bool wide = compBytes == 2;
	auto makeTex =
		[&](GLuint& tex, GLint internal, GLenum fmt, unsigned tw, unsigned th) {
			glGenTextures(1, &tex);
			glBindTexture(GL_TEXTURE_2D, tex);
			glTexImage2D(
				GL_TEXTURE_2D,
				0,
				internal,
				GLint(tw),
				GLint(th),
				0,
				fmt,
				wide ? GL_UNSIGNED_SHORT : GL_UNSIGNED_BYTE,
				nullptr);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
			glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
		};
	makeTex(m_texY, wide ? GL_R16 : GL_R8, GL_RED, w, h);
	makeTex(m_texUV, wide ? GL_RG16 : GL_RG8, GL_RG, w / 2, h / 2);

	CUresult r = cuGraphicsGLRegisterImage(
		&m_resY,
		m_texY,
		GL_TEXTURE_2D,
		CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD);
	if (r != CUDA_SUCCESS) {
		const char* es = nullptr;
		cuGetErrorString(r, &es);
		qCWarning(
			videoLog,
			"VideoPlayer: register luma tex failed (%s)",
			es ? es : "?");
		m_texSize = QSize();
		cuCtxPopCurrent(&prev);
		return false;
	}
	r = cuGraphicsGLRegisterImage(
		&m_resUV,
		m_texUV,
		GL_TEXTURE_2D,
		CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD);
	if (r != CUDA_SUCCESS) {
		qCWarning(videoLog, "VideoPlayer: register chroma tex failed");
		cuGraphicsUnregisterResource(m_resY);
		m_resY = nullptr;
		m_texSize = QSize();
		cuCtxPopCurrent(&prev);
		return false;
	}

	m_texSize = sz;
	m_texComp = compBytes;
	cuCtxPopCurrent(&prev);
	return true;
}

bool VideoPlayer::importCurrentFrame() {
	AVFrame* frame = nullptr;
	{
		std::lock_guard<std::mutex> lock(m_mtx);
		if (!m_pending || m_pending->format == AV_PIX_FMT_NONE) return false;
		frame = m_pending;
		m_pending = nullptr;
	}

	if (frame->format != AV_PIX_FMT_CUDA || !frame->hw_frames_ctx) {
		static int logged = 0;
		if (logged++ < 3)
			qCWarning(
				videoLog,
				"VideoPlayer: frame not on GPU (fmt=%s hw_frames=%d)",
				av_get_pix_fmt_name(AVPixelFormat(frame->format)),
				frame->hw_frames_ctx != nullptr);
		av_frame_unref(frame);
		return false;
	}

	const AVHWFramesContext* hf =
		reinterpret_cast<const AVHWFramesContext*>(frame->hw_frames_ctx->data);
	const AVPixelFormat sw = hf->sw_format;
	const bool tenBit = sw == AV_PIX_FMT_P010LE || sw == AV_PIX_FMT_P010BE ||
						sw == AV_PIX_FMT_P012LE || sw == AV_PIX_FMT_P016LE;
	const int compBytes = tenBit ? 2 : 1;
	const size_t compSize = size_t(compBytes);

	const unsigned w = unsigned(frame->width);
	const unsigned h = unsigned(frame->height);
	if (w == 0 || h == 0) {
		av_frame_unref(frame);
		return false;
	}

	size_t pitchY = size_t(frame->linesize[0]);
	size_t pitchUV = frame->linesize[1] > 0 ? size_t(frame->linesize[1])
											: pitchY;
	if (pitchY == 0) pitchY = size_t(w) * compSize;

	switch (frame->colorspace) {
	case AVCOL_SPC_BT470BG:
	case AVCOL_SPC_SMPTE170M:
		m_matrix = 0;
		break;
	case AVCOL_SPC_BT709:
		m_matrix = 1;
		break;
	default:
		m_matrix = h >= 720 ? 1 : 0;
		break;
	}
	m_range = frame->color_range == AVCOL_RANGE_JPEG ? 1 : 0;
	m_size = QSize(int(w), int(h));

	if (!ensureTextures(w, h, compBytes)) {
		av_frame_unref(frame);
		return false;
	}

	CUcontext prev = nullptr;
	cuCtxPushCurrent(m_cuCtx);
	CUgraphicsResource res[2] = {m_resY, m_resUV};
	CUresult r = cuGraphicsMapResources(2, res, nullptr);
	if (r != CUDA_SUCCESS) {
		cuCtxPopCurrent(&prev);
		av_frame_unref(frame);
		qCWarning(videoLog, "VideoPlayer: cuGraphicsMapResources failed");
		return false;
	}

	CUarray arrY = nullptr, arrUV = nullptr;
	cuGraphicsSubResourceGetMappedArray(&arrY, m_resY, 0, 0);
	cuGraphicsSubResourceGetMappedArray(&arrUV, m_resUV, 0, 0);

	const size_t dstPitch = size_t(w) * compSize;
	CUDA_MEMCPY2D mc{};
	mc.srcMemoryType = CU_MEMORYTYPE_DEVICE;
	mc.dstMemoryType = CU_MEMORYTYPE_ARRAY;
	mc.srcDevice = CUdeviceptr(frame->data[0]);
	mc.srcPitch = pitchY;
	mc.dstArray = arrY;
	mc.dstPitch = dstPitch;
	mc.WidthInBytes = std::min(pitchY, dstPitch);
	mc.Height = h;
	r = cuMemcpy2D(&mc);
	if (r == CUDA_SUCCESS) {
		mc.srcDevice = CUdeviceptr(frame->data[1]);
		mc.srcPitch = pitchUV;
		mc.dstArray = arrUV;
		mc.dstPitch = dstPitch;
		mc.WidthInBytes = std::min(pitchUV, dstPitch);
		mc.Height = h / 2;
		r = cuMemcpy2D(&mc);
	}

	cuGraphicsUnmapResources(2, res, nullptr);
	cuCtxPopCurrent(&prev);

	if (r != CUDA_SUCCESS) {
		av_frame_unref(frame);
		qCWarning(videoLog, "VideoPlayer: cuMemcpy2D failed");
		return false;
	}

	{
		std::lock_guard<std::mutex> lock(m_mtx);
		if (m_displayed) av_frame_unref(m_displayed);
		m_displayed = frame;
	}

	return true;
}

#endif