/*
 * MediaBridge dovi_convert bitstream filter (plan spec 7).
 *
 * Dolby Vision Profile 7 (BL+EL+RPU) → Profile 8.1 (BL+RPU, HDR10-compatible)
 * without re-encoding: per access unit, NAL 62 (RPU) is converted with libdovi
 * (dovi_tool's dolby_vision crate, MIT) and NAL 63 (enhancement layer) is
 * dropped. AVPlayer plays 8.1; it can never play 7.
 *
 * Copied into libavcodec/ by build-ffmpeg.sh (kept out of the FFmpeg tarball
 * so the patch stays tiny); registered in bitstream_filters.c by the same
 * script. Compiled only when CONFIG_DOVI_CONVERT_BSF is enabled.
 *
 * Failure ladder (mandate: never a black screen):
 *  - first RPU fails to parse/convert → latch session-wide strip immediately
 *    (equivalent to the plan's probe test-convert: DV is never advertised on
 *    a stream we can't deliver — BridgeSession greps the log marker before
 *    the master playlist is served).
 *  - ≥5 consecutive failures later → latch strip.
 *  - stripped output = plain HDR10 base layer (P7 BL is bl_compat_id 1).
 * Marker string "dovi-strip-latched" is grepped by BridgeSession — keep it.
 *
 * mode option: 81 (default) = convert RPU to 8.1;  0 = strip 62/63 entirely
 * (used for P7 dual-track sources where the RPU references a separate EL
 * track we deliberately drop).
 */

#include "libavutil/intreadwrite.h"
#include "libavutil/mem.h"
#include "libavutil/opt.h"
#include "libavutil/dovi_meta.h"
#include "bsf.h"
#include "bsf_internal.h"

#include <libdovi/rpu_parser.h>

#define DOVI_MODE_STRIP 0
#define DOVI_MODE_81    81

typedef struct DoviConvertContext {
    const AVClass *class;
    int mode;          // DOVI_MODE_*
    int length_size;   // 0 = Annex B start codes, else hvcC NAL length prefix
    int seen_rpu;      // a NAL 62 was processed successfully at least once
    int fail_streak;   // consecutive RPU failures
    int latched_strip; // permanent downgrade to strip for this session
} DoviConvertContext;

// Growing output buffer: fatal-free append with av_fast_realloc bookkeeping.
typedef struct GrowBuf {
    uint8_t *data;
    unsigned cap;
    size_t   len;
} GrowBuf;

static int grow_append(GrowBuf *b, const uint8_t *src, size_t n)
{
    uint8_t *p = av_fast_realloc(b->data, &b->cap, b->len + n);
    if (!p)
        return AVERROR(ENOMEM);
    b->data = p;
    memcpy(b->data + b->len, src, n);
    b->len += n;
    return 0;
}

static void latch_strip(AVBSFContext *ctx, const char *why)
{
    DoviConvertContext *s = ctx->priv_data;
    if (s->latched_strip)
        return;
    s->latched_strip = 1;
    // BridgeSession matches this marker (see header comment) — do not reword.
    av_log(ctx, AV_LOG_WARNING, "dovi-strip-latched: %s\n", why);
}

/* Convert one NAL 62 (buffer includes the 2-byte 0x7C01 NAL header, payload
 * still emulation-escaped — libdovi handles both). Appends the replacement
 * NAL payload (0x7C01 + escaped RPU) to `out` on success.
 * Returns 0 on success, <0 on conversion failure (caller drops the NAL). */
static int convert_rpu(AVBSFContext *ctx, const uint8_t *nal, size_t len, GrowBuf *out)
{
    DoviConvertContext *s = ctx->priv_data;
    DoviRpuOpaque *rpu = dovi_parse_unspec62_nalu(nal, len);
    const DoviData *data = NULL;
    int ret = AVERROR_INVALIDDATA;

    if (!rpu)
        return AVERROR(ENOMEM);
    if (dovi_rpu_get_error(rpu))
        goto done;
    if (dovi_convert_rpu_with_mode(rpu, 2) != 0)   // 2 = to profile 8.1
        goto done;
    data = dovi_write_unspec62_nalu(rpu);
    if (!data || dovi_rpu_get_error(rpu))
        goto done;
    ret = grow_append(out, data->data, data->len);

done:
    if (ret == AVERROR_INVALIDDATA) {
        const char *err = dovi_rpu_get_error(rpu);
        av_log(ctx, AV_LOG_WARNING, "dovi_convert: RPU rejected (%s)\n",
               err ? err : "unknown");
    }
    if (data)
        dovi_data_free(data);
    dovi_rpu_free(rpu);
    return ret;
}

