#pragma once
/// TRACES: IR-001 | SR-003
#include "types.hpp"
#include "config.hpp"
#include "track_registry.hpp"
#ifdef SAE_SCENE_XGB
#include "inference/xgb_scene_boundary.hpp"
#include "inference/audio_logpsd.hpp"
#endif

#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": 2, "movie": "...",
//     "extraction": { "sample_fps": ..., "extinction_sec": ..., "gallery_scope": ... },
//     "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. The extraction block reports
            /// track_extinction_sec, which bounds re-association -- not the
            /// withdrawn actor keep-alive that shared its name.
            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)
        Route  route{Route::live};   ///< how it was identified (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, c.route});
        }

        // Flood-fill: snap each claim to the shot it sits in, so an actor seen
        // once in a scene is reported across the whole scene. Bounded by real
        // TransNetV2 boundaries — a window never crosses one — and a no-op when
        // scene detection found no boundaries (nothing to snap to).
        if (cfg_.presence_mode == PresenceMode::flood) {
            const std::vector<double> bounds = scene_boundaries();
            if (!bounds.empty())
                for (auto& [idx, aw] : by_actor)
                    for (auto& w : aw.scenes) {
                        w.start = boundary_at_or_before(bounds, w.start);
                        w.end   = boundary_after(bounds, w.end);
                    }
        }

        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;
    }

    // Sorted, de-duplicated boundary timestamps seen this run, framed by the
    // film's own extent so the first and last shots are closed intervals. Derived
    // from frames_ rather than a separate accumulator: the frames are already
    // retained and this runs once.
    //
    // Prefers TransNetV2 shot boundaries (is_scene_boundary) when a scene detector
    // populated them; otherwise falls back to the always-on histogram cuts
    // (is_cut, camera_position_change_detector). On this ROCm box the scene
    // detector cannot run in-process (see the dumper note), so is_cut is what
    // flood-fill actually snaps to — coarser than true shot boundaries (cuts also
    // fire on in-shot angle changes) but present with no extra pass.
    std::vector<double> scene_boundaries() const {
        std::vector<double> b;
        b.push_back(0.0);

        // Preferred: the learned XGBoost scene detector, run once here post-EOF
        // (the knee threshold needs the whole film, so this is inherently a final
        // step — like flood-fill itself). Measured best flood boundary source.
        std::vector<double> learned = xgb_boundaries();
        if (!learned.empty()) {
            for (double t : learned) b.push_back(t);
        } else {
            // Fallback: TransNetV2 shot boundaries if present, else histogram cuts.
            bool have_scene = false;
            for (const auto& sa : frames_)
                if (sa.is_scene_boundary) { have_scene = true; break; }
            for (const auto& sa : frames_) {
                const bool boundary = have_scene ? sa.is_scene_boundary : sa.is_cut;
                if (boundary) b.push_back(sa.timestamp_sec);
            }
        }
        b.push_back(last_ts_ + 1.0);   // a right edge past the final sample
        std::sort(b.begin(), b.end());
        b.erase(std::unique(b.begin(), b.end()), b.end());
        return b;
    }

    // Run the learned scene-boundary detector over the collected per-frame RGB
    // histograms + per-second audio log-PSD (decoded once from the movie). Returns
    // {} when no model is configured, the build lacks XGBoost, or no rgb_hist was
    // stamped (camera-position node only does so when a model is set).
    std::vector<double> xgb_boundaries() const {
#ifdef SAE_SCENE_XGB
        if (cfg_.scene_xgb_model.empty()) return {};
        std::vector<std::vector<float>> hist;
        std::vector<double> ts;
        hist.reserve(frames_.size()); ts.reserve(frames_.size());
        for (const auto& sa : frames_) {
            if (sa.rgb_hist.empty()) return {};   // hist not stamped → bail to fallback
            hist.push_back(sa.rgb_hist);
            ts.push_back(sa.timestamp_sec);
        }
        if (hist.size() < 16) return {};
        try {
            auto audio = AudioLogPSD::extract(cfg_.movie_path);   // [T'][B], aligned per second
            if ((int)audio.size() != (int)hist.size())
                audio.resize(hist.size(),
                             std::vector<float>(audio.empty() ? 57 : audio[0].size(), 0.f));
            XGBSceneBoundary det(cfg_.scene_xgb_model);
            auto b = det.boundaries(hist, ts, audio);
            std::cerr << "[result_sink] XGBoost scene detector: " << b.size()
                      << " boundaries\n";
            return b;
        } catch (const std::exception& e) {
            std::cerr << "[result_sink] scene detector failed (" << e.what()
                      << "), falling back to histogram cuts\n";
            return {};
        }
#else
        return {};
#endif
    }

    // The boundary opening the shot that contains t (largest boundary ≤ t).
    static double boundary_at_or_before(const std::vector<double>& b, double t) {
        auto it = std::upper_bound(b.begin(), b.end(), t);
        return (it == b.begin()) ? b.front() : *(it - 1);
    }
    // The boundary closing the shot that contains t (smallest boundary > t).
    static double boundary_after(const std::vector<double>& b, double t) {
        auto it = std::upper_bound(b.begin(), b.end(), t);
        return (it == b.end()) ? b.back() : *it;
    }

    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", route_name(w.route)}});
            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
};