#include <fcntl.h>
#include <stdbool.h>
#include <stdio.h>
#include <unistd.h>

#include <drm_fourcc.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/hwcontext.h>
#include <libavutil/hwcontext_drm.h>
#include <libavutil/pixdesc.h>
#include <xf86drm.h>
#include <xf86drmMode.h>

static enum AVPixelFormat hw_pix_fmt = AV_PIX_FMT_NONE;

static const struct {
	uint32_t format;
	int nb_layers;
	uint32_t layers[AV_DRM_MAX_PLANES];
} drm_format_map[] = {
	{ DRM_FORMAT_NV12, 2, { DRM_FORMAT_R8, DRM_FORMAT_GR88 } },
};

static uint32_t get_drm_frame_format(const AVDRMFrameDescriptor *drm_frame_desc) {
	if (drm_frame_desc->nb_layers == 1) {
		return drm_frame_desc->layers[0].format;
	}
	for (size_t i = 0; i < sizeof(drm_format_map) / sizeof(drm_format_map[0]); i++) {
		if (drm_format_map[i].nb_layers != drm_frame_desc->nb_layers) {
			continue;
		}
		bool match = true;
		for (int j = 0; j < drm_frame_desc->nb_layers; j++) {
			match &= drm_frame_desc->layers[j].format == drm_format_map[i].layers[j];
		}
		if (match) {
			return drm_format_map[i].format;
		}
	}
	return DRM_FORMAT_INVALID;
}

static int check_hw_device_type(enum AVHWDeviceType type) {
	enum AVHWDeviceType t = AV_HWDEVICE_TYPE_NONE;
	while (1) {
		t = av_hwdevice_iterate_types(t);
		if (t == AV_HWDEVICE_TYPE_NONE) {
			break;
		}
		if (t == type) {
			return 0;
		}
	}
	return -1;
}

static enum AVPixelFormat get_hw_format(AVCodecContext *ctx,
		const enum AVPixelFormat *pix_fmts) {
	for (size_t i = 0; pix_fmts[i] != AV_PIX_FMT_NONE; i++) {
		if (pix_fmts[i] == hw_pix_fmt) {
			return hw_pix_fmt;
		}
	}

	fprintf(stderr, "Failed to find HW pixel format\n");
	return AV_PIX_FMT_NONE;
}

static uint32_t conn_id = 0;

static uint32_t pick_crtc(int drm_fd) {
	drmModeRes *res = drmModeGetResources(drm_fd);
	if (!res) {
		perror("drmModeGetResources");
		return 0;
	}

	drmModeConnector *conn = NULL;
	for (int i = 0; i < res->count_connectors; i++) {
		drmModeConnector *c = drmModeGetConnector(drm_fd, res->connectors[i]);
		if (c->connection == DRM_MODE_CONNECTED) {
			conn = c;
			break;
		}
		drmModeFreeConnector(c);
	}
	if (!conn) {
		fprintf(stderr, "Failed to find connected connector\n");
		return 0;
	}
	if (conn->encoder_id == 0) {
		fprintf(stderr, "Connector is not connected to an encoder\n");
		return 0;
	}

	conn_id = conn->connector_id;

	drmModeFreeResources(res);

	drmModeEncoder *enc = drmModeGetEncoder(drm_fd, conn->encoder_id);
	if (enc->crtc_id == 0) {
		fprintf(stderr, "Encoder is not connected to a CRTC\n");
		return 0;
	}
	uint32_t crtc_id = enc->crtc_id;
	drmModeFreeEncoder(enc);

	return crtc_id;
}

static int decode_frame(AVFormatContext *input_ctx, int video_stream,
		AVCodecContext *decoder_ctx, AVFrame **frame_ptr) {
	int ret;
	while (1) {
		AVPacket packet;
		ret = av_read_frame(input_ctx, &packet);
		if (ret < 0) {
			fprintf(stderr, "Failed to read frames into packet\n");
			return ret;
		}

		if (video_stream != packet.stream_index) {
			av_packet_unref(&packet);
			continue;
		}

		ret = avcodec_send_packet(decoder_ctx, &packet);
		if (ret < 0) {
			fprintf(stderr, "Failed to send packet to decoder\n");
			return ret;
		}
		av_packet_unref(&packet);

		AVFrame *frame = av_frame_alloc();
		if (!frame) {
			fprintf(stderr, "Failed to allocate frame\n");
			return AVERROR(ENOMEM);
		}

		while (1) {
			ret = avcodec_receive_frame(decoder_ctx, frame);
			if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
				break;
			} else if (ret < 0) {
				fprintf(stderr, "Failed to receive frame from decoder\n");
				return ret;
			}

			if (frame->format != hw_pix_fmt) {
				continue;
			}

			*frame_ptr = frame;
			return 0;
		}

		av_frame_free(&frame);
	}

	return AVERROR_EOF;
}

