/*
 * src/audio.c
 * https://gitlab.com/bztsrc/assetpack
 *
 * Copyright (C) 2024 bzt, MIT license
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL ANY
 * DEVELOPER OR DISTRIBUTOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
 * IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *
 * @brief Handle audio formats
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "audiocvt.h"

#define FREQ 44100

#define MINIMP3_NO_SIMD
#define MINIMP3_NONSTANDARD_BUT_LOGICAL
#define MINIMP3_IMPLEMENTATION
#include "minimp3.h"
extern int stb_vorbis_decode_memory(const unsigned char *mem, int len, int *channels, int *sample_rate, short **output);

/*** lots of stuff that should be in minimp3.h really ***/
static void mp3dec_skip_id3v1(const uint8_t *buf, size_t *pbuf_size)
{
    size_t buf_size = *pbuf_size;
    const uint8_t *tag;
    uint32_t tag_size;
#ifndef MINIMP3_NOSKIP_ID3V1
    if (buf_size >= 128 && !memcmp(buf + buf_size - 128, "TAG", 3))
    {
        buf_size -= 128;
        if (buf_size >= 227 && !memcmp(buf + buf_size - 227, "TAG+", 4))
            buf_size -= 227;
    }
#endif
#ifndef MINIMP3_NOSKIP_APEV2
    if (buf_size > 32 && !memcmp(buf + buf_size - 32, "APETAGEX", 8))
    {
        buf_size -= 32;
        tag = buf + buf_size + 8 + 4;
        tag_size = (uint32_t)(tag[3] << 24) | (tag[2] << 16) | (tag[1] << 8) | tag[0];
        if (buf_size >= tag_size)
            buf_size -= tag_size;
    }
#endif
    *pbuf_size = buf_size;
}

static size_t mp3dec_skip_id3v2(const uint8_t *buf, size_t buf_size)
{
#define MINIMP3_ID3_DETECT_SIZE 10
#ifndef MINIMP3_NOSKIP_ID3V2
    if (buf_size >= MINIMP3_ID3_DETECT_SIZE && !memcmp(buf, "ID3", 3) && !((buf[5] & 15) || (buf[6] & 0x80) || (buf[7] & 0x80) || (buf[8] & 0x80) || (buf[9] & 0x80)))
    {
        size_t id3v2size = (((buf[6] & 0x7f) << 21) | ((buf[7] & 0x7f) << 14) | ((buf[8] & 0x7f) << 7) | (buf[9] & 0x7f)) + 10;
        if ((buf[5] & 16))
            id3v2size += 10; /* footer */
        return id3v2size;
    }
#endif
    return 0;
}

static void mp3dec_skip_id3(const uint8_t **pbuf, size_t *pbuf_size)
{
    uint8_t *buf = (uint8_t *)(*pbuf);
    size_t buf_size = *pbuf_size;
    size_t id3v2size = mp3dec_skip_id3v2(buf, buf_size);
    if (id3v2size)
    {
        if (id3v2size >= buf_size)
            id3v2size = buf_size;
        buf      += id3v2size;
        buf_size -= id3v2size;
    }
    mp3dec_skip_id3v1(buf, &buf_size);
    *pbuf = (const uint8_t *)buf;
    *pbuf_size = buf_size;
}

static int mp3dec_check_vbrtag(const uint8_t *frame, int frame_size, uint32_t *frames, int *delay, int *padding)
{
    static const char g_xing_tag[4] = { 'X', 'i', 'n', 'g' };
    static const char g_info_tag[4] = { 'I', 'n', 'f', 'o' };
    const uint8_t *tag;
    int flags;
#define FRAMES_FLAG     1
#define BYTES_FLAG      2
#define TOC_FLAG        4
#define VBR_SCALE_FLAG  8
    /* Side info offsets after header:
    /                Mono  Stereo
    /  MPEG1          17     32
    /  MPEG2 & 2.5     9     17*/
    bs_t bs[1];
    L3_gr_info_t gr_info[4];
    bs_init(bs, frame + HDR_SIZE, frame_size - HDR_SIZE);
    if (HDR_IS_CRC(frame))
        get_bits_mp3(bs, 16);
    if (L3_read_side_info(bs, gr_info, frame) < 0)
        return 0; /* side info corrupted */

    tag = frame + HDR_SIZE + bs->pos/8;
    if (memcmp(g_xing_tag, tag, 4) && memcmp(g_info_tag, tag, 4))
        return 0;
    flags = tag[7];
    if (!((flags & FRAMES_FLAG)))
        return -1;
    tag += 8;
    *frames = (uint32_t)(tag[0] << 24) | (tag[1] << 16) | (tag[2] << 8) | tag[3];
    tag += 4;
    if (flags & BYTES_FLAG)
        tag += 4;
    if (flags & TOC_FLAG)
        tag += 100;
    if (flags & VBR_SCALE_FLAG)
        tag += 4;
    *delay = *padding = 0;
    if (*tag)
    {   /* extension, LAME, Lavc, etc. Should be the same structure. */
        tag += 21;
        if (tag - frame + 14 >= frame_size)
            return 0;
        *delay   = ((tag[0] << 4) | (tag[1] >> 4)) + (528 + 1);
        *padding = (((tag[1] & 0xF) << 8) | tag[2]) - (528 + 1);
    }
    return 1;
}

