#pragma once
#include "types.hpp"
#include "config.hpp"

#include <functional>

#include <cmath>
#include <cstdio>
#include <iostream>
#include <limits>
#include <map>
#include <string>
#include <vector>

#ifdef SAE_DEBUG
#include <opencv2/imgcodecs.hpp>
#endif
#include <opencv2/core.hpp>
#include <filesystem>

// ── TrackGallery ──────────────────────────────────────────────────────────────
// Per-film gallery expansion driven by track continuity.
//
// A single uncut face track is, by construction, one physical person: the
// face_tracker links detections frame-to-frame and clears all tracks on a scene
// cut, so a track ID never spans a cut. That continuity is a same-identity label
// the baked gallery does not have. TrackGallery exploits it in two stages:
//
//   1. Diversity buffer (per track). Every frame's embedding is offered to a
//      fixed-capacity buffer for that track. When full, the member whose best
//      similarity to the *owning actor's* gallery references is HIGHEST is
//      dropped — i.e. the pose the gallery already recognises well is the least
//      informative, so the buffer is continuously biased toward the gallery-far
//      (novel-pose) embeddings on which recognition currently fails.
//
//   2. Promotion (on confirmation). A track is "owned" by actor A once ≥N frames
//      have been accepted (by the matcher's calibrated posterior) as A. On
//      confirmation the retained buffer — the hard, gallery-far poses — is
//      promoted into A's per-film annex, subject to one safety gate: the band's
//      lower bound, re-applied across the whole store (see `store_coherence`).
//
// There is exactly one threshold here, the AR-018 band, and it is a calibrated
// probability. Novelty is no longer a threshold at all — the eviction policy
// above *orders* by gallery similarity rather than cutting at a constant, and
// the band's upper bound refuses the redundant views at the door. The raw
// cosines this replaces, expand_novelty_sim and expand_track_spread_max, are
// retired under AR-024.
//
// The annex is CPU-side and in-memory: it is small (tens of embeddings) so the
// matcher scans it with a scalar loop, and it is discarded when the process
// exits. Promoted embeddings only help SUBSEQUENT frames and later tracks of A —
// the pipeline stays streaming, no emitted output is buffered or relabelled.

struct TrackGallery {
    // One promoted reference view held in the per-actor annex.
    struct AnnexEntry {
        Embedding emb;
        int       actor_idx{-1};
    };

    explicit TrackGallery(const Config& cfg)
        : enabled_(cfg.expand_gallery)
        , buffer_size_(std::max(1, cfg.expand_buffer_size))
        , band_lo_(cfg.expand_band_lo)
        , band_hi_(cfg.expand_band_hi)
        , min_anchor_frames_(std::max(1, cfg.expand_min_anchor_frames))
        , debug_dir_(cfg.expand_debug_dir)
    {
        if (!enabled_) return;
        std::cerr << "[track_gallery] per-film expansion ON"
                  << "  buffer=" << buffer_size_
                  << "  band=[" << band_lo_ << ", " << band_hi_ << "]"
                  << "  min_anchor_frames=" << min_anchor_frames_;
        if (!debug_dir_.empty()) {
            std::filesystem::create_directories(debug_dir_);
            std::cerr << "  debug_dir=" << debug_dir_;
        }
        std::cerr << "\n";
    }

    bool enabled() const { return enabled_; }

    // Current annex contents (empty when disabled). The matcher scans these
    // alongside the baked gallery so a promoted view can win best-of-N for its
    // actor. Returned by const-ref; only grows, never reordered.
    const std::vector<AnnexEntry>& annex() const { return annex_; }

    // Offer one observed face to its track's diversity buffer.
    //   track_id     : face_tracker track (−1 = untracked, ignored)
    //   emb          : this frame's raw embedding
    //   best_actor   : actor with the highest gallery similarity for this face
    //   best_gal_sim : that similarity (best sim to best_actor's baked+annex
    //                  refs) — a raw cosine, the last one in this class: it is
    //                  calibrated on entry and only the probability is stored
    //   accepted     : true if the matcher accepted this face as best_actor
    //   crop         : aligned crop, retained only when debug dumping is on
    void observe(int track_id, const Embedding& emb,
                 int best_actor, float best_gal_sim, bool accepted,
                 const cv::Mat& crop)
    {
        if (!enabled_ || track_id < 0) return;

        TrackState& ts = tracks_[track_id];

        // Vote toward ownership: only accepted frames name an actor, and a track
        // that flip-flops between actors is ambiguous, so we tally per actor and
        // pick the plurality winner at confirmation time.
        if (accepted && best_actor >= 0) {
            ts.actor_votes[best_actor]++;
            ts.accepted_frames++;
        }

        insert_into_buffer(ts, emb, best_gal_sim, crop);

        // Confirm and promote as soon as the anchor threshold is met, once.
        if (!ts.promoted && ts.accepted_frames >= min_anchor_frames_)
            promote(track_id, ts);
    }

