/* ========================== *
 * Microsoft Media Foundation (MMF) mp3 decoder
 * ========================== */

#include "mp3.h"
#include "arena.h"
#include "playback.h"



#pragma warning(push, 0)
# define COBJMACROS
# define WIN32_LEAN_AND_MEAN
# define UNICODE
# include <Windows.h>
# include <uuids.h>
# include <mfapi.h>
# include <mfidl.h>
# include <mfreadwrite.h>
# include <Shlwapi.h>
#pragma warning(pop)

#pragma comment(lib, "mfplat")
#pragma comment(lib, "mfreadwrite")
#pragma comment(lib, "shlwapi")

struct mp3_decode_result mp3_decode(struct arena *arena, struct string encoded, u32 flags)
{
    struct mp3_decode_result res = ZI;

    u64 sample_rate = PLAYBACK_SAMPLE_RATE;
    u64 bytes_per_sample = 2;

    u64 channel_count;
    u32 channel_mask;
    if (flags & MP3_DECODE_FLAG_STEREO) {
        channel_count = 2;
        channel_mask = SPEAKER_FRONT_LEFT | SPEAKER_FRONT_RIGHT;
    } else {
        channel_count = 1;
        channel_mask = SPEAKER_FRONT_CENTER;
    }

    /* ========================== *
     * Startup
     * ========================== */

    MFStartup(MF_VERSION, MFSTARTUP_LITE);

    /* Create IStream from encoded string */
    IStream *i_stream = SHCreateMemStream(encoded.text, encoded.len);

    /* Create IMFByteStream from IStream */
    IMFByteStream *byte_stream = 0;
    MFCreateMFByteStreamOnStream(i_stream, &byte_stream);

    /* Create reader from IMFByteStream */
    IMFSourceReader *reader;
    MFCreateSourceReaderFromByteStream(byte_stream, 0, &reader);

    /* ========================== *
     * Get media type
     * ========================== */

    /* Read only first audio stream */
    IMFSourceReader_SetStreamSelection(reader, (DWORD)MF_SOURCE_READER_ALL_STREAMS, 0);
    IMFSourceReader_SetStreamSelection(reader, (DWORD)MF_SOURCE_READER_FIRST_AUDIO_STREAM, 1);

    WAVEFORMATEXTENSIBLE format = {
        .Format = {
            .wFormatTag = WAVE_FORMAT_EXTENSIBLE,
            .nChannels = (WORD)channel_count,
            .nSamplesPerSec = (WORD)sample_rate,
            .nAvgBytesPerSec = (DWORD)(sample_rate * channel_count * bytes_per_sample),
            .nBlockAlign = (WORD)(channel_count * bytes_per_sample),
            .wBitsPerSample = (WORD)(8 * bytes_per_sample),
            .cbSize = sizeof(format) - sizeof(format.Format)
        },
        .Samples.wValidBitsPerSample = 8 * bytes_per_sample,
        .dwChannelMask = channel_mask,
        .SubFormat = MEDIASUBTYPE_PCM
    };

    /* Media Foundation in Windows 8+ allows reader to convert output to different format than native */
    IMFMediaType *type;
    MFCreateMediaType(&type);
    MFInitMediaTypeFromWaveFormatEx(type, &format.Format, sizeof(format));
    IMFSourceReader_SetCurrentMediaType(reader, (DWORD)MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, type);
    IMFMediaType_Release(type);

    /* ========================== *
     * Read
     * ========================== */

    res.pcm.samples = arena_push_dry(arena, i16);
    u64 sample_bytes_read = 0;
    for (;;) {
        IMFSample *sample;
        DWORD sample_flags = 0;
        HRESULT hr = IMFSourceReader_ReadSample(reader, (DWORD)MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, 0, &sample_flags, 0, &sample);
        if (FAILED(hr)) {
            break;
        }

        /* Check if done */
        if (sample_flags & MF_SOURCE_READERF_ENDOFSTREAM) {
            res.success = 1;
            break;
        }
        ASSERT(sample_flags == 0);

        /* Read samples */
        IMFMediaBuffer *buffer;
        IMFSample_ConvertToContiguousBuffer(sample, &buffer);

        BYTE *src;
        DWORD size_bytes;
        IMFMediaBuffer_Lock(buffer, &src, 0, &size_bytes);
        {
            i16 *dst = arena_push_array_no_zero(arena, i16, (size_bytes + 1) >> 1);
            MEMCPY(dst, src, size_bytes);
            sample_bytes_read += size_bytes;
        }
        IMFMediaBuffer_Unlock(buffer);

        IMediaBuffer_Release(buffer);
        IMFSample_Release(sample);
    }

    res.pcm.count = sample_bytes_read / bytes_per_sample;

    /* ========================== *
     * Cleanup
     * ========================== */

    IMFSourceReader_Release(reader);
    IMFByteStream_Close(byte_stream);
    IStream_Release(i_stream);
    MFShutdown();

    return res;
}