typedef struct
{
    int16_t *buffer;
    size_t samples; /* channels included, byte size = samples*sizeof(mp3d_sample_t) */
    int channels, hz, layer, avg_bitrate_kbps;
} mp3dec_file_info_t;

#define MP3D_E_MEMORY   -1
#define MP3D_E_PARAM    -2
#define MP3D_E_DECODE   -3
int mp3dec_load_buf(mp3dec_t *dec, uint8_t *buf, size_t buf_size, mp3dec_file_info_t *info)
{
    uint32_t iframes;
    int i, delay, padding, free_format_bytes = 0, frame_size = 0;
    const uint8_t *hdr;
    uint64_t detected_samples = 0;
    int to_skip = 0, ret = 0;
    mp3dec_frame_info_t frame_info;
    size_t allocated;
    size_t avg_bitrate_kbps = 0, frames = 0;
    mp3d_sample_t *alloc_buf;
    size_t skip;
    int samples;

    memset(info, 0, sizeof(*info));
    memset(&frame_info, 0, sizeof(frame_info));
    if (!dec || !buf || !info || (size_t)-1 == buf_size)
        return MP3D_E_PARAM;

    /* skip id3 */
    mp3dec_skip_id3((const uint8_t **)&buf, &buf_size);
    if (!buf_size)
        return 0;
    /* try to make allocation size assumption by first frame or vbr tag */
    mp3dec_init(dec);
    do
    {
        free_format_bytes = 0; frame_size = 0;
        i = mp3d_find_frame(buf, buf_size, &free_format_bytes, &frame_size);
        buf      += i;
        buf_size -= i;
        hdr = buf;
        if (i && !frame_size)
            continue;
        if (!frame_size)
            return 0;
        frame_info.channels = HDR_IS_MONO(hdr) ? 1 : 2;
        frame_info.hz = hdr_sample_rate_hz(hdr);
        frame_info.layer = 4 - HDR_GET_LAYER(hdr);
        frame_info.bitrate_kbps = hdr_bitrate_kbps(hdr);
        frame_info.frame_bytes = frame_size;
        samples = hdr_frame_samples(hdr)*frame_info.channels;
        if (3 != frame_info.layer)
            break;
        ret = mp3dec_check_vbrtag(hdr, frame_size, &iframes, &delay, &padding);
        if (ret > 0)
        {
            padding *= frame_info.channels;
            to_skip = delay*frame_info.channels;
            detected_samples = samples*(uint64_t)iframes;
            if (detected_samples >= (uint64_t)to_skip)
                detected_samples -= to_skip;
            if (padding > 0 && detected_samples >= (uint64_t)padding)
                detected_samples -= padding;
            if (!detected_samples)
                return 0;
        }
        if (ret)
        {
            buf      += frame_size;
            buf_size -= frame_size;
        }
        break;
    } while(1);
    allocated = MINIMP3_MAX_SAMPLES_PER_FRAME*sizeof(mp3d_sample_t);
    if (detected_samples)
        allocated += detected_samples*sizeof(mp3d_sample_t);
    else
        allocated += (buf_size/frame_info.frame_bytes)*samples*sizeof(mp3d_sample_t);
    info->buffer = (mp3d_sample_t*)malloc(allocated);
    if (!info->buffer)
        return MP3D_E_MEMORY;
    /* save info */
    info->channels = frame_info.channels;
    info->hz       = frame_info.hz;
    info->layer    = frame_info.layer;
    /* decode all frames */
    avg_bitrate_kbps = 0; frames = 0;
    do
    {
        if ((allocated - info->samples*sizeof(mp3d_sample_t)) < MINIMP3_MAX_SAMPLES_PER_FRAME*sizeof(mp3d_sample_t))
        {
            allocated *= 2;
            alloc_buf = (mp3d_sample_t*)realloc(info->buffer, allocated);
            if (!alloc_buf)
                return MP3D_E_MEMORY;
            info->buffer = alloc_buf;
        }
        samples = mp3dec_decode_frame(dec, buf, MINIMP3_MIN(buf_size, (size_t)INT_MAX), info->buffer + info->samples, &frame_info);
        buf      += frame_info.frame_bytes;
        buf_size -= frame_info.frame_bytes;
        if (samples)
        {
            if (info->hz != frame_info.hz || info->layer != frame_info.layer)
            {
                ret = MP3D_E_DECODE;
                break;
            }
            if (info->channels && info->channels != frame_info.channels)
            {
                ret = MP3D_E_DECODE;
                break;
            }
            samples *= frame_info.channels;
            if (to_skip)
            {
                skip = MINIMP3_MIN(samples, to_skip);
                to_skip -= skip;
                samples -= skip;
                memmove(info->buffer, info->buffer + skip, samples*sizeof(mp3d_sample_t));
            }
            info->samples += samples;
            avg_bitrate_kbps += frame_info.bitrate_kbps;
            frames++;
        }
    } while (frame_info.frame_bytes);
    if (detected_samples && info->samples > detected_samples)
        info->samples = detected_samples; /* cut padding */
    /* reallocate to normal buffer size */
    if (allocated != info->samples*sizeof(mp3d_sample_t) && info->samples > 1)
    {
        alloc_buf = (mp3d_sample_t*)realloc(info->buffer, info->samples*sizeof(mp3d_sample_t));
        if (!alloc_buf && info->samples)
            return MP3D_E_MEMORY;
        info->buffer = alloc_buf;
    }
    if (frames)
        info->avg_bitrate_kbps = avg_bitrate_kbps/frames;
    return ret;
}