    // Drop a track's buffer when the face_tracker expires it or on a scene cut,
    // so stale/cross-cut embeddings can never be promoted later. Called by the
    // matcher when it observes a cut or track disappearance.
    void forget(int track_id) { tracks_.erase(track_id); }

    /// TRACES: AR-019 | SR-005
    /// The registry's verdict on who this track is. Authoritative: it comes from
    /// the Bayesian accumulation (AR-025), where the local tally counted raw
    /// accepted frames and so weighted thirty near-identical looks the same as
    /// thirty distinct ones.
    void set_owner(int track_id, int actor_idx) {
        if (track_id < 0 || actor_idx < 0) return;
        tracks_[track_id].registry_owner = actor_idx;
    }

    /// TRACES: AR-024 | SR-005
    /// Supply the calibration belonging to the active embedder. Without it the
    /// band falls back to treating cosine as probability, which is wrong but
    /// bounded — and the default is loud in the header rather than silent.
    void set_calibration(std::function<float(float)> c) { calibrate_ = std::move(c); }

    /// Embeddings the band refused. A store that admits nothing is as wrong as
    /// one that admits everything, and neither is visible without this.
    std::size_t band_rejected() const { return rejected_; }

    // Drop every track buffer (scene cut / EOF). Mirrors face_tracker's clear.
    void clear_tracks() { tracks_.clear(); }

private:
    struct BufEntry {
        Embedding emb;
        /// P(same person) against the owning actor's refs when observed —
        /// calibrated at the door (AR-024), so the eviction ordering below is a
        /// comparison of probabilities and the struct holds no bare cosine.
        float     gal_p{0.f};
        cv::Mat   crop;           // populated only when debug_dir_ set
    };

    struct TrackState {
        std::vector<BufEntry> buf;
        std::map<int, int>    actor_votes;   // actor_idx → accepted-frame count
        int                   accepted_frames{0};
        bool                  promoted{false};
        int                   registry_owner{-1};   ///< AR-019: authoritative
    };

    /// TRACES: AR-018, AR-024 | SR-005
    /// Banded admission: an embedding joins the store only if its similarity to
    /// something already there falls **inside a band**.
    ///
    ///   above the upper bound → redundant. It is another look at a pose the
    ///     store already covers, and adding it teaches the annex nothing while
    ///     costing a slot that a novel view could have used.
    ///   below the lower bound → suspect. Within one track every face is the
    ///     same person by construction, so an embedding unlike everything else
    ///     on the track is evidence the construction failed — a track-ID
    ///     collision or a bad detection. Admitting it is how an actor's annex
    ///     gets poisoned with someone else's face.
    ///
    /// Both bounds are calibrated probabilities, never raw cosines (AR-024): a
    /// bare similarity threshold means something different for every model and
    /// every face size, and this gate has to hold across both.
    ///
    /// The first embedding is always admitted — there is nothing for it to be
    /// redundant with, and nothing to contradict it.
    bool admit(const TrackState& ts, const Embedding& emb) const {
        if (ts.buf.empty()) return true;

        float p_max = 0.f;
        for (const auto& b : ts.buf)
            p_max = std::max(p_max, calibrate_(cosine_similarity(b.emb, emb)));

        return p_max >= band_lo_ && p_max <= band_hi_;
    }

    void insert_into_buffer(TrackState& ts, const Embedding& emb,
                            float gal_sim, const cv::Mat& crop)
    {
        if (!admit(ts, emb)) { ++rejected_; return; }

        BufEntry e;
        e.emb   = emb;
        e.gal_p = calibrate_(gal_sim);
        if (!debug_dir_.empty() && !crop.empty()) e.crop = crop.clone();

        if (static_cast<int>(ts.buf.size()) < buffer_size_) {
            ts.buf.push_back(std::move(e));
            return;
        }

        // Buffer full: evict the member the gallery recognises best (highest
        // gal_p) — least informative — but only if the newcomer is at least as
        // novel. Keeping the most gallery-far views is the whole point.
        //
        // This is an *ordering*, not a threshold: there is no constant to tune,
        // and novelty-seeking lives here rather than in a cutoff. It ranks
        // probabilities, so it says the same thing across models (AR-024).
        int   worst_i = -1;
        float worst_p = e.gal_p;   // newcomer's probability is the bar to beat
        for (int i = 0; i < static_cast<int>(ts.buf.size()); ++i) {
            if (ts.buf[i].gal_p > worst_p) {
                worst_p = ts.buf[i].gal_p;
                worst_i   = i;
            }
        }
        // worst_i == −1 → every buffered view is already more novel than the
        // newcomer; drop the newcomer instead of a better sample.
        if (worst_i >= 0) ts.buf[worst_i] = std::move(e);
    }

