path: root/src/ffmpeg_backend.cpp

#if SPECIAL
#include "main.h"
#endif

// workaround to make libav error printing work
#ifdef av_err2str
#undef av_err2str
#include <string>
av_always_inline std::string av_err2string(int errnum) {
    char str[AV_ERROR_MAX_STRING_SIZE];
    return av_make_error_string(str, AV_ERROR_MAX_STRING_SIZE, errnum);
}
#define av_err2str(err) av_err2string(err).c_str()
#endif  // av_err2str

bool32 AV_TryFrame(av_info *AV, AVCodecContext *CodecContext, int32 *err, bool32 *EndOfFile, int StreamIndex)
{
    *err = av_read_frame(AV->FileFormatContext, AV->Packet);
    bool32 CheckForStream = (StreamIndex != -1) ? (AV->Packet->stream_index == StreamIndex) : 1;
    if (!(*err >= 0 && CheckForStream)) {
        av_packet_unref(AV->Packet);
        return 0;
    }

    if (*err < 0)
        *err = avcodec_send_packet(CodecContext, AV->Packet);
    else {
        *err = avcodec_send_packet(CodecContext, AV->Packet);
    }
    av_packet_unref(AV->Packet);

    if (*err == AVERROR_EOF)
    {
        *err = 0;
        *EndOfFile = true;
        av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_BACKWARD);
        avcodec_flush_buffers(CodecContext);
        return 0;
    }

    if (*err < 0)
    {
        fprintf(stderr, "Libav *error: (%s)\n", av_err2str(*err));
        Assert(0);
    }

    while (*err >= 0) {
        *err = avcodec_receive_frame(CodecContext, AV->Frame);
        if (*err == AVERROR_EOF) {
            *EndOfFile = true;
            av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_BACKWARD);
            avcodec_flush_buffers(CodecContext);
            *err = 0;
            break;
        } else if (*err == AVERROR(EAGAIN)) {
            *err = 0;
            break;
        } else if (*err < 0) {
            Assert(0);
        }
        return 1;
    }
    return 0;
}

void AV_IsFileSupported(char *filename, bool32 *IsVideo, bool32 *HasAudio)
{
    int32 err = 0;

    AVFormatContext *temp = avformat_alloc_context();
    err = avformat_open_input(&temp, filename, NULL, NULL);;

    // if (*err == AVERROR_) {
    // }

    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
            avformat_free_context(temp);
            return;
    }

    err = avformat_find_stream_info(temp,  NULL);

    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
            avformat_free_context(temp);
            return;
    }

    for (uint32 i = 0; i < temp->nb_streams; i++)
    {
        AVCodecParameters *LocalCodecParameters = NULL;
        LocalCodecParameters = temp->streams[i]->codecpar;
        const AVCodec *codec = avcodec_find_decoder(LocalCodecParameters->codec_id);
        if (LocalCodecParameters->codec_type == AVMEDIA_TYPE_AUDIO) {
            if (codec) {
                *HasAudio = true;
            } else {
                fprintf(stderr, "Codec not found.\n");
            }
        }
        if (LocalCodecParameters->codec_type == AVMEDIA_TYPE_VIDEO) {
            if (temp->streams[i]->duration == 1) {
                fprintf(stderr, "TODO(fox): Make single-frame FFmpeg work!\n");
            } else if (codec) {
                *IsVideo = true;
            } else {
                fprintf(stderr, "Codec not found.\n");
            }
        }
    }

    avformat_free_context(temp);
}

void AV_Dealloc(av_info *AV)
{
    Assert(AV);
    avformat_free_context(AV->FileFormatContext);
    avcodec_free_context(&AV->Video.CodecContext);
    avcodec_free_context(&AV->Audio.CodecContext);
    av_packet_free(&AV->Packet);
    av_frame_free(&AV->Frame);
};

void AV_InitStream(av_stream_info *Stream)
{
    int32 err = 0;
    Stream->Codec = avcodec_find_decoder(Stream->CodecParameters->codec_id);

    if (!Stream->Codec) {
        printf("Libav error: Suitable decoder could not be identified for codec.\n");
    }

    Stream->CodecContext = avcodec_alloc_context3(Stream->Codec);
    if (!Stream->CodecContext) {
        printf("Libav error: Decoder context allocation failed.");
    }
    err = avcodec_parameters_to_context(Stream->CodecContext, Stream->CodecParameters);
    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
    }
    avcodec_open2(Stream->CodecContext, Stream->Codec, NULL);
    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
    }

}