typedef struct {
    char str_riff[4];
    int wav_size;
    char str_wave[8];
    int fmt_chunk_size;
    short audio_format;
    short channels;
    int sample_rate;
    int byte_rate;
    short frame_size;
    short bit_depth;
    char str_data[4];
    int data_bytes;
} wav_t;

/**
 * Decode an audio asset into a 16-bit stereo PCM data
 */
uint8_t *audio_decode(uint8_t *buf, uint64_t size, int ismusic, uint64_t *len)
{
    wav_t *wav = (wav_t*)buf;
    SDL_AudioCVT cvt;
    mp3dec_file_info_t info;
    mp3dec_t mp3d;
    uint8_t *ret = NULL;
    int ch;

    if(!buf || !size || !len) return NULL;
    *len = 0; memset(&info, 0, sizeof(info));
    if(!memcmp(buf, "RIFF", 4) && !memcmp(buf + 8, "WAVE", 4) && wav->data_bytes && wav->channels && wav->bit_depth) {
        ch = ismusic ? wav->channels : 1;
        SDL_BuildAudioCVT(&cvt, wav->bit_depth == 8 ? AUDIO_U8 : (wav->bit_depth == 16 ? AUDIO_S16 :
            (wav->audio_format == 3 ? AUDIO_F32 : AUDIO_S32)), wav->channels, wav->sample_rate, AUDIO_S16, ch, FREQ);
        if(sizeof(wav_t) + wav->data_bytes > size) wav->data_bytes = size - sizeof(wav_t);
        cvt.len = wav->data_bytes;
        info.channels = wav->channels;
        info.samples = (wav->data_bytes / (wav->channels * wav->bit_depth / 8));
        info.buffer = (int16_t*)(buf + sizeof(wav_t));
        goto doconv;
    } else
    if(mp3dec_load_buf(&mp3d, buf, size, &info) >= 0 && info.channels && info.samples && info.buffer) {
        ch = ismusic ? info.channels : 1;
        SDL_BuildAudioCVT(&cvt, AUDIO_S16, info.channels, info.hz, AUDIO_S16, ch, FREQ);
        cvt.len = info.samples * 2;
        goto doconv;
    } else
    if((info.samples = stb_vorbis_decode_memory(buf, size, &info.channels, &info.hz, &info.buffer)) && info.buffer) {
        if(!info.channels || !info.samples) { free(info.buffer); return NULL; }
        info.samples *= info.channels;
        ch = ismusic ? info.channels : 1;
        SDL_BuildAudioCVT(&cvt, AUDIO_S16, info.channels, info.hz, AUDIO_S16, ch, FREQ);
        cvt.len = info.samples * 2;
doconv: if(!cvt.len_mult) cvt.len_mult = 1;
        if(!(ret = (uint8_t*)malloc(sizeof(wav_t) + cvt.len * cvt.len_mult))) {
            if(info.buffer && info.buffer != (int16_t*)(buf + sizeof(wav_t))) free(info.buffer);
            return NULL;
        }
        /* convert to 16 bit stereo / mono 44100Hz */
        cvt.buf = ret + sizeof(wav_t);
        memcpy(cvt.buf, info.buffer, cvt.len);
        if(info.buffer && info.buffer != (int16_t*)(buf + sizeof(wav_t))) free(info.buffer);
        if(SDL_ConvertAudio(&cvt)) { free(ret); return NULL; }
        memset(ret, 0, sizeof(wav_t));
        /* add wav header */
        wav = (wav_t*)ret;
        memcpy(&wav->str_riff, "RIFF", 4);
        wav->wav_size = wav->data_bytes = info.samples * sizeof(int16_t);
        wav->wav_size += sizeof(wav_t);
        memcpy(&wav->str_wave, "WAVEfmt ", 8);
        wav->fmt_chunk_size = 16;
        wav->audio_format = 1;
        wav->channels = ch;
        wav->sample_rate = FREQ;
        wav->byte_rate = FREQ * ch * sizeof(int16_t);
        wav->frame_size = ch * sizeof(int16_t);
        wav->bit_depth = sizeof(int16_t) << 3;
        memcpy(&wav->str_data, "data", 4);
        *len = wav->wav_size;
    }
    return ret;
}