static void print_drm_frame_desc(const AVDRMFrameDescriptor *drm_frame_desc) {
	fprintf(stderr, "DRM frame:\n");
	for (int i = 0; i < drm_frame_desc->nb_objects; i++) {
		const AVDRMObjectDescriptor *drm_object = &drm_frame_desc->objects[i];
		fprintf(stderr, "  object #%i: fd = %d, size = %zu, "
			"format_modifier = 0x%lX\n", i, drm_object->fd,
			drm_object->size, drm_object->format_modifier);
	}
	for (int i = 0; i < drm_frame_desc->nb_layers; i++) {
		const AVDRMLayerDescriptor *drm_layer = &drm_frame_desc->layers[i];
		fprintf(stderr, "  layer #%d: format = 0x%X\n",
			i, drm_layer->format);
		for (int j = 0; j < drm_layer->nb_planes; j++) {
			const AVDRMPlaneDescriptor *drm_plane = &drm_layer->planes[j];
			fprintf(stderr, "    plane #%d: object_index = %d, "
				"offset = %ld, pitch = %ld\n", j, drm_plane->object_index,
				drm_plane->offset, drm_plane->pitch);
		}
	}
}

static uint32_t import_drm_frame_desc(const AVDRMFrameDescriptor *drm_frame_desc,
		int drm_fd, int width, int height) {
	// VA-API drivers may use separate layers with one plane each, or a single
	// layer with multiple planes. We need to handle both.
	uint32_t drm_format = get_drm_frame_format(drm_frame_desc);
	if (drm_format == DRM_FORMAT_INVALID) {
		fprintf(stderr, "Failed to get DRM frame format\n");
		return 0;
	}
	fprintf(stderr, "DRM format: 0x%X\n", drm_format);

	uint32_t handles[4];
	uint32_t pitches[4];
	uint32_t offsets[4];
	int k = 0;
	for (int i = 0; i < drm_frame_desc->nb_layers; i++) {
		const AVDRMLayerDescriptor *drm_layer = &drm_frame_desc->layers[i];
		for (int j = 0; j < drm_layer->nb_planes; j++) {
			const AVDRMPlaneDescriptor *drm_plane = &drm_layer->planes[j];
			const AVDRMObjectDescriptor *drm_object = &drm_frame_desc->objects[drm_plane->object_index];

			uint32_t handle = 0;
			int ret = drmPrimeFDToHandle(drm_fd, drm_object->fd, &handle);
			if (ret < 0) {
				perror("drmPrimeFDToHandle");
				return 0;
			}

			handles[k] = handle;
			pitches[k] = drm_plane->pitch;
			offsets[k] = drm_plane->offset;
			k++;
		}
	}

	uint32_t fb_id = 0;
	int ret = drmModeAddFB2(drm_fd, width, height, drm_format,
		handles, pitches, offsets, &fb_id, 0);
	if (ret < 0) {
		perror("drmModeAddFB2");
		return 0;
	}

	return fb_id;
}