/* Handle one NAL: append it (possibly converted) to `out`, or drop it.
 * `nal`/`len` exclude any start code / length prefix. Emits Annex B or
 * length-prefixed framing to match the input. */
static int handle_nal(AVBSFContext *ctx, const uint8_t *nal, size_t len, GrowBuf *out)
{
    DoviConvertContext *s = ctx->priv_data;
    static const uint8_t startcode[4] = { 0, 0, 0, 1 };
    int type = len ? (nal[0] >> 1) & 0x3F : -1;
    int ret;

    if (type == 63)                    // enhancement layer: always dropped
        return 0;

    if (type == 62) {
        if (s->mode == DOVI_MODE_STRIP || s->latched_strip)
            return 0;
        GrowBuf conv = { 0 };
        ret = convert_rpu(ctx, nal, len, &conv);
        if (ret == AVERROR(ENOMEM)) {
            av_freep(&conv.data);
            return ret;
        }
        if (ret < 0) {                 // conversion failed → drop this RPU
            av_freep(&conv.data);
            if (!s->seen_rpu)          // first RPU already broken → give up now
                latch_strip(ctx, "first RPU rejected");
            else if (++s->fail_streak >= 5)
                latch_strip(ctx, "5 consecutive RPU failures");
            return 0;
        }
        s->seen_rpu = 1;
        s->fail_streak = 0;
        if (s->length_size) {
            uint8_t pre[4];
            for (int i = 0; i < s->length_size; i++)
                pre[i] = (uint8_t)(conv.len >> (8 * (s->length_size - 1 - i)));
            ret = grow_append(out, pre, s->length_size);
        } else {
            ret = grow_append(out, startcode, 4);
        }
        if (ret >= 0)
            ret = grow_append(out, conv.data, conv.len);
        av_freep(&conv.data);
        return ret;
    }

    // Every other NAL: copy through with matching framing.
    if (s->length_size) {
        uint8_t pre[4];
        for (int i = 0; i < s->length_size; i++)
            pre[i] = (uint8_t)(len >> (8 * (s->length_size - 1 - i)));
        ret = grow_append(out, pre, s->length_size);
    } else {
        ret = grow_append(out, startcode, 4);
    }
    if (ret >= 0)
        ret = grow_append(out, nal, len);
    return ret;
}

