//
// Created by Matthew on 2025/3/1.
//

#include "RTSPRecorder.h"
#include <chrono>
#include <thread>
#include <android/log.h>
extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/opt.h>
}


#define LOG_TAG "libcurl"

#define LOGV(...) __android_log_print(ANDROID_LOG_VERBOSE, LOG_TAG, __VA_ARGS__)
#define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)


void dumpRtmpToMp4(const char* rtmpUrl, const char* outputPath, uint32_t duration)
{
    AVFormatContext* inputFormatContext = nullptr;
    AVFormatContext* outputFormatContext = nullptr;
    AVPacket packet;

    av_register_all();
    avformat_network_init();

    // Open input RTMP stream
    if (avformat_open_input(&inputFormatContext, rtmpUrl, nullptr, nullptr) != 0) {
        fprintf(stderr, "Could not open input file '%s'\n", rtmpUrl);
        return;
    }

    // Retrieve input stream information
    if (avformat_find_stream_info(inputFormatContext, nullptr) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        avformat_close_input(&inputFormatContext);
        return;
    }

    // Open output MP4 file
    if (avformat_alloc_output_context2(&outputFormatContext, nullptr, "mp4", outputPath) < 0) {
        fprintf(stderr, "Could not create output context\n");
        avformat_close_input(&inputFormatContext);
        return;
    }

    // Copy stream information from input to output
    for (unsigned int i = 0; i < inputFormatContext->nb_streams; i++) {
        AVStream* inStream = inputFormatContext->streams[i];
        AVStream* outStream = avformat_new_stream(outputFormatContext, nullptr);
        if (!outStream) {
            fprintf(stderr, "Failed to allocate output stream\n");
            avformat_close_input(&inputFormatContext);
            avformat_free_context(outputFormatContext);
            return;
        }

        if (avcodec_parameters_copy(outStream->codecpar, inStream->codecpar) < 0) {
            fprintf(stderr, "Failed to copy codec parameters\n");
            avformat_close_input(&inputFormatContext);
            avformat_free_context(outputFormatContext);
            return;
        }
        outStream->codecpar->codec_tag = 0;
    }

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

    // Write output file header
    if (avformat_write_header(outputFormatContext, nullptr) < 0) {
        fprintf(stderr, "Error occurred when writing header to output file\n");
        avformat_close_input(&inputFormatContext);
        avformat_free_context(outputFormatContext);
        return;
    }

    // Start a thread to stop the streaming after the specified duration
    std::thread stop_thread([&]() {
        std::this_thread::sleep_for(std::chrono::milliseconds(duration));
        av_read_pause(inputFormatContext);
    });

    // Read packets from input and write them to output
    while (av_read_frame(inputFormatContext, &packet) >= 0) {
        AVStream* inStream = inputFormatContext->streams[packet.stream_index];
        AVStream* outStream = outputFormatContext->streams[packet.stream_index];

        packet.pts = av_rescale_q_rnd(packet.pts, inStream->time_base, outStream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.dts = av_rescale_q_rnd(packet.dts, inStream->time_base, outStream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.duration = av_rescale_q(packet.duration, inStream->time_base, outStream->time_base);
        packet.pos = -1;

        if (av_interleaved_write_frame(outputFormatContext, &packet) < 0) {
            fprintf(stderr, "Error muxing packet\n");
            break;
        }

        av_packet_unref(&packet);
    }

    stop_thread.join();

    // Write output file trailer
    av_write_trailer(outputFormatContext);

    // Clean up
    avformat_close_input(&inputFormatContext);
    if (outputFormatContext && !(outputFormatContext->oformat->flags & AVFMT_NOFILE)) {
        avio_closep(&outputFormatContext->pb);
    }
    avformat_free_context(outputFormatContext);
}


void dumpRtspToMp4(const char* rtspUrl, const char* outputPath, uint32_t duration)
{
    AVFormatContext* inputFormatContext = nullptr;
    AVFormatContext* outputFormatContext = nullptr;
    AVPacket packet;

    av_register_all();
    avformat_network_init();

    // Open input RTSP stream
    if (avformat_open_input(&inputFormatContext, rtspUrl, nullptr, nullptr) != 0) {
        fprintf(stderr, "Could not open input file '%s'\n", rtspUrl);
        return;
    }

    // Retrieve input stream information
    if (avformat_find_stream_info(inputFormatContext, nullptr) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        avformat_close_input(&inputFormatContext);
        return;
    }

    // Open output MP4 file
    if (avformat_alloc_output_context2(&outputFormatContext, nullptr, "mp4", outputPath) < 0) {
        fprintf(stderr, "Could not create output context\n");
        avformat_close_input(&inputFormatContext);
        return;
    }

    // Copy stream information from input to output
    for (unsigned int i = 0; i < inputFormatContext->nb_streams; i++) {
        AVStream* inStream = inputFormatContext->streams[i];
        AVStream* outStream = avformat_new_stream(outputFormatContext, nullptr);
        if (!outStream) {
            fprintf(stderr, "Failed to allocate output stream\n");
            avformat_close_input(&inputFormatContext);
            avformat_free_context(outputFormatContext);
            return;
        }

        if (avcodec_parameters_copy(outStream->codecpar, inStream->codecpar) < 0) {
            fprintf(stderr, "Failed to copy codec parameters\n");
            avformat_close_input(&inputFormatContext);
            avformat_free_context(outputFormatContext);
            return;
        }
        outStream->codecpar->codec_tag = 0;
    }

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

    // Write output file header
    if (avformat_write_header(outputFormatContext, nullptr) < 0) {
        fprintf(stderr, "Error occurred when writing header to output file\n");
        avformat_close_input(&inputFormatContext);
        avformat_free_context(outputFormatContext);
        return;
    }

    // Start a thread to stop the streaming after the specified duration
    std::thread stop_thread([&]() {
        std::this_thread::sleep_for(std::chrono::milliseconds(duration));
        av_read_pause(inputFormatContext);
    });

    // Read packets from input and write them to output
    while (av_read_frame(inputFormatContext, &packet) >= 0) {
        AVStream* inStream = inputFormatContext->streams[packet.stream_index];
        AVStream* outStream = outputFormatContext->streams[packet.stream_index];

        packet.pts = av_rescale_q_rnd(packet.pts, inStream->time_base, outStream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.dts = av_rescale_q_rnd(packet.dts, inStream->time_base, outStream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.duration = av_rescale_q(packet.duration, inStream->time_base, outStream->time_base);
        packet.pos = -1;

        if (av_interleaved_write_frame(outputFormatContext, &packet) < 0) {
            fprintf(stderr, "Error muxing packet\n");
            break;
        }

        av_packet_unref(&packet);
    }

    stop_thread.join();

    // Write output file trailer
    av_write_trailer(outputFormatContext);

    // Clean up
    avformat_close_input(&inputFormatContext);
    if (outputFormatContext && !(outputFormatContext->oformat->flags & AVFMT_NOFILE)) {
        avio_closep(&outputFormatContext->pb);
    }
    avformat_free_context(outputFormatContext);
}