// The duration isn't always reported in AVStream, but seeking towards the end
// and advancing until we hit EOF seems to be accurate.
void AV_GetDuration(av_info *AV, av_stream_info *Stream, uint64 *Duration, real32 *SecondCount)
{
    if (Stream->Stream->duration > 0) {
        *Duration = Stream->Stream->duration;
    } else {
        avformat_seek_file(AV->FileFormatContext, -1, INT64_MIN, 10000000000, INT64_MAX, 0);

        int32 err = 0;
        bool32 EndOfFile = 0;
        uint64 TestDuration = 0;
        while (err >= 0) {
            if (AV_TryFrame(AV, Stream->CodecContext, &err, &EndOfFile, Stream->Index))
            {
                TestDuration = AV->Frame->pts;
            }
            av_frame_unref(AV->Frame);
            if (EndOfFile)
                break;
        }
        *Duration = TestDuration;
    }
    real32 FPS = (real32)AV->Video.Stream->r_frame_rate.num / AV->Video.Stream->r_frame_rate.den;
    if (Stream->Stream->nb_frames > 0) {
        *SecondCount = Stream->Stream->nb_frames / FPS;
    } else if (AV->Video.CodecContext) {
        Assert(AV->FileFormatContext->duration > 0);
        *SecondCount = (real32)AV->FileFormatContext->duration / 1000000LL;
    }
}


void AV_Init(block_source *Source, av_info *AV, memory *Memory)
{
    int32 err = 0;

    // enum AVHWDeviceType type;
    // while((type = av_hwdevice_iterate_types(type)) != AV_HWDEVICE_TYPE_NONE)
    //     printf("%s\n", av_hwdevice_get_type_name(type));

    // The two calls below theoretically shouldn't fail since we already tested them in IsFileSupported.

    AV->FileFormatContext = avformat_alloc_context();
    block_string *String = (block_string *)Memory_Block_AddressAtIndex(Memory, F_Strings, Source->Path_String_Index);
    err = avformat_open_input(&AV->FileFormatContext, String->Char, NULL, NULL);;
    if (err < 0) {
        fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
        Assert(0);
    }

    err = avformat_find_stream_info(AV->FileFormatContext,  NULL);

    if (err < 0) {
        fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
        Assert(0);
    }

    for (uint32 i = 0; i < AV->FileFormatContext->nb_streams; i++)
    {
        AVCodecParameters *LocalCodecParameters = NULL;
        LocalCodecParameters = AV->FileFormatContext->streams[i]->codecpar;
        if (LocalCodecParameters->codec_type == AVMEDIA_TYPE_AUDIO) {
            AV->Audio.Index = i;
            AV->Audio.CodecParameters = LocalCodecParameters;
            AV->Audio.Stream =  AV->FileFormatContext->streams[i];
        }
        if (LocalCodecParameters->codec_type == AVMEDIA_TYPE_VIDEO) {
            AV->Video.Index = i;
            AV->Video.CodecParameters = LocalCodecParameters;
            AV->Video.Stream =  AV->FileFormatContext->streams[i];
        }
        if (AV->Video.Stream && AV->Audio.Stream)
            break;
    }

    if (!AV->Video.CodecParameters && !AV->Audio.CodecParameters) {
        printf("Libav error: No tracks found.");
    }

    if (AV->Video.CodecParameters)
        AV_InitStream(&AV->Video);
    if (AV->Audio.CodecParameters)
        AV_InitStream(&AV->Audio);

    AV->Packet = av_packet_alloc();
    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
    }
    AV->Frame = av_frame_alloc();
    if (err < 0) {
            fprintf(stderr, "Libav error: (%s)\n", av_err2str(err));
    }

    if (AV->Video.CodecParameters) {
        Source->BytesPerPixel = 4;
        Source->Width = AV->Video.CodecContext->width;
        Source->Height = AV->Video.CodecContext->height;
        Source->FPS = (real32)AV->Video.Stream->r_frame_rate.num / AV->Video.Stream->r_frame_rate.den;
        AV_GetDuration(AV, &AV->Video, &AV->PTSDuration, &AV->SecondCount);
        AV->LastFrameRendered = -1;
        av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_BACKWARD);
    } else {
        AV_GetDuration(AV, &AV->Audio, &AV->PTSDuration, &AV->SecondCount);
        av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_BACKWARD);
    }
    avcodec_flush_buffers(AV->Video.CodecContext);
    avcodec_flush_buffers(AV->Audio.CodecContext);
};

