scene-actor-extraction/src/nodes/result_sink_node.hpp

#pragma once
/// TRACES: IR-001 | SR-003
#include "types.hpp"
#include "config.hpp"
#include "track_registry.hpp"

#include <nlohmann/json.hpp>
#include <algorithm>
#include <atomic>
#include <cmath>
#include <fstream>
#include <iostream>
#include <map>
#include <mutex>
#include <functional>
#include <string>
#include <vector>

using json = nlohmann::json;

// ── ResultSinkFunc ────────────────────────────────────────────────────────────
// KPN sink node: accumulates SceneAnnotations and writes the final JSON on EOF.
//
// Verbosity::minimal  — merges per-frame presence into contiguous time windows.
//   Output: {
//     "schema_version": 1, "movie": "...", "sample_fps": ..., "anneal_sec": ...,
//     "actors": [{ "name", "imdb_id", "tmdb_id", "jellyfin_id", "scenes": [[t0,t1], ...] }]
//   }
//   An optional top-level "jellyfin_item_id" (the analysed title's Jellyfin item
//   GUID) may also be present: scene_analyze doesn't know it, so it's stamped in
//   by run_from_jellyfin.py after analysis. Downstream tools (cameo detection)
//   use it to check cast membership in Jellyfin's own id space — see
//   scripts/cameo_jellyfin.py.
//   This is the spec consumed by the Jellyfin plugin: each actor carries every
//   identity key the gallery knows (empty string if not resolved). The plugin
//   should prefer "jellyfin_id" (direct Person item GUID) when non-empty, and
//   otherwise resolve "imdb_id"/"tmdb_id" against the item's People ProviderIds.
//   To find who's on screen at timestamp t, scan each actor's "scenes" for a
//   window where start <= t <= end.
//
// Verbosity::standard — per-frame detail including bboxes, similarity, unknowns.
//   Output: { "frames": [{ "t", "identified": [...], "unknowns": [...] }] }
//
// eof signal: sets done_ = true so the main thread can call net.stop().

struct ResultSinkFunc {
    static constexpr std::string_view label() { return "result_sink"; }

    /// TRACES: AR-012, AR-017, IR-002 | SR-002, SR-003
    /// A finished presence claim from the registry. Called from inside the
    /// registry's reap while it holds its own lock, so this must stay a cheap
    /// push and must never re-enter the registry.
    void add_claim(const DeadTrack& d) {
        if (d.actor_idx < 0) return;              // never owned: nothing to claim
        std::lock_guard<std::mutex> g(claims_mu_);
        claims_.push_back(d);
    }

    ResultSinkFunc(const Config& cfg, std::atomic<bool>& done)
        : cfg_(cfg), done_(done)
    {}

    /// TRACES: AR-016 | SR-002
    /// Runs immediately before the output is written, with the last timestamp
    /// seen. Used to flush tracks still live at EOF, which have not timed out
    /// and would otherwise never be emitted.
    void set_pre_write_hook(std::function<void(double)> fn) { pre_write_ = std::move(fn); }

    void operator()(SceneAnnotation sa) {
        if (sa.eof) {
            flush();
            return;
        }

        // Progress to stderr
        std::cerr << "\r[result_sink] t=" << sa.timestamp_sec << "s"
                  << "  active=" << count_known(sa.visible_actors)
                  << "  unknowns=" << count_unknown(sa.visible_actors)
                  << std::flush;

        for (const auto& ia : sa.visible_actors) {
            if (ia.actor_idx < 0) continue;
            auto& m = actor_meta_[ia.actor_idx];
            if (m.name.empty())
                m = {ia.name, ia.imdb_id, ia.tmdb_id, ia.jellyfin_id};
        }
        last_ts_ = sa.timestamp_sec;
        frames_.push_back(std::move(sa));
    }

    // Write accumulated results and signal done. Safe to call more than once.
    void flush() {
        if (written_.exchange(true)) return;
        if (pre_write_) pre_write_(last_ts_);
        write_output();
        done_.store(true, std::memory_order_release);
    }

private:
    // Bump when the minimal/standard output JSON structure changes in a way
    // the Jellyfin plugin needs to detect.
    static constexpr int kSchemaVersion = 2;   // SR-003 coordinated bump

    static int count_known(const std::vector<IdentifiedActor>& v) {
        int n = 0;
        for (const auto& a : v) if (a.actor_idx >= 0) ++n;
        return n;
    }
    static int count_unknown(const std::vector<IdentifiedActor>& v) {
        int n = 0;
        for (const auto& a : v) if (a.actor_idx < 0) ++n;
        return n;
    }

    void write_output() {
        std::cerr << "\n[result_sink] writing " << cfg_.output_path << "\n";

        json root;
        if (cfg_.verbosity == Verbosity::xray) {
            root = build_xray();
        } else {
            /// TRACES: IR-002 | SR-003
            /// schema_version 2, per jRay/SPEC.md JR-002. anneal_sec is REMOVED
            /// rather than zeroed: a field naming a mechanism the pipeline no
            /// longer has is actively misleading, and would outlive everyone who
            /// remembers why it reads 0. extinction_sec succeeds it as the
            /// parameter that actually shapes window extent.
            root["schema_version"] = kSchemaVersion;
            root["movie"]          = cfg_.movie_path;
            root["extraction"] = {
                {"sample_fps",     cfg_.sample_fps},
                {"extinction_sec", cfg_.track_extinction_sec},
                {"gallery_scope",  cfg_.gallery_scope},
            };
            root["actors"]     = build_epochs();
            if (cfg_.verbosity == Verbosity::standard)
                root["frames"] = build_standard();
        }

        std::ofstream f(cfg_.output_path);
        if (!f.is_open()) {
            std::cerr << "[result_sink] ERROR: cannot write " << cfg_.output_path << "\n";
            return;
        }
        f << root.dump(2) << "\n";
        std::cerr << "[result_sink] done.\n";
    }

