scene-actor-extraction/src/audio_signature.cpp

// ── JRay audio signature, v1 — implementation ────────────────────────────────
//
/// TRACES: IR-004, IR-007, IR-008 | SR-003
//
// The contract this implements is documented in full in audio_signature.hpp;
// read that before changing anything here. Every constant is load-bearing: the
// JRay Jellyfin plugin computes the same bytes in C#, and a signature that
// differs in any parameter simply does not match.
//
// Audio decode is a *second stream from an existing dependency* — the pipeline
// already links libavformat/libavcodec/libavutil for video (ffmpeg_decoder.hpp);
// this adds libswresample for the downmix+resample, no new project dependency.
// The FFT is written out here rather than pulled from a library for the same
// reason the plugin vendors one: it is a fixed, fully specified transform, and
// a dependency whose version could change the numerics is a liability when the
// output has to be bit-identical across two languages.

#include "audio_signature.hpp"

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/channel_layout.h>
#include <libavutil/opt.h>
#include <libavutil/samplefmt.h>
#include <libswresample/swresample.h>
}

#include <algorithm>
#include <cmath>
#include <cstring>
#include <vector>

namespace sae::audio {
namespace {

constexpr double kPi = 3.14159265358979323846;

// ── Band table ───────────────────────────────────────────────────────────────
// edge[b] = 300 * 10^(b/32); band b owns FFT bins [k_lo[b], k_lo[b+1]).
// ceil() of the edge in bins, so membership is decided once by integers rather
// than by a float comparison per bin per frame. The bands tile [112, 1115)
// contiguously with no gap and no overlap, which is what lets the frame energy
// below be accumulated from the per-band sums.
std::array<std::pair<int, int>, kNumBands> build_band_table() {
    const double hz_per_bin = static_cast<double>(kSampleRate) / kFrameSize;
    std::array<int, kNumBands + 1> k{};
    for (int b = 0; b <= kNumBands; ++b) {
        const double edge = kBandLoHz * std::pow(kBandHiHz / kBandLoHz,
                                                 static_cast<double>(b) / kNumBands);
        k[b] = static_cast<int>(std::ceil(edge / hz_per_bin));
    }
    std::array<std::pair<int, int>, kNumBands> tbl{};
    for (int b = 0; b < kNumBands; ++b) tbl[b] = {k[b], k[b + 1]};
    return tbl;
}

// Hann, periodic: w[n] = 0.5 * (1 - cos(2*pi*n/N)). Not the symmetric (N-1)
// variant — the two differ, and the difference is observable.
const std::vector<double>& hann_window() {
    static const std::vector<double> w = [] {
        std::vector<double> v(kFrameSize);
        for (int n = 0; n < kFrameSize; ++n)
            v[n] = 0.5 * (1.0 - std::cos(2.0 * kPi * n / kFrameSize));
        return v;
    }();
    return w;
}

// ── Radix-2 decimation-in-time complex FFT, in place, no normalisation ──────
// Twiddles are precomputed per stage from cos/sin of -2*pi*j/len so the angle
// is an exactly reproducible double in any language and only the libm rounding
// of cos/sin (≤1 ulp) can differ — orders of magnitude below the decision
// margins in the golden fixture.
struct FftTables {
    std::vector<int> rev;                       // bit-reversal permutation
    std::vector<std::vector<double>> wr, wi;    // per stage
};

const FftTables& fft_tables() {
    static const FftTables t = [] {
        FftTables f;
        f.rev.resize(kFrameSize);
        int bits = 0;
        while ((1 << bits) < kFrameSize) ++bits;
        for (int i = 0; i < kFrameSize; ++i) {
            int r = 0;
            for (int b = 0; b < bits; ++b)
                if (i & (1 << b)) r |= 1 << (bits - 1 - b);
            f.rev[i] = r;
        }
        for (int len = 2; len <= kFrameSize; len <<= 1) {
            const int half = len / 2;
            std::vector<double> cr(half), ci(half);
            for (int j = 0; j < half; ++j) {
                const double ang = -2.0 * kPi * j / len;
                cr[j] = std::cos(ang);
                ci[j] = std::sin(ang);
            }
            f.wr.push_back(std::move(cr));
            f.wi.push_back(std::move(ci));
        }
        return f;
    }();
    return t;
}

void fft_4096(std::vector<double>& re, std::vector<double>& im) {
    const FftTables& t = fft_tables();
    for (int i = 0; i < kFrameSize; ++i) {
        const int j = t.rev[i];
        if (i < j) { std::swap(re[i], re[j]); std::swap(im[i], im[j]); }
    }
    int stage = 0;
    for (int len = 2; len <= kFrameSize; len <<= 1, ++stage) {
        const int half = len / 2;
        const std::vector<double>& wr = t.wr[stage];
        const std::vector<double>& wi = t.wi[stage];
        for (int base = 0; base < kFrameSize; base += len) {
            for (int j = 0; j < half; ++j) {
                const int a = base + j;
                const int b = a + half;
                const double tr = re[b] * wr[j] - im[b] * wi[j];
                const double ti = re[b] * wi[j] + im[b] * wr[j];
                re[b] = re[a] - tr;  im[b] = im[a] - ti;
                re[a] = re[a] + tr;  im[a] = im[a] + ti;
            }
        }
    }
}

int energy_class(double r) {
    if (r < kEnergyClassEdges[0]) return 0;
    if (r < kEnergyClassEdges[1]) return 1;
    if (r < kEnergyClassEdges[2]) return 2;
    return 3;
}

// ── FFmpeg RAII ─────────────────────────────────────────────────────────────
struct DecodeCtx {
    AVFormatContext* fmt  = nullptr;
    AVCodecContext*  dec  = nullptr;
    SwrContext*      swr  = nullptr;
    AVFrame*         frm  = nullptr;
    AVPacket*        pkt  = nullptr;
    ~DecodeCtx() {
        if (swr) swr_free(&swr);
        if (frm) av_frame_free(&frm);
        if (pkt) av_packet_free(&pkt);
        if (dec) avcodec_free_context(&dec);
        if (fmt) avformat_close_input(&fmt);
    }
};

bool open_resampler(DecodeCtx& c, const AVFrame* f) {
#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(57, 24, 100)
    AVChannelLayout out_layout;
    av_channel_layout_default(&out_layout, 1);          // mono
    AVChannelLayout in_layout;
    if (av_channel_layout_copy(&in_layout, &f->ch_layout) < 0) return false;
    if (in_layout.nb_channels <= 0) {
        av_channel_layout_uninit(&in_layout);
        av_channel_layout_default(&in_layout, 1);
    }
    const int rc = swr_alloc_set_opts2(
        &c.swr,
        &out_layout, AV_SAMPLE_FMT_FLT, kSampleRate,
        &in_layout, static_cast<AVSampleFormat>(f->format),
        f->sample_rate ? f->sample_rate : kSampleRate,
        0, nullptr);
    av_channel_layout_uninit(&in_layout);
    av_channel_layout_uninit(&out_layout);
    if (rc < 0 || !c.swr) return false;
#else
    const int64_t in_layout = f->channel_layout
        ? static_cast<int64_t>(f->channel_layout)
        : av_get_default_channel_layout(f->channels ? f->channels : 1);
    c.swr = swr_alloc_set_opts(
        nullptr,
        AV_CH_LAYOUT_MONO, AV_SAMPLE_FMT_FLT, kSampleRate,
        in_layout, static_cast<AVSampleFormat>(f->format),
        f->sample_rate ? f->sample_rate : kSampleRate,
        0, nullptr);
    if (!c.swr) return false;
#endif
    return swr_init(c.swr) >= 0;
}

// Push one decoded frame (or a flush) through the resampler, dropping the
// leading `to_skip` output samples, and append to `out`.
void drain(SwrContext* swr, const AVFrame* f, int in_rate,
           std::size_t& to_skip, std::vector<float>& out) {
    const int64_t delay   = swr_get_delay(swr, in_rate ? in_rate : kSampleRate);
    const int     in_n    = f ? f->nb_samples : 0;
    const int     max_out = static_cast<int>(av_rescale_rnd(
        delay + in_n, kSampleRate, in_rate ? in_rate : kSampleRate, AV_ROUND_UP)) + 32;
    if (max_out <= 0) return;

    std::vector<float> buf(static_cast<std::size_t>(max_out));
    uint8_t* dst = reinterpret_cast<uint8_t*>(buf.data());
    const int n = swr_convert(swr, &dst, max_out,
                              f ? const_cast<const uint8_t**>(f->extended_data) : nullptr,
                              in_n);
    if (n <= 0) return;

    std::size_t produced = static_cast<std::size_t>(n);
    std::size_t off = 0;
    if (to_skip) {
        const std::size_t drop = std::min(to_skip, produced);
        to_skip -= drop;
        off = drop;
        produced -= drop;
    }
    if (produced)
        out.insert(out.end(), buf.begin() + off, buf.begin() + off + produced);
}

}  // namespace

// ── Public surface ──────────────────────────────────────────────────────────

const std::array<std::pair<int, int>, kNumBands>& band_fft_bins() {
    static const std::array<std::pair<int, int>, kNumBands> tbl = build_band_table();
    return tbl;
}

std::string base64_encode(const std::uint8_t* data, std::size_t n) {
    static constexpr char kAlphabet[] =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::string out;
    out.reserve(((n + 2) / 3) * 4);
    std::size_t i = 0;
    for (; i + 3 <= n; i += 3) {
        const std::uint32_t v = (std::uint32_t(data[i]) << 16) |
                                (std::uint32_t(data[i + 1]) << 8) |
                                 std::uint32_t(data[i + 2]);
        out += kAlphabet[(v >> 18) & 0x3F];
        out += kAlphabet[(v >> 12) & 0x3F];
        out += kAlphabet[(v >> 6) & 0x3F];
        out += kAlphabet[v & 0x3F];
    }
    if (i < n) {
        const bool two = (n - i) == 2;
        const std::uint32_t v = (std::uint32_t(data[i]) << 16) |
                                (two ? (std::uint32_t(data[i + 1]) << 8) : 0u);
        out += kAlphabet[(v >> 18) & 0x3F];
        out += kAlphabet[(v >> 12) & 0x3F];
        out += two ? kAlphabet[(v >> 6) & 0x3F] : '=';
        out += '=';
    }
    return out;
}

std::uint64_t fnv1a64(const void* data, std::size_t n) {
    const auto* p = static_cast<const std::uint8_t*>(data);
    std::uint64_t h = 0xcbf29ce484222325ULL;
    for (std::size_t i = 0; i < n; ++i) {
        h ^= p[i];
        h *= 0x100000001b3ULL;
    }
    return h;
}

/// TRACES: IR-004
std::vector<std::uint8_t> pack_frames(const std::vector<float>& mono) {
    if (mono.size() < static_cast<std::size_t>(kFrameSize)) return {};

    const std::size_t nframes = 1 + (mono.size() - kFrameSize) / kHopSize;
    const auto& bands = band_fft_bins();
    const auto& win   = hann_window();
    const int   k_lo  = bands.front().first;
    const int   k_hi  = bands.back().second;      // exclusive
    const double bin_count = static_cast<double>(k_hi - k_lo);

    std::vector<double>       re(kFrameSize), im(kFrameSize);
    std::vector<std::uint8_t> peak(nframes);
    std::vector<double>       energy(nframes);

    for (std::size_t f = 0; f < nframes; ++f) {
        const float* src = mono.data() + f * kHopSize;
        for (int n = 0; n < kFrameSize; ++n) {
            re[n] = static_cast<double>(src[n]) * win[n];
            im[n] = 0.0;
        }
        fft_4096(re, im);

        // Per-band mean magnitude; the bands tile the 300–3000 Hz range with no
        // gaps, so the frame's band-limited energy is the sum of the band sums.
        double best = -1.0, total = 0.0;
        int    best_b = 0;
        for (int b = 0; b < kNumBands; ++b) {
            double sum = 0.0;
            for (int k = bands[b].first; k < bands[b].second; ++k)
                sum += std::sqrt(re[k] * re[k] + im[k] * im[k]);
            total += sum;
            const double mean = sum / (bands[b].second - bands[b].first);
            if (mean > best) { best = mean; best_b = b; }   // ties → lowest index
        }
        peak[f]   = static_cast<std::uint8_t>(best_b);
        energy[f] = total / bin_count;
    }

    // Reference is the upper median of the frame energies: an actually observed
    // value (no averaging of the two middle samples), so it is bit-reproducible,
    // gain-invariant and barely moves when the window is trimmed.
    std::vector<double> sorted = energy;
    std::sort(sorted.begin(), sorted.end());
    const double ref = sorted[sorted.size() / 2];

    std::vector<std::uint8_t> out(nframes);
    for (std::size_t f = 0; f < nframes; ++f) {
        const double r = std::log10((energy[f] + kEnergyEps) / (ref + kEnergyEps));
        out[f] = static_cast<std::uint8_t>(((peak[f] & 0x1F) << 2) |
                                           (energy_class(r) & 0x03));
    }
    return out;
}

/// TRACES: IR-004, IR-008
std::optional<std::string> signature_from_mono(const std::vector<float>& mono) {
    const std::vector<std::uint8_t> packed = pack_frames(mono);
    if (packed.empty()) return std::nullopt;
    return std::string(kVersionPrefix) + base64_encode(packed.data(), packed.size());
}

/// TRACES: IR-004, IR-007
std::optional<std::vector<float>> decode_centre_window(const std::string& path) {
    av_log_set_level(AV_LOG_ERROR);

    DecodeCtx c;
    if (avformat_open_input(&c.fmt, path.c_str(), nullptr, nullptr) < 0)
        return std::nullopt;
    if (avformat_find_stream_info(c.fmt, nullptr) < 0) return std::nullopt;
    if (c.fmt->duration == AV_NOPTS_VALUE) return std::nullopt;

    const double duration = static_cast<double>(c.fmt->duration) / AV_TIME_BASE;

    // IR-007 — the window underflows, so there is no signature and no sync
    // offset downstream. The plugin applies the identical rule.
    if (duration < kWindowSec) return std::nullopt;

    const int idx = av_find_best_stream(c.fmt, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
    if (idx < 0) return std::nullopt;                     // no audio → no signature

    AVStream* st = c.fmt->streams[idx];
    const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id);
    if (!codec) return std::nullopt;
    c.dec = avcodec_alloc_context3(codec);
    if (!c.dec) return std::nullopt;
    if (avcodec_parameters_to_context(c.dec, st->codecpar) < 0) return std::nullopt;
    c.dec->thread_count = 0;
    if (avcodec_open2(c.dec, codec, nullptr) < 0) return std::nullopt;

    const double start_sec = duration / 2.0 - kWindowSec / 2.0;

    // Seek to a packet at or before the window start; the exact start is then
    // reached by discarding the leading output samples, which is what
    // `ffmpeg -ss <t> -i <file>` does and therefore what the plugin sees.
    if (start_sec > 0.0) {
        const int64_t tgt = av_rescale_q(
            static_cast<int64_t>(start_sec * AV_TIME_BASE), AV_TIME_BASE_Q, st->time_base);
        if (av_seek_frame(c.fmt, idx, tgt, AVSEEK_FLAG_BACKWARD) >= 0)
            avcodec_flush_buffers(c.dec);
    }

    c.frm = av_frame_alloc();
    c.pkt = av_packet_alloc();
    if (!c.frm || !c.pkt) return std::nullopt;

    std::vector<float> mono;
    mono.reserve(kWindowSamples + kSampleRate);
    std::size_t to_skip  = 0;
    bool        have_swr = false;
    int         in_rate  = kSampleRate;
    bool        eof      = false;

    while (mono.size() < kWindowSamples && !eof) {
        const int rr = av_read_frame(c.fmt, c.pkt);
        if (rr < 0) {
            eof = true;
            avcodec_send_packet(c.dec, nullptr);          // flush the decoder
        } else if (c.pkt->stream_index != idx) {
            av_packet_unref(c.pkt);
            continue;
        } else {
            avcodec_send_packet(c.dec, c.pkt);
            av_packet_unref(c.pkt);
        }

        while (avcodec_receive_frame(c.dec, c.frm) == 0) {
            if (!have_swr) {
                if (!open_resampler(c, c.frm)) return std::nullopt;
                have_swr = true;
                in_rate  = c.frm->sample_rate ? c.frm->sample_rate : kSampleRate;

                int64_t pts = c.frm->best_effort_timestamp;
                if (pts == AV_NOPTS_VALUE) pts = c.frm->pts;
                const double t0 = (pts == AV_NOPTS_VALUE)
                    ? start_sec : av_q2d(st->time_base) * static_cast<double>(pts);
                const double lead = start_sec - t0;
                to_skip = lead > 0.0
                    ? static_cast<std::size_t>(std::llround(lead * kSampleRate)) : 0;
            }
            drain(c.swr, c.frm, in_rate, to_skip, mono);
            av_frame_unref(c.frm);
            if (mono.size() >= kWindowSamples) break;
        }
    }

    if (have_swr && mono.size() < kWindowSamples)
        drain(c.swr, nullptr, in_rate, to_skip, mono);     // flush the resampler

    if (mono.empty()) return std::nullopt;
    // Truncate to exactly 120.000 s so the frame count is 1288 for every input
    // and does not wobble with seek granularity or the resampler tail.
    if (mono.size() > kWindowSamples) mono.resize(kWindowSamples);
    return mono;
}

/// TRACES: IR-004, IR-005, IR-007, IR-008
std::optional<std::string> compute_signature(const std::string& path) {
    const std::optional<std::vector<float>> mono = decode_centre_window(path);
    if (!mono) return std::nullopt;
    return signature_from_mono(*mono);
}

}  // namespace sae::audio