//
// Created by Matthew on 2025/2/28.
//

#include "RTSPToMP4.h"
#include <android/native_window.h>
#include <android/native_window_jni.h>
#include <jni.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <cstring>
#include <limits>

int32_t getMaxInputSize(AMediaExtractor* extractor, size_t trackIndex)
{
    AMediaFormat* format = AMediaExtractor_getTrackFormat(extractor, trackIndex);
    int32_t maxInputSize = 0;
    if (AMediaFormat_getInt32(format, AMEDIAFORMAT_KEY_MAX_INPUT_SIZE, &maxInputSize)) {
        // LOGI("Max input size for track %zu: %d", trackIndex, maxInputSize);
    } else {
        // LOGE("Failed to get max input size for track %zu", trackIndex);
    }
    AMediaFormat_delete(format);
    return maxInputSize;
}

RTSPToMP4::RTSPToMP4(const char* rtspUrl, const char* outputPath, uint64_t durationInMs/* = 0*/)
        : fd(-1), codec(nullptr), extractor(nullptr), muxer(nullptr), videoTrackIndex(-1), durationInMs(durationInMs), running(false) {
    initExtractor(rtspUrl);
    initCodec("video/avc");
    initMuxer(outputPath);
}

RTSPToMP4::~RTSPToMP4() {
    if (codec) AMediaCodec_delete(codec);
    if (extractor) AMediaExtractor_delete(extractor);
    if (muxer) AMediaMuxer_delete(muxer);

    if (fd != -1)
    {
        fdatasync(fd);
        close(fd);
        fd = -1;
    }
}

void RTSPToMP4::initCodec(const char* mime) {
    codec = AMediaCodec_createDecoderByType(mime);
    AMediaFormat* format = AMediaFormat_new();
    AMediaFormat_setString(format, AMEDIAFORMAT_KEY_MIME, mime);
    // Set other format parameters as needed
    // ...
    AMediaCodec_configure(codec, format, nullptr, nullptr, 0);
    AMediaFormat_delete(format);
}

void RTSPToMP4::initExtractor(const char* rtspUrl) {
    extractor = AMediaExtractor_new();
    media_status_t status = AMediaExtractor_setDataSource(extractor, rtspUrl);
    if (status != AMEDIA_OK) {
        // Handle error
        // ...
    }
}

void RTSPToMP4::initMuxer(const char* outputPath) {
    fd = open(outputPath, O_CREAT | O_WRONLY, 0644);
    muxer = AMediaMuxer_new(fd, AMEDIAMUXER_OUTPUT_FORMAT_MPEG_4);

    int numTracks = AMediaExtractor_getTrackCount(extractor);
    if (numTracks <= 0) {
        // LOGE("No tracks found in RTSP stream");
        AMediaExtractor_delete(extractor);
        return;
    }

    for (int i = 0; i < numTracks; ++i) {
        AMediaFormat* format = AMediaExtractor_getTrackFormat(extractor, i);
        const char* mime;
        if (AMediaFormat_getString(format, AMEDIAFORMAT_KEY_MIME, &mime) && strncmp(mime, "video/", 6) == 0) {
            videoTrackIndex = AMediaMuxer_addTrack(muxer, format);
            AMediaExtractor_selectTrack(extractor, i);
        }
        AMediaFormat_delete(format);
    }

    if (videoTrackIndex == -1) {
        // LOGE("No video track found in RTSP stream");
        AMediaExtractor_delete(extractor);
        AMediaMuxer_delete(muxer);
        return;
    }

    int32_t maxInputSize = getMaxInputSize(extractor, videoTrackIndex);
    if (maxInputSize <= 0) {
        // LOGE("Invalid max input size");
        // releaseMediaExtractor(extractor);
        sampleData.resize(1920 * 1080 * 4, 0);
        return;
    }

    sampleData.resize(maxInputSize, 0);
}

void RTSPToMP4::startDecodingAndMuxing() {
    AMediaCodec_start(codec);
    size_t bufferSize = sampleData.size();
    uint8_t* buffer = &sampleData[0];
    int64_t sampleTime = 0;
    int64_t startTime = 0;
    bool firstSampleData = true;

    int64_t durationTime = (durationInMs == 0) ? std::numeric_limits<int64_t>::max() : (int64_t)durationInMs * 1000;


    while (running) {
        // Extract data from RTSP stream
        ssize_t sampleSize = AMediaExtractor_readSampleData(extractor, buffer, bufferSize);
        if (sampleSize < 0) {
            break; // End of stream
        }

        sampleTime = AMediaExtractor_getSampleTime(extractor);
        if (firstSampleData)
        {
            startTime = sampleTime;
            firstSampleData = false;
        }

        sampleTime -= startTime;

        // Feed data to codec
        size_t inputBufferIndex;
        uint8_t* inputBuffer = AMediaCodec_getInputBuffer(codec, inputBufferIndex, &bufferSize);
        memcpy(inputBuffer, buffer, sampleSize);
        AMediaCodec_queueInputBuffer(codec, inputBufferIndex, 0, sampleSize, sampleTime, 0);

        // Retrieve decoded frames and write to muxer
        AMediaCodecBufferInfo bufferInfo;
        ssize_t outputBufferIndex = AMediaCodec_dequeueOutputBuffer(codec, &bufferInfo, 0);
        if (outputBufferIndex >= 0) {

            bufferInfo.offset = 0;
            bufferInfo.size = sampleSize;
            bufferInfo.presentationTimeUs = sampleTime;
            bufferInfo.flags = AMediaExtractor_getSampleFlags(extractor);

            uint8_t* outputBuffer = AMediaCodec_getOutputBuffer(codec, outputBufferIndex, &bufferSize);
            AMediaMuxer_writeSampleData(muxer, videoTrackIndex, outputBuffer, &bufferInfo);
            AMediaCodec_releaseOutputBuffer(codec, outputBufferIndex, false);
        }

        AMediaExtractor_advance(extractor);

        if (sampleTime > durationTime)
        {
            break;
        }
    }

    AMediaCodec_stop(codec);
    AMediaMuxer_stop(muxer);

    if (fd != -1)
    {
        fdatasync(fd);
        close(fd);
        fd = -1;
    }
}

void RTSPToMP4::start() {
    // Add video track to muxer
    AMediaFormat* format = AMediaExtractor_getTrackFormat(extractor, 0);
    videoTrackIndex = AMediaMuxer_addTrack(muxer, format);
    running = true;
    AMediaMuxer_start(muxer);

    startDecodingAndMuxing();
}

void RTSPToMP4::stop() {
    running = false;
}