// B54 shim. LGPL-2.1-or-later. See LibavShim.h.
#include "LibavShim.h"
#include <libavcodec/avcodec.h>
#include <libavutil/channel_layout.h>
#include <libavutil/version.h>
#include <math.h>
#include <stdlib.h>

struct ls_decoder {
    AVCodecContext *ctx;
    AVPacket *pkt;
    AVFrame *frame;
};

const char *ls_version(void) { return av_version_info(); }

ls_decoder *ls_open(int codec, int sample_rate, int channels) {
    enum AVCodecID id = codec == LS_CODEC_TRUEHD ? AV_CODEC_ID_TRUEHD
                      : codec == LS_CODEC_DTS ? AV_CODEC_ID_DTS : AV_CODEC_ID_NONE;
    const AVCodec *c = avcodec_find_decoder(id);
    if (!c) return NULL;
    ls_decoder *d = calloc(1, sizeof *d);
    if (!d) return NULL;
    d->ctx = avcodec_alloc_context3(c);
    d->pkt = av_packet_alloc();
    d->frame = av_frame_alloc();
    if (!d->ctx || !d->pkt || !d->frame) { ls_close(d); return NULL; }
    d->ctx->sample_rate = sample_rate;
    if (channels > 0) av_channel_layout_default(&d->ctx->ch_layout, channels);
    if (avcodec_open2(d->ctx, c, NULL) < 0) { ls_close(d); return NULL; }
    return d;
}

void ls_reset(ls_decoder *d) { if (d && d->ctx) avcodec_flush_buffers(d->ctx); }

void ls_close(ls_decoder *d) {
    if (!d) return;
    avcodec_free_context(&d->ctx);
    av_packet_free(&d->pkt);
    av_frame_free(&d->frame);
    free(d);
}

// Right-justified, sign-extended integer at `bits`. Integer formats keep the decoder's own
// depth; float (lossy DTS core) becomes 24-bit with clamping.
static int32_t sample_at(const AVFrame *f, int ch, int i, int bits) {
    int channels = f->ch_layout.nb_channels;
    switch (f->format) {
    case AV_SAMPLE_FMT_S16:  return ((const int16_t *)f->data[0])[i * channels + ch];
    case AV_SAMPLE_FMT_S16P: return ((const int16_t *)f->data[ch])[i];
    case AV_SAMPLE_FMT_S32:  return ((const int32_t *)f->data[0])[i * channels + ch] >> (32 - bits);
    case AV_SAMPLE_FMT_S32P: return ((const int32_t *)f->data[ch])[i] >> (32 - bits);
    case AV_SAMPLE_FMT_FLT:
    case AV_SAMPLE_FMT_FLTP: {
        float x = f->format == AV_SAMPLE_FMT_FLT ? ((const float *)f->data[0])[i * channels + ch]
                                                 : ((const float *)f->data[ch])[i];
        long v = lrintf(x * 8388608.0f);
        return (int32_t)(v > 8388607 ? 8388607 : v < -8388608 ? -8388608 : v);
    }
    default: return 0;
    }
}

static int bits_for(const AVCodecContext *ctx, int format) {
    if (format == AV_SAMPLE_FMT_S16 || format == AV_SAMPLE_FMT_S16P) return 16;
    if (format == AV_SAMPLE_FMT_FLT || format == AV_SAMPLE_FMT_FLTP) return 24;
    return ctx->bits_per_raw_sample == 16 ? 16 : 24;
}

int ls_decode(ls_decoder *d, const uint8_t *data, int size,
              int32_t *out, int out_capacity_frames,
              int *out_frames, int *out_channels, uint64_t *out_layout_mask, int *out_bits) {
    if (!d || !out_frames || !out_channels || !out_layout_mask || !out_bits) return LS_ERR_ARG;
    *out_frames = 0; *out_channels = 0; *out_layout_mask = 0; *out_bits = 0;
    d->pkt->data = (uint8_t *)data;
    d->pkt->size = size;
    if (avcodec_send_packet(d->ctx, d->pkt) < 0) return LS_ERR_DECODE;
    for (;;) {
        int r = avcodec_receive_frame(d->ctx, d->frame);
        if (r == AVERROR(EAGAIN) || r == AVERROR_EOF) break;
        if (r < 0) return LS_ERR_DECODE;
        AVFrame *f = d->frame;
        int fmt = f->format;
        if (fmt != AV_SAMPLE_FMT_S16 && fmt != AV_SAMPLE_FMT_S16P && fmt != AV_SAMPLE_FMT_S32 &&
            fmt != AV_SAMPLE_FMT_S32P && fmt != AV_SAMPLE_FMT_FLT && fmt != AV_SAMPLE_FMT_FLTP) {
            av_frame_unref(f);
            return LS_ERR_FORMAT;
        }
        int channels = f->ch_layout.nb_channels;
        // The capacity argument counts FRAMES and the caller sizes `out` for at most 8 channels
        // per frame (LibavDecoder: capacityFrames * 8 Int32s). Refuse anything wider BEFORE any
        // write, so a stream that decodes to more channels cannot run past the buffer.
        if (channels < 1 || channels > 8) { av_frame_unref(f); return LS_ERR_FORMAT; }
        int bits = bits_for(d->ctx, fmt);
        if (*out_frames + f->nb_samples > out_capacity_frames) { av_frame_unref(f); return LS_ERR_CAPACITY; }
        for (int i = 0; i < f->nb_samples; i++)
            for (int ch = 0; ch < channels; ch++)
                out[(*out_frames + i) * channels + ch] = sample_at(f, ch, i, bits);
        *out_frames += f->nb_samples;
        *out_channels = channels;
        *out_bits = bits;
        *out_layout_mask = f->ch_layout.order == AV_CHANNEL_ORDER_NATIVE ? f->ch_layout.u.mask : 0;
        av_frame_unref(f);
    }
    return LS_OK;
}