#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
#include <stdbool.h>
#include <time.h>

bool is_frame_all_zero(AVFrame *frame, int width, int height) {
	int num_bytes = width * height * 3; // 3 bytes per pixel (RGB)
	for (int i = 0; i < num_bytes; i++) {
		if (frame->data[0][i] != 0) {
			return false;
		}
	}
	return true;
}

char* get_current_iso_utc_date() {
	time_t now = time(NULL);
	struct tm *tm_info = gmtime(&now);
	char *buffer = malloc(20 * sizeof(char));
	strftime(buffer, 20, "%Y-%m-%dT%H-%M-%SZ", tm_info);
	return buffer;
}

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

	const char *input_filename = argv[1];

	// Open the input file
	AVFormatContext *pFormatContext = avformat_alloc_context();
	if (avformat_open_input(&pFormatContext, input_filename, NULL, NULL) != 0) {
		fprintf(stderr, "Could not open input file '%s'\n", input_filename);
		return -1;
	}

	// Retrieve stream information
	if (avformat_find_stream_info(pFormatContext, NULL) < 0) {
		fprintf(stderr, "Could not find stream information\n");
		return -1;
	}

	// Open the codec for the video stream
	AVCodec *pCodec = NULL;
	AVCodecParameters *pCodecParameters = NULL;
	int video_stream_index = -1;
	for (unsigned int i = 0; i < pFormatContext->nb_streams; i++) {
		AVCodecParameters *pLocalCodecParameters = pFormatContext->streams[i]->codecpar;
		AVCodec *pLocalCodec = avcodec_find_decoder(pLocalCodecParameters->codec_id);
		if (pLocalCodec == NULL) {
			continue;
		}
		if (pLocalCodecParameters->codec_type == AVMEDIA_TYPE_VIDEO) {
			video_stream_index = i;
			pCodec = pLocalCodec;
			pCodecParameters = pLocalCodecParameters;
			break;
		}
	}

	if (video_stream_index == -1) {
		fprintf(stderr, "Could not find a video stream\n");
		return -1;
	}

	AVCodecContext *pCodecContext = avcodec_alloc_context3(pCodec);
	if (avcodec_parameters_to_context(pCodecContext, pCodecParameters) < 0) {
		fprintf(stderr, "Could not copy codec parameters to context\n");
		return -1;
	}

	if (avcodec_open2(pCodecContext, pCodec, NULL) < 0) {
		fprintf(stderr, "Could not open codec\n");
		return -1;
	}

	// Allocate frame and packet
	AVFrame *pFrame = av_frame_alloc();
	AVPacket *pPacket = av_packet_alloc();

	// Initialize SWS context for software scaling
	struct SwsContext *sws_ctx = sws_getContext(
			pCodecContext->width,
			pCodecContext->height,
			pCodecContext->pix_fmt,
			pCodecContext->width,
			pCodecContext->height,
			AV_PIX_FMT_RGB24,
			SWS_BILINEAR,
			NULL,
			NULL,
			NULL
			);

	// Allocate buffer for RGB image
	int num_bytes = av_image_get_buffer_size(AV_PIX_FMT_RGB24, pCodecContext->width, pCodecContext->height, 1);
	uint8_t *buffer = (uint8_t *)av_malloc(num_bytes * sizeof(uint8_t));
	AVFrame *pFrameRGB = av_frame_alloc();
	av_image_fill_arrays(pFrameRGB->data, pFrameRGB->linesize, buffer, AV_PIX_FMT_RGB24, pCodecContext->width, pCodecContext->height, 1);

	int zero_frame_count = 0;
	bool start_writing = false;
	AVFormatContext *outFormatContext = NULL;

	// Read frames
	while (av_read_frame(pFormatContext, pPacket) >= 0) {
		if (pPacket->stream_index == video_stream_index) {
			if (avcodec_send_packet(pCodecContext, pPacket) >= 0) {
				while (avcodec_receive_frame(pCodecContext, pFrame) >= 0) {
					// Convert the image from its native format to RGB
					sws_scale(
							sws_ctx,
							(uint8_t const *const *)pFrame->data,
							pFrame->linesize,
							0,
							pCodecContext->height,
							pFrameRGB->data,
							pFrameRGB->linesize
						 );

					// Check if all pixels in the frame are zero
					bool all_zero = is_frame_all_zero(pFrameRGB, pCodecContext->width, pCodecContext->height);
					if (all_zero) {
						zero_frame_count++;
						if (zero_frame_count >= 5 && !start_writing) {
							start_writing = true;
							char *current_iso_utc_date = get_current_iso_utc_date();
							char output_filename[256];
							snprintf(output_filename, sizeof(output_filename), "%s.mp4", current_iso_utc_date);
							free(current_iso_utc_date);

							// Open the output file
							avformat_alloc_output_context2(&outFormatContext, NULL, NULL, output_filename);
							if (!outFormatContext) {
								fprintf(stderr, "Could not create output context\n");
								return -1;
							}

							// Create a stream for each input stream in the output file
							for (unsigned int i = 0; i < pFormatContext->nb_streams; i++) {
								AVStream *in_stream = pFormatContext->streams[i];
								AVStream *out_stream = avformat_new_stream(outFormatContext, NULL);
								if (!out_stream) {
									fprintf(stderr, "Failed allocating output stream\n");
									return -1;
								}
								if (avcodec_parameters_copy(out_stream->codecpar, in_stream->codecpar) < 0) {
									fprintf(stderr, "Failed to copy codec parameters\n");
									return -1;
								}
								out_stream->codecpar->codec_tag = 0;
							}

							// Open output file
							if (!(outFormatContext->oformat->flags & AVFMT_NOFILE)) {
								if (avio_open(&outFormatContext->pb, output_filename, AVIO_FLAG_WRITE) < 0) {
									fprintf(stderr, "Could not open output file '%s'\n", output_filename);
									return -1;
								}
							}

							// Write the stream header
							if (avformat_write_header(outFormatContext, NULL) < 0) {
								fprintf(stderr, "Error occurred when opening output file\n");
								return -1;
							}
						} else if (start_writing) {
							// Stop writing when an all-zero frame is detected
							if (outFormatContext) {
								// Write trailer to output file
								av_write_trailer(outFormatContext);

								if (!(outFormatContext->oformat->flags & AVFMT_NOFILE)) {
									avio_close(outFormatContext->pb);
								}
								avformat_free_context(outFormatContext);
								outFormatContext = NULL;
							}
							start_writing = false;
						}
					} else {
						zero_frame_count = 0;
					}
				}
			}
		}

		// Write packet to output file if the start_writing flag is set
		if (start_writing && outFormatContext) {
			AVStream *in_stream = pFormatContext->streams[pPacket->stream_index];
			AVStream *out_stream = outFormatContext->streams[pPacket->stream_index];
			pPacket->pts = av_rescale_q_rnd(pPacket->pts, in_stream->time_base, out_stream->time_base, AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
			pPacket->dts = av_rescale_q_rnd(pPacket->dts, in_stream->time_base, out_stream->time_base, AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
			pPacket->duration = av_rescale_q(pPacket->duration, in_stream->time_base, out_stream->time_base);
			pPacket->pos = -1;

			if (av_interleaved_write_frame(outFormatContext, pPacket) < 0) {
				fprintf(stderr, "Error while writing frame to output file\n");
				return -1;
			}
		}

		av_packet_unref(pPacket);
	}

	if (outFormatContext) {
		// Write trailer to output file
		av_write_trailer(outFormatContext);

		if (!(outFormatContext->oformat->flags & AVFMT_NOFILE)) {
			avio_close(outFormatContext->pb);
		}
		avformat_free_context(outFormatContext);
	}

	// Free memory
	av_free(buffer);
	av_frame_free(&pFrameRGB);
	av_frame_free(&pFrame);
	av_packet_free(&pPacket);
	avcodec_free_context(&pCodecContext);
	avformat_close_input(&pFormatContext);
	sws_freeContext(sws_ctx);

	return 0;
}