uint32 AV_AudioTest(av_info *AV, void *Data, uint32 Size, real32 FPS, int32 InitialFrameToSeek)
{
    AVChannelLayout ChannelLayout = {};
    av_channel_layout_default(&ChannelLayout, 2);

    av_stream_info *Stream = &AV->Audio;
    Assert(Stream->CodecContext->ch_layout.nb_channels == 2);

    real32 TotalFrames = AV->SecondCount * FPS;
    int64 SeekPTS = (int64)(((real64)InitialFrameToSeek / TotalFrames) * AV->PTSDuration);
    int64 AveragePTS = AV->PTSDuration / TotalFrames;

    int32 err = 0;
    bool32 EndOfFile = 0;
    while (err >= 0) {
        if (AV_TryFrame(AV, Stream->CodecContext, &err, &EndOfFile, Stream->Index))
        {
            if (AV->Frame->pts < Stream->Stream->start_time || AV->Frame->pts < SeekPTS) {
                av_frame_unref(AV->Frame);
                continue;
            }
            // NOTE(fox): The docs say Frame->format corresponds to AVSampleFormat, but ffplay doesn't need the cast.
            SwrContext *Swr = swr_alloc();
            swr_alloc_set_opts2(&Swr, &ChannelLayout, AV_SAMPLE_FMT_S16, 48000,
                                &AV->Frame->ch_layout, (AVSampleFormat)AV->Frame->format, AV->Frame->sample_rate,
                                0, NULL);
            swr_init(Swr);

            uint8 **Out = (uint8 **)&Data;
            const uint8 **In = (const uint8_t **)AV->Frame->extended_data;
            int len = swr_convert(Swr, Out, Size, In, AV->Frame->nb_samples);
            if (len < 0) {
                fprintf(stderr, "Libav error: (%s)\n", av_err2str(len));
                av_log(NULL, AV_LOG_ERROR, "swr_convert() failed\n");
                Assert(0);
            }
            /*
            int BytesPerSample = av_get_bytes_per_sample(Stream->CodecContext->sample_fmt);
            Assert(BytesPerSample == 2 || BytesPerSample == 4);
            int Offset = (BytesPerSample == 4) ? 2 : 1;
            for (int i = 0; i < (Stream->Frame->nb_samples * 2); i+=2) {
                uint16 *Channel_L = (uint16 *)Data + i + 0;
                uint16 *Channel_R = (uint16 *)Data + i + 1;
                *Channel_L = (*(uint16 *)(Stream->Frame->data[0] + BytesPerSample*i)) / Offset;
                *Channel_R = (*(uint16 *)(Stream->Frame->data[1] + BytesPerSample*i)) / Offset;
                int aa = 0;
            }
            */
            return len;
        }
        av_frame_unref(AV->Frame);
    }
    return 0;
}

void AV_SeekAudio(av_info *AV, real32 FPS, int32 FrameToSeek)
{
    Assert(FrameToSeek > -1)
    int64 SeekSeconds = (int64)(FrameToSeek / (int32)(FPS + 0.5f) * AV_TIME_BASE);
    av_seek_frame(AV->FileFormatContext, -1, SeekSeconds, AVSEEK_FLAG_BACKWARD);
}