    struct Window {
        double start{0.0};
        double end{0.0};
        float  belief{0.f};   ///< the posterior that justified the claim (AR-017)
    };
    struct ActorWindow {
        std::string name, imdb_id, tmdb_id, jellyfin_id;
        std::vector<Window> scenes;
    };
    struct ActorMeta { std::string name, imdb_id, tmdb_id, jellyfin_id; };

    // Core logic: merge per-frame detections into annealed [start, end] windows.
    /// TRACES: AR-012, IR-002 | SR-002
    /// A claim already IS a window — `[first_seen, last_seen]` of a track the
    /// actor owned. There is no annealing pass: `anneal_sec` existed to bridge
    /// gaps between isolated accepted frames, and a track that survives its own
    /// gaps leaves it nothing to do (see the AR-012 withdrawal note).
    std::vector<ActorWindow> build_actor_windows() {
        std::lock_guard<std::mutex> g(claims_mu_);

        std::map<int, ActorWindow> by_actor;
        for (const auto& c : claims_) {
            auto& aw = by_actor[c.actor_idx];
            if (aw.name.empty()) {
                auto it = actor_meta_.find(c.actor_idx);
                if (it != actor_meta_.end()) {
                    aw.name        = it->second.name;
                    aw.imdb_id     = it->second.imdb_id;
                    aw.tmdb_id     = it->second.tmdb_id;
                    aw.jellyfin_id = it->second.jellyfin_id;
                }
            }
            aw.scenes.push_back({c.first_seen, c.last_seen, c.belief});
        }

        std::vector<ActorWindow> result;
        for (auto& [idx, aw] : by_actor) {
            std::sort(aw.scenes.begin(), aw.scenes.end(),
                      [](const Window& a, const Window& b) { return a.start < b.start; });
            result.push_back(std::move(aw));
        }
        return result;
    }

    json build_epochs() {
        json actors = json::array();
        for (const auto& aw : build_actor_windows()) {
            // Objects, not float pairs: a window carries the belief that
            // justified it and the route by which it was identified (AR-017),
            // so a consumer can caveat or filter rather than treating every
            // window as equally certain.
            json windows = json::array();
            for (const auto& w : aw.scenes)
                windows.push_back({{"start", w.start},
                                   {"end",   w.end},
                                   {"belief", w.belief},
                                   {"route", "live"}});
            json ja;
            ja["name"]        = aw.name;
            ja["imdb_id"]     = aw.imdb_id;
            ja["tmdb_id"]     = aw.tmdb_id;
            ja["jellyfin_id"] = aw.jellyfin_id;
            ja["scenes"]      = std::move(windows);
            actors.push_back(std::move(ja));
        }
        return actors;
    }

    // Jellyfin-Xray format: { "second": ["Actor", ...] }
    // Expands each annealed window into every integer second so coverage is dense
    // regardless of sample rate. Seconds between scenes have no key → overlay clears.
    json build_xray() {
        std::map<int, std::vector<std::string>> xray;
        for (const auto& aw : build_actor_windows()) {
            for (const auto& w : aw.scenes) {
                int t0 = static_cast<int>(std::floor(w.start));
                int t1 = static_cast<int>(std::ceil(w.end));
                for (int t = t0; t <= t1; ++t)
                    xray[t].push_back(aw.name);
            }
        }

        json root = json::object();
        for (const auto& [t, names] : xray)
            root[std::to_string(t)] = names;
        return root;
    }

    json build_standard() {
        json frames = json::array();
        for (const auto& frame : frames_) {
            json jf;
            jf["t"]           = frame.timestamp_sec;
            jf["identified"]  = json::array();
            jf["unknowns"]    = json::array();

            for (const auto& ia : frame.visible_actors) {
                const auto& b = ia.bbox;
                json jbox = {b.x, b.y, b.width, b.height};

                if (ia.actor_idx >= 0) {
                    json ja;
                    ja["name"]        = ia.name;
                    ja["imdb_id"]     = ia.imdb_id;
                    ja["tmdb_id"]     = ia.tmdb_id;
                    ja["jellyfin_id"] = ia.jellyfin_id;
                    ja["similarity"] = ia.similarity;
                    ja["track_id"]   = ia.track_id;
                    ja["bbox"]       = jbox;
                    jf["identified"].push_back(std::move(ja));
                } else {
                    json ju;
                    ju["bbox"]       = jbox;
                    ju["track_id"]   = ia.track_id;
                    ju["confidence"] = ia.similarity;  // reuse field; 0 for unknowns
                    jf["unknowns"].push_back(std::move(ju));
                }
            }
            frames.push_back(std::move(jf));
        }
        return frames;
    }

    const Config&                cfg_;
    std::atomic<bool>&           done_;
    std::atomic<bool>            written_{false};
    std::vector<SceneAnnotation> frames_;
    std::function<void(double)>  pre_write_;
    double                       last_ts_{0.0};
    std::mutex                   claims_mu_;
    std::vector<DeadTrack>       claims_;
    std::map<int, ActorMeta>     actor_meta_;   ///< actor_idx → identity keys
};