#include <fstream>
#include <iostream>

extern "C" {
#include <libavformat/avformat.h>
}

const std::string input = "data\\stream_0.h265";
const std::string output = "data\\stream_0_out.mp4";

const struct nalUnit
{
    int start;
    int size;
    const long long pts;
} nalUnits[] = {
    {0, 27, 0}, {27, 67, 0}, {94, 14, 0}, {108, 10091, 0}, {10199, 6571, 0}, {16770, 5021, 0}, {21791, 614, 40000},
    {22405, 392, 40000}, {22797, 432, 40000}, {23229, 490, 80000}, {23719, 281, 80000}, {24000, 434, 80000},
    {24434, 488, 120000}, {24922, 220, 120000}, {25142, 314, 120000}, {25456, 420, 160000}, {25876, 233, 160000},
    {26109, 405, 160000}, {26514, 513, 200000}, {27027, 284, 200000}, {27311, 401, 200000}, {27712, 512, 240000},
    {28224, 243, 240000}, {28467, 355, 240000}, {28822, 407, 280000}, {29229, 277, 280000}, {29506, 381, 280000},
    {29887, 372, 320000}, {30259, 245, 320000}, {30504, 365, 320000}, {30869, 444, 360000}, {31313, 230, 360000},
    {31543, 385, 360000}, {31928, 444, 400000}, {32372, 236, 400000}, {32608, 315, 400000}, {32923, 398, 440000},
    {33321, 304, 440000}, {33625, 300, 440000}, {33925, 465, 480000}, {34390, 159, 480000}, {34549, 371, 480000},
    {34920, 382, 520000}, {35302, 241, 520000}, {35543, 357, 520000}, {35900, 338, 560000}, {36238, 152, 560000},
    {36390, 278, 560000}, {36668, 442, 600000}, {37110, 256, 600000}, {37366, 355, 600000}, {37721, 418, 640000},
    {38139, 207, 640000}, {38346, 379, 640000}, {38725, 334, 680000}, {39059, 210, 680000}, {39269, 349, 680000},
    {39618, 330, 720000}, {39948, 260, 720000}, {40208, 346, 720000}, {40554, 372, 760000}, {40926, 247, 760000},
    {41173, 316, 760000}, {41489, 284, 800000}, {41773, 228, 800000}, {42001, 265, 800000}, {42266, 405, 840000},
    {42671, 229, 840000}, {42900, 369, 840000}, {43269, 349, 880000}, {43618, 330, 880000}, {43948, 237, 880000},
    {44185, 411, 920000}, {44596, 192, 920000}, {44788, 278, 920000}, {45066, 397, 960000}, {45463, 235, 960000},
    {45698, 335, 960000}, {46033, 27, 1000000}, {46060, 67, 1000000}, {46127, 14, 1000000}, {46141, 10130, 1000000},
    {56271, 6612, 1000000}, {62883, 5110, 1000000}, {67993, 584, 1040000}, {68577, 368, 1040000}, {68945, 388, 1040000},
    {69333, 512, 1080000}, {69845, 223, 1080000}, {70068, 375, 1080000}, {70443, 465, 1120000}, {70908, 265, 1120000},
    {71173, 431, 1120000}, {71604, 464, 1160000}, {72068, 278, 1160000}, {72346, 445, 1160000}, {72791, 368, 1200000},
    {73159, 245, 1200000}, {73404, 377, 1200000}, {73781, 449, 1240000}, {74230, 248, 1240000}, {74478, 386, 1240000},
    {74864, 491, 1280000}, {75355, 226, 1280000}, {75581, 328, 1280000}, {75909, 403, 1320000}, {76312, 223, 1320000},
    {76535, 454, 1320000}, {76989, 457, 1360000}, {77446, 272, 1360000}, {77718, 371, 1360000}, {78089, 431, 1400000},
    {78520, 253, 1400000}, {78773, 303, 1400000}, {79076, 397, 1440000}, {79473, 294, 1440000}, {79767, 366, 1440000},
    {80133, 405, 1480000}, {80538, 253, 1480000}, {80791, 394, 1480000}, {81185, 361, 1520000}, {81546, 273, 1520000},
    {81819, 344, 1520000}, {82163, 386, 1560000}, {82549, 239, 1560000}, {82788, 353, 1560000}, {83141, 395, 1600000},
    {83536, 279, 1600000}, {83815, 340, 1600000}, {84155, 449, 1640000}, {84604, 199, 1640000}, {84803, 366, 1640000},
    {85169, 381, 1680000}, {85550, 267, 1680000}, {85817, 317, 1680000}, {86134, 412, 1720000}, {86546, 251, 1720000},
    {86797, 320, 1720000}, {87117, 325, 1760000}, {87442, 226, 1760000}, {87668, 270, 1760000}, {87938, 438, 1767511},
    {88376, 215, 1767511}, {88591, 255, 1767511}, {88846, 336, 1807511}, {89182, 224, 1807511}, {89406, 290, 1807511},
    {89696, 372, 1847511}, {90068, 236, 1847511}, {90304, 278, 1847511}, {90582, 389, 1887511}, {90971, 189, 1887511},
    {91160, 286, 1887511}, {91446, 405, 1927511}, {91851, 244, 1927511}, {92095, 280, 1927511}, {92375, 27, 1967511},
    {92402, 67, 1967511}, {92469, 14, 1967511}, {92483, 10136, 1967511}, {102619, 6585, 1967511},
    {109204, 5051, 1967511}, {114255, 639, 2007511}, {114894, 354, 2007511}, {115248, 416, 2007511},
    {115664, 523, 2047511}, {116187, 322, 2047511}, {116509, 414, 2047511}, {116923, 454, 2087511},
    {117377, 295, 2087511}, {117672, 474, 2087511}, {118146, 502, 2127511}, {118648, 290, 2127511},
    {118938, 359, 2127511}, {119297, 476, 2167511}, {119773, 283, 2167511}, {120056, 440, 2167511},
    {120496, 404, 2207511}, {120900, 279, 2207511}, {121179, 348, 2207511}, {121527, 470, 2247511},
    {121997, 225, 2247511}, {122222, 381, 2247511}, {122603, 420, 2287511}, {123023, 274, 2287511},
    {123297, 300, 2287511}, {123597, 389, 2327511}, {123986, 296, 2327511}, {124282, 405, 2327511},
    {124687, 408, 2367511}, {125095, 300, 2367511}, {125395, 311, 2367511}, {125706, 378, 2407511},
    {126084, 246, 2407511}, {126330, 433, 2407511}, {126763, 383, 2447511}, {127146, 270, 2447511},
    {127416, 344, 2447511}, {127760, 434, 2487511}, {128194, 229, 2487511}, {128423, 351, 2487511}
};