void AV_LoadVideoFrame(memory *Memory, block_source *Source, av_info *AV, int32 FrameToSeek, void *Buffer)
{
    int32 *LastFrameRendered = &AV->LastFrameRendered;

    int32 err = 0;

    Assert(FrameToSeek > -1)

    // NOTE(fox): The decoder automatically advances to the next frame, so we
    // don't need to call av_seek_frame under normal playback.
    // This function only seeks to the nearest "keyframe."

    if (*LastFrameRendered != FrameToSeek - 1) {
        int64 SeekSeconds = (int64)(FrameToSeek / (int32)(Source->FPS + 0.5f) * AV_TIME_BASE);
        av_seek_frame(AV->FileFormatContext, -1, SeekSeconds, AVSEEK_FLAG_BACKWARD);
        printf("Seek activated\n");
    } else if (*LastFrameRendered < 0) {
        av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_BACKWARD);
    }

    *LastFrameRendered = FrameToSeek;

    real32 TotalFrames = AV->SecondCount * Source->FPS;
    int64 SeekPTS = (int64)(((real64)FrameToSeek / TotalFrames) * AV->PTSDuration);
    int64 AveragePTS = AV->PTSDuration / TotalFrames;

    bool32 EndOfFile = 0;
    while (err >= 0) {
        if (AV_TryFrame(AV, AV->Video.CodecContext, &err, &EndOfFile, AV->Video.Index))
        {
            // The first frame that gets loaded isn't always the actual
            // first frame, so we need to check until it's correct.
            if (FrameToSeek == 0 && AV->Frame->pts < AV->Video.Stream->start_time) {
                av_frame_unref(AV->Frame);
                // printf("NON-START: avg: %li, real pts: %li", SeekPTS, AV->VideoFrame->pts);
                continue;
            }

            int64 Difference = AV->Frame->pts - SeekPTS;
            if (abs(Difference) < AveragePTS)
            {
                if (AV->PreviousPTS == -1) {
                    AV->PreviousPTS = AV->Frame->pts;
                    // printf("avg: %li, real pts: %li, difference: %li\n", SeekPTS, AV->Frame->pts, Difference);
                } else {
                    // printf("avg: %li, real pts: %li, difference: %li difference from last pts: %li\n", SeekPTS, AV->Frame->pts, AV->Frame->pts - SeekPTS, AV->Frame->pts - AV->PreviousPTS);
                    AV->PreviousPTS = AV->Frame->pts;
                }

                uint16 Width = Source->Width;
                uint16 Height = Source->Height;
                uint16 BytesPerPixel = Source->BytesPerPixel;
                int32 Pitch = Width*BytesPerPixel;

                int out_linesize[4] = { Pitch, Pitch, Pitch, Pitch };
                uint8 *dst_data[4] = { (uint8 *)Buffer, (uint8 *)Buffer + Width*Height*BytesPerPixel,
                                       (uint8 *)Buffer + Width*Height*BytesPerPixel*2, (uint8 *)Buffer + Width*Height*BytesPerPixel*3 };

                // NOTE(fox): This function will be replaced in the future.
                AV->RGBContext = sws_getContext(AV->Frame->width, AV->Frame->height, (AVPixelFormat)AV->Frame->format,
                                                AV->Frame->width, AV->Frame->height, AV_PIX_FMT_RGBA, SWS_BILINEAR,
                                                NULL, NULL, NULL);

                if(!AV->RGBContext) {
                    printf("Libav error: SwsContext creation failed.");
                }

                sws_scale(AV->RGBContext, AV->Frame->data, AV->Frame->linesize, 0, AV->Frame->height,
                          dst_data, out_linesize);

                av_frame_unref(AV->Frame);
                return;
            }
            else
            {
                // If this gets printed when not seeking, a frame has been skipped!
                // printf("FRAME SKIP: avg: %li, real pts: %li, difference: %li\n", SeekPTS, AV->VideoFrame->pts, Difference);
            }
            av_frame_unref(AV->Frame);
        }
    }
    Assert(0);
}

#if 0

cached_bitmap * AV_LoadStill(block_source *Source, layer_bitmap_info *BitmapInfo, memory *Memory)
{
    av_info *AV = (av_info *)BitmapInfo->AVInfo;
    int32 *CurrentlyRenderedFrame = &BitmapInfo->CurrentFrame;

    int32 err = 0;

    *CurrentlyRenderedFrame = 0;

    av_seek_frame(AV->FileFormatContext, -1, 0, AVSEEK_FLAG_FRAME);

    while (err >= 0) {
        if (AV_TryFrame(AV, &err))
        {
            uint16 Width = Source->Width;
            uint16 Height = Source->Height;
            uint16 BytesPerPixel = Source->BytesPerPixel;
            int32 Pitch = Width*BytesPerPixel;

            cached_bitmap *Bitmap = Memory_RollingBitmap(Memory, Source, 0);
            void *Buffer = Bitmap->Data;

            int out_linesize[4] = { Pitch, Pitch, Pitch, Pitch };
            uint8 *dst_data[4] = { (uint8 *)Buffer, (uint8 *)Buffer + Width*Height*BytesPerPixel,
                                   (uint8 *)Buffer + Width*Height*BytesPerPixel*2, (uint8 *)Buffer + Width*Height*BytesPerPixel*3 };

            // NOTE(fox): This function will be replaced in the future.
            AV->RGBContext = sws_getContext(AV->VideoFrame->width, AV->VideoFrame->height, (AVPixelFormat)AV->VideoFrame->format,
                                            AV->VideoFrame->width, AV->VideoFrame->height, AV_PIX_FMT_RGBA, SWS_BILINEAR,
                                            NULL, NULL, NULL);

            if(!AV->RGBContext) {
                printf("Libav error: SwsContext creation failed.");
            }

            sws_scale(AV->RGBContext, AV->VideoFrame->data, AV->VideoFrame->linesize, 0, AV->VideoFrame->height,
                      dst_data, out_linesize);

            av_frame_unref(AV->VideoFrame);

            Assert(BitmapInfo->BitmapBuffer);
            return Bitmap;
        }
        av_frame_unref(AV->VideoFrame);
    }
    return 0;
}

#endif