static int dovi_convert_filter(AVBSFContext *ctx, AVPacket *pkt)
{
    DoviConvertContext *s = ctx->priv_data;
    GrowBuf out = { 0 };
    int ret, touched = 0;

    ret = ff_bsf_get_packet_ref(ctx, pkt);
    if (ret < 0)
        return ret;

    // Fast scan: does this AU contain NAL 62/63 at all? Most don't grow, and
    // packets without DV NALs pass through untouched (zero-copy).
    const uint8_t *p = pkt->data, *end = pkt->data + pkt->size;
    if (s->length_size) {
        while (p + s->length_size <= end) {
            size_t n = 0;
            for (int i = 0; i < s->length_size; i++)
                n = (n << 8) | *p++;
            if (n == 0 || p + n > end)
                break;
            int type = (p[0] >> 1) & 0x3F;
            if (type == 62 || type == 63) { touched = 1; break; }
            p += n;
        }
    } else {
        for (const uint8_t *q = pkt->data; q + 4 <= end; q++) {
            if (q[0] == 0 && q[1] == 0 && q[2] == 1) {
                int type = (q[3] >> 1) & 0x3F;
                if (type == 62 || type == 63) { touched = 1; break; }
                q += 2;
            }
        }
    }
    if (!touched)
        return 0;

    // Rebuild the AU without EL, with the RPU converted (or stripped).
    if (s->length_size) {
        p = pkt->data;
        while (p + s->length_size <= end) {
            size_t n = 0;
            for (int i = 0; i < s->length_size; i++)
                n = (n << 8) | *p++;
            if (n == 0 || p + n > end)
                break;
            ret = handle_nal(ctx, p, n, &out);
            if (ret < 0)
                goto fail;
            p += n;
        }
    } else {
        // Annex B: NAL spans from after a start code to the next start code.
        const uint8_t *nal = NULL;
        for (p = pkt->data; p + 3 <= end; p++) {
            if (p[0] == 0 && p[1] == 0 && p[2] == 1) {
                if (nal) {
                    const uint8_t *e = p;
                    while (e > nal && e[-1] == 0)  // trailing_zero_8bits belong to the boundary
                        e--;
                    ret = handle_nal(ctx, nal, e - nal, &out);
                    if (ret < 0)
                        goto fail;
                }
                nal = p + 3;
                p += 2;
            }
        }
        if (nal) {
            ret = handle_nal(ctx, nal, end - nal, &out);
            if (ret < 0)
                goto fail;
        }
    }

    ret = av_buffer_replace(&pkt->buf, NULL);
    if (ret < 0)
        goto fail;
    pkt->buf = av_buffer_create(out.data, out.len, av_buffer_default_free, NULL, 0);
    if (!pkt->buf) {
        ret = AVERROR(ENOMEM);
        goto fail;
    }
    pkt->data = out.data;
    pkt->size = out.len;
    return 0;

fail:
    av_freep(&out.data);
    av_packet_unref(pkt);
    return ret;
}

static int dovi_convert_init(AVBSFContext *ctx)
{
    DoviConvertContext *s = ctx->priv_data;

    // hvcC extradata (configurationVersion == 1) → length-prefixed NALs;
    // anything else (Annex B TS) → start codes.
    if (ctx->par_in->extradata_size >= 23 && ctx->par_in->extradata[0] == 1)
        s->length_size = (ctx->par_in->extradata[21] & 3) + 1;
    else
        s->length_size = 0;

    // Rewrite the DV configuration record so movenc authors a correct 8.1
    // dvvC (profile 8, no EL, HDR10-compatible base layer). Strip mode
    // removes it entirely — the output is plain HDR10.
    AVPacketSideData *sd = (AVPacketSideData *)av_packet_side_data_get(
        ctx->par_out->coded_side_data, ctx->par_out->nb_coded_side_data,
        AV_PKT_DATA_DOVI_CONF);
    if (sd && s->mode == DOVI_MODE_81) {
        AVDOVIDecoderConfigurationRecord *cfg =
            (AVDOVIDecoderConfigurationRecord *)sd->data;
        if (cfg->dv_profile == 7) {
            cfg->dv_profile = 8;
            cfg->el_present_flag = 0;
            cfg->dv_bl_signal_compatibility_id = 1;
        }
    } else if (sd) {
        av_packet_side_data_remove(ctx->par_out->coded_side_data,
                                   &ctx->par_out->nb_coded_side_data,
                                   AV_PKT_DATA_DOVI_CONF);
    }
    return 0;
}

#define OFFSET(x) offsetof(DoviConvertContext, x)
static const AVOption dovi_convert_options[] = {
    { "mode", "81 = convert RPU to profile 8.1, 0 = strip DV NALs",
      OFFSET(mode), AV_OPT_TYPE_INT, { .i64 = DOVI_MODE_81 }, 0, 81, 0 },
    { NULL }
};

static const AVClass dovi_convert_class = {
    .class_name = "dovi_convert_bsf",
    .item_name  = av_default_item_name,
    .option     = dovi_convert_options,
    .version    = LIBAVUTIL_VERSION_INT,
};

static const enum AVCodecID dovi_convert_codec_ids[] = {
    AV_CODEC_ID_HEVC, AV_CODEC_ID_NONE,
};

const FFBitStreamFilter ff_dovi_convert_bsf = {
    .p.name         = "dovi_convert",
    .p.codec_ids    = dovi_convert_codec_ids,
    .p.priv_class   = &dovi_convert_class,
    .priv_data_size = sizeof(DoviConvertContext),
    .init           = dovi_convert_init,
    .filter         = dovi_convert_filter,
};