int main(int argc, char *argv[]) {
	if (argc < 2) {
		fprintf(stderr, "usage: %s <file>\n", argv[0]);
		return 1;
	}

	enum AVHWDeviceType type = AV_HWDEVICE_TYPE_VAAPI;
	int ret = check_hw_device_type(type);
	if (ret != 0) {
		fprintf(stderr, "VA-API not supported\n");
		return 1;
	}

	AVFormatContext *input_ctx = NULL;
	if (avformat_open_input(&input_ctx, argv[1], NULL, NULL) != 0) {
		fprintf(stderr, "Failed to open input file\n");
		return 1;
	}

	if (avformat_find_stream_info(input_ctx, NULL) < 0) {
		fprintf(stderr, "Failed to find input stream info\n");
		return 1;
	}

	AVCodec *decoder = NULL;
	int video_stream = av_find_best_stream(input_ctx, AVMEDIA_TYPE_VIDEO,
		-1, -1, &decoder, 0);
	if (video_stream < 0) {
		fprintf(stderr, "Failed to find video stream in input file\n");
		return 1;
	}
	fprintf(stderr, "Selected decoder %s\n", decoder->name);

	int i = 0;
	while (1) {
		const AVCodecHWConfig *config = avcodec_get_hw_config(decoder, i);
		if (!config) {
			fprintf(stderr, "Decoder %s doesn't support device type %s\n",
				decoder->name, av_hwdevice_get_type_name(type));
			return 1;
		}

		if ((config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) &&
				config->device_type == type) {
			hw_pix_fmt = config->pix_fmt;
			break;
		}

		i++;
	}
	fprintf(stderr, "Selected pixel format %s\n", av_get_pix_fmt_name(hw_pix_fmt));

	AVCodecContext *decoder_ctx = avcodec_alloc_context3(decoder);
	if (!decoder_ctx) {
		fprintf(stderr, "Failed to allocate decoder context\n");
		return 1;
	}

	AVStream *video = input_ctx->streams[video_stream];
	if (avcodec_parameters_to_context(decoder_ctx, video->codecpar) < 0) {
		fprintf(stderr, "Failed to fill decoder context from video parameters\n");
		return 1;
	}
	fprintf(stderr, "Stream geometry: %dx%d\n", decoder_ctx->width, decoder_ctx->height);

	decoder_ctx->get_format = get_hw_format;

	AVBufferRef *hw_device_ctx = NULL;
	ret = av_hwdevice_ctx_create(&hw_device_ctx, type, NULL, NULL, 0);
	if (ret < 0) {
		fprintf(stderr, "Failed to create HW device context\n");
		return 1;
	}
	decoder_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);

	ret = avcodec_open2(decoder_ctx, decoder, NULL);
	if (ret < 0) {
		fprintf(stderr, "Failed to open codec for stream #%u\n", video_stream);
		return 1;
	}

	int drm_fd = open("/dev/dri/card0", O_RDWR);
	if (drm_fd < 0) {
		fprintf(stderr, "Failed to open DRM primary node\n");
		return 1;
	}

	uint32_t crtc_id = pick_crtc(drm_fd);
	if (crtc_id == 0) {
		return 1;
	}

	while (1) {

	AVFrame *frame = NULL;
	ret = decode_frame(input_ctx, video_stream, decoder_ctx, &frame);
	if (ret < 0) {
		return 1;
	}

	fprintf(stderr, "Got %dx%d frame!\n", frame->width, frame->height);

	AVFrame *drm_frame = av_frame_alloc();
	drm_frame->format = AV_PIX_FMT_DRM_PRIME;
	drm_frame->hw_frames_ctx = av_buffer_ref(frame->hw_frames_ctx);

	// Convert the VA-API frame into a DMA-BUF frame
	ret = av_hwframe_map(drm_frame, frame, 0);
	if (ret < 0) {
		fprintf(stderr, "Failed to map frame: %s\n", av_err2str(ret));
		return 1;
	}

	AVDRMFrameDescriptor *drm_frame_desc = (void *)drm_frame->data[0];
	print_drm_frame_desc(drm_frame_desc);

	uint32_t fb_id = import_drm_frame_desc(drm_frame_desc, drm_fd,
		drm_frame->width, drm_frame->height);
	if (!fb_id) {
		return 1;
	}

	drmModeCrtc *crtc = drmModeGetCrtc(drm_fd, crtc_id);
	if (!crtc->mode_valid) {
		fprintf(stderr, "No mode set\n");
		return 1;
	}

	static bool first = true;
	if (first) {
		ret = drmModeSetCrtc(drm_fd, crtc_id, fb_id, 0, 0, &conn_id, 1, &crtc->mode);
	} else {
		ret = drmModePageFlip(drm_fd, crtc_id, fb_id, 0, NULL);
	}
	if (ret < 0) {
		perror("drmModePageFlip");
		return 1;
	}

	av_frame_free(&drm_frame);
	av_frame_free(&frame);

	//usleep(16);
	}

	avcodec_free_context(&decoder_ctx);
	avformat_close_input(&input_ctx);
	av_buffer_unref(&hw_device_ctx);

	return 0;
}