const size_t fileLength = 128423+351;

const AVRational timebase = { 1,1000000 };

bool LoadInputFormat(AVFormatContext*& inputFormatContext)
{
    int ret;
    // create an AVFormatContext object using the video file. This can be used
    // to query information such as sample aspect ratio of the input file.
    if ((ret = avformat_open_input(&inputFormatContext, input.c_str(), nullptr, nullptr)) < 0)
    {
        std::cout << "Error: Unable to open input file: " << input.c_str() << ". Error code: " << ret <<std::endl;
        return true;
    }

    // search the file for stream info and add to the context
    if ((ret = avformat_find_stream_info(inputFormatContext, nullptr)) < 0)
    {
        std::cout << "Error: Unable to find stream info: " << input.c_str() << ". Error code: " << ret <<std::endl;
        return true;
    }

    av_dump_format(inputFormatContext, 0, input.c_str(), 0);
    return false;
}

unsigned char* ReadData()
{
    FILE* f;
    if (fopen_s(&f, input.c_str(), "rb") != 0 || f == nullptr)
        return nullptr;
    auto* const data = new uint8_t[fileLength];
    size_t index = 0;
    while (index < fileLength) {
        const auto count = fread(data + index, sizeof uint8_t, fileLength - index, f);
        std::cout << "read " << index << "+" << count << " - " << fileLength << std::endl;
        index += count;
    }
    return data;
}