    void promote(int track_id, TrackState& ts) {
        ts.promoted = true;   // idempotent: never promote a track twice

        int actor = owning_actor(ts);
        if (actor < 0) return;

        // ── Safety gate: the band's lower bound, across the whole store ──────
        float worst = store_coherence(ts.buf);
        if (worst < band_lo_) {
            std::cerr << "[track_gallery] track " << track_id
                      << " → actor " << actor
                      << " REJECTED (worst pairwise P=" << worst
                      << " < " << band_lo_ << ", likely ID collision)\n";
            return;
        }

        int added = 0;
        for (const auto& be : ts.buf) {
            annex_.push_back({be.emb, actor});
            if (!debug_dir_.empty() && !be.crop.empty())
                dump_mugshot(track_id, actor, added, be);
            ++added;
        }

        std::cerr << "[track_gallery] track " << track_id
                  << " confirmed actor " << actor
                  << " (" << ts.accepted_frames << " accepted frames, worst "
                  << "pairwise P=" << worst << ") — promoted " << added
                  << " views; annex now " << annex_.size() << "\n";
    }

    /// Prefer the registry's verdict; fall back to the local tally only when no
    /// registry is attached (unit tests, replay harness).
    static int owning_actor(const TrackState& ts) {
        if (ts.registry_owner >= 0) return ts.registry_owner;
        return plurality_actor(ts);
    }

    static int plurality_actor(const TrackState& ts) {
        int best = -1, best_votes = 0;
        for (const auto& [ai, v] : ts.actor_votes) {
            if (v > best_votes) { best_votes = v; best = ai; }
        }
        return best;
    }

    /// TRACES: AR-018, AR-024 | SR-005
    /// The store's weakest pairwise P(same person) — the band's lower bound
    /// asked of every pair, not just of the best match at the door.
    ///
    /// `admit` compares a newcomer against its *closest* existing member, so a
    /// track that drifts gradually can chain A→B→C with every step inside the
    /// band while A and C are strangers. That is precisely the shape a track-ID
    /// collision takes when two people are merged over a slow pan, so the bound
    /// is re-asked here across all pairs before anything reaches an actor's
    /// annex. Same bound, same probability space — not a second constant.
    ///
    /// A store of one has no pair to disagree; it is coherent by construction,
    /// hence 1.
    float store_coherence(const std::vector<BufEntry>& buf) const {
        float worst = std::numeric_limits<float>::max();
        for (size_t i = 0; i < buf.size(); ++i)
            for (size_t j = i + 1; j < buf.size(); ++j)
                worst = std::min(worst,
                                 calibrate_(cosine_similarity(buf[i].emb, buf[j].emb)));
        if (worst == std::numeric_limits<float>::max()) return 1.f;
        return worst;
    }

    void dump_mugshot(int track_id, int actor, int idx, const BufEntry& be) {
#ifdef SAE_DEBUG
        char name[64];
        std::snprintf(name, sizeof(name), "trk%d_actor%d_%d_p%.3f.jpg",
                      track_id, actor, idx, be.gal_p);
        cv::imwrite((std::filesystem::path(debug_dir_) / name).string(), be.crop);
#else
        (void)track_id; (void)actor; (void)idx; (void)be;
#endif
    }

    /// cosine → P(same person). The one probability space the pipeline reasons
    /// in; see gallery_calibration.hpp's same_person_probability.
    std::function<float(float)> calibrate_{[](float c) { return std::max(0.f, c); }};
    std::size_t             rejected_{0};   ///< admissions refused by the band

    bool                    enabled_;
    int                     buffer_size_;
    float                   band_lo_;       ///< AR-018, from cfg.expand_band_lo
    float                   band_hi_;       ///< AR-018, from cfg.expand_band_hi
    int                     min_anchor_frames_;
    std::string             debug_dir_;

    std::map<int, TrackState> tracks_;
    std::vector<AnnexEntry>   annex_;
};