#ifdef CAMU_HAVE_FFMPEG
#include "../codec/ffmpeg/demuxer.h"
#include "../codec/ffmpeg/decoder.h"
#include "../codec/ffmpeg/resampler.h"
#endif

#include "process.h"

struct raw_samples_data {
    struct camu_resampler *resampler;
    u32 channels;
    u32 bps;
    struct nn_buffer buf;
    u32 index;
};

#ifdef CAMU_HAVE_FFMPEG
static u32 append_samples(struct raw_samples_data *raw, u8 *data, u32 sample_count)
{
    u32 bytes = sample_count * raw->channels * raw->bps;
    nn_buffer_write(&raw->buf, data, raw->index, bytes);
    return bytes;
}

static void data_callback(void *userdata, struct camu_codec_frame *frame)
{
    struct raw_samples_data *raw = (struct raw_samples_data *)userdata;
    u8 **data = frame->av.frame->data;
    s32 sample_count = frame->av.frame->nb_samples;
    sample_count = raw->resampler->convert(raw->resampler, (const u8 **)data, sample_count);
    data = raw->resampler->get_data(raw->resampler);
    raw->index += append_samples(raw, data[0], (u32)sample_count);
}

static void flush_resampler(struct camu_resampler *resampler, struct raw_samples_data *raw)
{
    s32 sample_count = resampler->flush(resampler);
    u8 **data = resampler->get_data(resampler);
    raw->index += append_samples(raw, data[0], (u32)sample_count);
}

bool lia_prepare_visual_data(struct cch_handle *handle, struct lia_visual_data *data)
{
    struct camu_demuxer *demux = camu_ff_demuxer_create();
    struct camu_decoder *dec = camu_ff_decoder_create();
    struct camu_resampler *resampler = camu_ff_resampler_create();
    data->samples = NULL;

    if (!demux->init(demux, handle)) {
        goto out;
    }

    struct camu_codec_stream *stream = NULL;
    al_array_foreach_ptr(demux->streams, i, stream) {
        if (stream->type == CAMU_STREAM_AUDIO) {
            demux->subscribed = 1 << stream->index;
            break;
        }
    }
    if (!stream || demux->subscribed == 0) goto out;

    struct camu_resampler_format fmt;
    AVCodecParameters *codecpar = stream->av.stream->codecpar;
    fmt.in.format = codecpar->format;
    fmt.req.format = CAMU_SAMPLE_FORMAT_S16;
    fmt.in.sample_rate = codecpar->sample_rate;
    fmt.req.sample_rate = codecpar->sample_rate;
    camu_copy_channel_layout(&fmt.in.channel_layout, &AV_CODECPAR_CHANNEL_LAYOUT(codecpar));
    camu_copy_channel_layout(&fmt.req.channel_layout, &AV_CODECPAR_CHANNEL_LAYOUT(codecpar));
    fmt.in.channel_count = AV_CODECPAR_CHANNELS(codecpar);
    fmt.req.channel_count = AV_CODECPAR_CHANNELS(codecpar);

    struct raw_samples_data raw;
    raw.resampler = resampler;
    raw.channels = fmt.req.channel_count;
    raw.bps = camu_audio_format_bytes_per_sample(&fmt.req);
    nn_buffer_init(&raw.buf);
    raw.index = 0;

    if (!dec->init(dec, NULL, stream, data_callback, &raw)) {
        goto out;
    }

    fmt.resampler_needed = true;
    if (!resampler->init(resampler, &fmt)) {
        goto out;
    }

    AVPacket *pkt = av_packet_alloc();
    struct camu_codec_packet packet = { .av.pkt = pkt };
    s32 status;
    do {
        status = demux->get_packet(demux, &packet);
        if (status == CAMU_OK) {
            dec->push_av_packet(dec, pkt);
            dec->process(dec);
            av_packet_unref(pkt);
        }
    } while (status == CAMU_OK);
    dec->flush(dec);
    flush_resampler(resampler, &raw);

    data->size = raw.index;
    data->samples = raw.buf.data;
    data->count = raw.index / raw.channels / raw.bps;
    camu_audio_format_copy(&data->fmt, &fmt.req);

out:
    demux->free(&demux);
    dec->free(&dec);
    resampler->free(&resampler);

    return data->samples != NULL;
}
#endif