int Initialise(const std::string& outputFile, AVCodecParameters* encoderParameters, AVFormatContext*& outputContainer,
               AVStream*& outStream)
{
    auto ret = avformat_alloc_output_context2(&outputContainer, nullptr, nullptr, outputFile.c_str());
    if (!outputContainer || ret < 0) // both check for error. Should always be consistent
    {
        std::cout << "Error: Unable to open output file: " << outputFile << " Error code: " << ret << std::endl;
        return -1;
    }

    outStream = avformat_new_stream(outputContainer, nullptr);

    if (outStream == nullptr)
    {
        std::cout << "Error: Unable to create video stream for output file" << std::endl;
        return -1;
    }

    if ((ret = avcodec_parameters_copy(outStream->codecpar, encoderParameters)) < 0)
    {
        std::cout << "Error: Unable to initialise writer context. Error code: " << ret << std::endl;
        return -1;
    }

    outStream->sample_aspect_ratio = encoderParameters->sample_aspect_ratio;

    if ((ret = avio_open(&outputContainer->pb, outputFile.c_str(), AVIO_FLAG_WRITE)) < 0)
    {
        std::cout << "Error: Unable open output for writing: " << outputFile.c_str() << ". Error code: " << ret <<
            std::endl;
        return -1;
    }

    if ((ret = avformat_init_output(outputContainer, nullptr)) < 0)
    {
        std::cout << "Error: Unable to initialise output. Error code: " << ret << std::endl;
        return -1;
    }

    av_dump_format(outputContainer, 0, outputFile.c_str(), 1);

    if ((ret = avformat_write_header(outputContainer, nullptr)) < 0)
    {
        std::cout << "Error: Unable to write header for output. Error code: " << ret << std::endl;
        return -1;
    }
    return 0;
}

int WritePacket(const nalUnit& nalUnit, unsigned char* const data, AVFormatContext* outputContainer, AVStream* outStream, int64_t& lastRawPts, int64_t& lastPts)
{

    AVPacket packet;
    auto ret = av_new_packet(&packet, nalUnit.size);
    if (ret != 0)
    {
        std::cout << "Error: Unable to allocate packet: " << ret <<std::endl;
        return -1;
    }

    memcpy(packet.data, data + nalUnit.start, nalUnit.size);
    std::cout << nalUnit.start << " "<<(int)packet.data[0] << " " << (int)packet.data[1] << " " << (int)packet.data[2] << " " << (int)packet.data[3] << std::endl;
    const auto pts = av_rescale_q(nalUnit.pts, timebase, outStream->time_base);

    if (lastRawPts == pts)
        packet.pts = packet.dts = ++lastPts;
    else
        packet.pts = packet.dts = lastPts = lastRawPts = pts;
    packet.duration = 0;
    packet.stream_index = 0;
    packet.flags = 0;
    packet.pos = -1;

    if ((ret = av_write_frame(outputContainer, &packet)) < 0)
    {
        std::cout << "Error: Unable to write frame. Error code: " << ret << std::endl;
        return -1;
    }

    av_packet_unref(&packet);
    return 0;
}

int main()
{
    AVFormatContext* inputFormatContext = nullptr;
    if (LoadInputFormat(inputFormatContext))
        return -1;

    auto* const data = ReadData();
    if (data == nullptr) return -1;

    AVFormatContext* outputContainer;
    AVStream* outStream;

    auto ret = Initialise(output, inputFormatContext->streams[0]->codecpar, outputContainer, outStream);
    if (ret != 0)
        return ret;

    avformat_close_input(&inputFormatContext);

    int64_t lastRawPts = -1;
    int64_t lastPts = -1;
    for (const auto& nalUnit : nalUnits) {
        ret = WritePacket(nalUnit, data, outputContainer, outStream, lastRawPts, lastPts);
        if (ret != 0)
            return ret;
    }

    ret = av_write_frame(outputContainer, nullptr);
    if (ret < 0)
    {
        std::cout << "Error: Unable to flush frames. Error code: " << ret << std::endl;
        return -1;
    }

    av_write_trailer(outputContainer);
    avio_closep(&outputContainer->pb);

    if (outputContainer)
        avformat_free_context(outputContainer);

    return 0;
}