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

#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, after two safety gates:
//        • novelty: only embeddings whose best sim to A's refs is below
//          expand_novelty_sim are added (skip poses already covered);
//        • spread: if the retained buffer's internal spread (1 − min pairwise
//          cosine sim) exceeds expand_track_spread_max the whole track is
//          rejected — such spread signals a track-ID collision merging two
//          people, whose embeddings must never enter A's annex.
//
// 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))
        , novelty_sim_(cfg.expand_novelty_sim)
        , spread_max_(cfg.expand_track_spread_max)
        , 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_
                  << "  novelty_sim<" << novelty_sim_
                  << "  spread_max=" << spread_max_
                  << "  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)
    //   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); }

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

private:
    struct BufEntry {
        Embedding emb;
        float     gal_sim{0.f};   // best sim to owning actor's refs when observed
        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};
    };

    void insert_into_buffer(TrackState& ts, const Embedding& emb,
                            float gal_sim, const cv::Mat& crop)
    {
        BufEntry e;
        e.emb     = emb;
        e.gal_sim = 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_sim) — least informative — but only if the newcomer is at least as
        // novel. Keeping the most gallery-far views is the whole point.
        int   worst_i   = -1;
        float worst_sim = e.gal_sim;   // newcomer's sim is the bar to beat
        for (int i = 0; i < static_cast<int>(ts.buf.size()); ++i) {
            if (ts.buf[i].gal_sim > worst_sim) {
                worst_sim = ts.buf[i].gal_sim;
                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 = plurality_actor(ts);
        if (actor < 0) return;

        // ── Safety gate: internal spread ─────────────────────────────────────
        // A legitimate single-person track varies in pose but stays reasonably
        // self-similar. Large spread signals two people merged under one track
        // ID — reject the whole track rather than poison the actor's annex.
        float spread = buffer_spread(ts.buf);
        if (spread > spread_max_) {
            std::cerr << "[track_gallery] track " << track_id
                      << " → actor " << actor
                      << " REJECTED (spread " << spread
                      << " > " << spread_max_ << ", likely ID collision)\n";
            return;
        }

        int added = 0;
        for (const auto& be : ts.buf) {
            // ── Safety gate: novelty ─────────────────────────────────────────
            // Skip poses the gallery already covers; only gallery-far views are
            // worth the annex slot (and the extra per-frame scan cost).
            if (be.gal_sim >= novelty_sim_) continue;
            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, spread "
                  << spread << ") — promoted " << added << "/"
                  << ts.buf.size() << " views; annex now "
                  << annex_.size() << "\n";
    }

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

    // Spread = 1 − min pairwise cosine similarity over the buffer (0 when <2).
    static float buffer_spread(const std::vector<BufEntry>& buf) {
        float min_sim = std::numeric_limits<float>::max();
        for (size_t i = 0; i < buf.size(); ++i)
            for (size_t j = i + 1; j < buf.size(); ++j)
                min_sim = std::min(min_sim, cosine_similarity(buf[i].emb, buf[j].emb));
        if (min_sim == std::numeric_limits<float>::max()) return 0.f;
        return 1.f - min_sim;
    }

    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_sim%.3f.jpg",
                      track_id, actor, idx, be.gal_sim);
        cv::imwrite((std::filesystem::path(debug_dir_) / name).string(), be.crop);
#else
        (void)track_id; (void)actor; (void)idx; (void)be;
#endif
    }

    bool                    enabled_;
    int                     buffer_size_;
    float                   novelty_sim_;
    float                   spread_max_;
    int                     min_anchor_frames_;
    std::string             debug_dir_;

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