#include <stdio.h>
#include <stdlib.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/imgutils.h>

// Function prototypes
int save_frame_as_jpeg(AVFrame *pFrame, int width, int height, const char *filename);

int main(int argc, char *argv[]) {
    if (argc < 3) {
        printf("Usage: %s <filename> <timestamp(ms)>\n", argv[0]);
        return -1;
    }

    const char *filename = argv[1];
    int64_t timestamp = atoll(argv[2]);

    av_register_all();

    AVFormatContext *pFormatContext = avformat_alloc_context();
    if (!pFormatContext) {
        printf("Could not allocate memory for Format Context\n");
        return -1;
    }

    if (avformat_open_input(&pFormatContext, filename, NULL, NULL) != 0) {
        printf("Could not open the file\n");
        return -1;
    }

    if (avformat_find_stream_info(pFormatContext, NULL) < 0) {
        printf("Could not get the stream info\n");
        return -1;
    }

    int video_stream_index = -1;
    AVCodecParameters *pCodecParameters = NULL;
    for (int i = 0; i < pFormatContext->nb_streams; i++) {
        if (pFormatContext->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            pCodecParameters = pFormatContext->streams[i]->codecpar;
            break;
        }
    }

    if (video_stream_index == -1) {
        printf("Could not find video stream in the input file\n");
        return -1;
    }

    AVCodec *pCodec = avcodec_find_decoder(pCodecParameters->codec_id);
    if (!pCodec) {
        printf("Could not find the codec\n");
        return -1;
    }

    AVCodecContext *pCodecContext = avcodec_alloc_context3(pCodec);
    if (!pCodecContext) {
        printf("Could not allocate memory for Codec Context\n");
        return -1;
    }

    if (avcodec_parameters_to_context(pCodecContext, pCodecParameters) < 0) {
        printf("Could not copy codec params to codec context\n");
        return -1;
    }

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

    AVFrame *pFrame = av_frame_alloc();
    if (!pFrame) {
        printf("Could not allocate memory for AVFrame\n");
        return -1;
    }

    int64_t timestamp_in_stream_timebase = av_rescale_q(timestamp, AV_TIME_BASE_Q, pFormatContext->streams[video_stream_index]->time_base);
    av_seek_frame(pFormatContext, video_stream_index, timestamp_in_stream_timebase, AVSEEK_FLAG_BACKWARD);

    AVPacket *packet = av_packet_alloc();
    if (!packet) {
        printf("Could not allocate memory for AVPacket\n");
        return -1;
    }

    int frame_finished = 0;
    while (av_read_frame(pFormatContext, packet) >= 0) {
        if (packet->stream_index == video_stream_index) {
            avcodec_send_packet(pCodecContext, packet);
            if (avcodec_receive_frame(pCodecContext, pFrame) == 0) {
                save_frame_as_jpeg(pFrame, pCodecContext->width, pCodecContext->height, "output.jpg");
                break;
            }
        }
        av_packet_unref(packet);
    }

    av_packet_free(&packet);
    av_frame_free(&pFrame);
    avcodec_free_context(&pCodecContext);
    avformat_close_input(&pFormatContext);

    return 0;
}

int save_frame_as_jpeg(AVFrame *pFrame, int width, int height, const char *filename) {
    AVCodecContext *pJpegCodecContext = NULL;
    AVCodec *pJpegCodec = NULL;
    AVPacket packet;
    int ret;

    pJpegCodec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!pJpegCodec) {
        printf("Could not find JPEG codec\n");
        return -1;
    }

    pJpegCodecContext = avcodec_alloc_context3(pJpegCodec);
    if (!pJpegCodecContext) {
        printf("Could not allocate JPEG codec context\n");
        return -1;
    }

    pJpegCodecContext->pix_fmt = AV_PIX_FMT_YUVJ420P;
    pJpegCodecContext->height = height;
    pJpegCodecContext->width = width;
    pJpegCodecContext->time_base = (AVRational){1, 25};

    if (avcodec_open2(pJpegCodecContext, pJpegCodec, NULL) < 0) {
        printf("Could not open JPEG codec\n");
        return -1;
    }

    av_init_packet(&packet);
    packet.data = NULL;
    packet.size = 0;

    ret = avcodec_send_frame(pJpegCodecContext, pFrame);
    if (ret < 0) {
        printf("Error sending frame to codec context\n");
        return ret;
    }

    ret = avcodec_receive_packet(pJpegCodecContext, &packet);
    if (ret < 0) {
        printf("Error receiving packet from codec context\n");
        return ret;
    }

    FILE *JPEGFile = fopen(filename, "wb");
    if (!JPEGFile) {
        printf("Could not open JPEG file\n");
        return -1;
    }

    fwrite(packet.data, 1, packet.size, JPEGFile);
    fclose(JPEGFile);

    av_packet_unref(&packet);
    avcodec_free_context(&pJpegCodecContext);

    return 0;
}