#pragma once
#include "types.hpp"
#include "config.hpp"
#include "inference/similarity.hpp"
#include "gallery/gallery_store.hpp"
#include "gallery/gallery_calibration.hpp"

#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>

// ── IdentityMatcherFunc ───────────────────────────────────────────────────────
// KPN node: compares each embedding against every reference embedding in the
// actor gallery using cosine similarity.
//
// Matching strategy — two modes selected at construction time:
//
//   Calibrated (preferred): gallery calibration fits a sigmoid
//     P(match) = σ(a·similarity + b) from intra/inter-class pairs.
//     A face is accepted if P(match | best_actor) > prob_threshold.
//
//   Fallback (no calibration): dual-criterion accept —
//     (a) best cosine distance < match_threshold, OR
//     (b) ratio test: best_dist/second_best_dist < match_ratio
//         AND best_dist < match_ratio_ceil.
//
// In both modes, per-actor best similarity is determined by scanning
// reference embeddings and taking the closest (best-of-N).
//
// Gallery scan: the full reference set (tens of thousands of 512-dim
// embeddings) is uploaded to the GPU once at construction time and stays
// resident there. Per frame, only the small query matrix (n_faces x 512) is
// uploaded and a single SGEMM computes the full similarity matrix in well under
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.

struct IdentityMatcherFunc {
    static constexpr std::string_view label() { return "identity_matcher"; }

    // Max faces handled per frame without reallocating GPU buffers.
    static constexpr int kMaxFaces = 32;

    IdentityMatcherFunc(const ActorGallery& gallery, const Config& cfg)
        : gallery_(gallery)
        , prob_threshold_(cfg.prob_threshold)
        , log_prior_odds_(std::log(cfg.match_prior / (1.f - cfg.match_prior)))
        , threshold_(cfg.match_threshold)
        , ratio_(cfg.match_ratio)
        , ratio_ceil_(cfg.match_ratio_ceil)
    {
        std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
        for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
            for (const auto& emb : gallery_.actors[ai].embeddings) {
                flat_emb_.push_back(emb);
                flat_actor_.push_back(ai);
            }
        }
        n_gallery_ = static_cast<int>(flat_emb_.size());

        std::cerr << "[identity_matcher] starting calibration ("
                  << flat_emb_.size() << " embeddings)...\n";
        cal_ = calibrate_gallery_cached(flat_emb_, flat_actor_,
                                         cfg.gallery_path + ".calib_cache.json");

        if (cal_.valid) {
            std::cerr << "[identity_matcher] calibrated Bayesian matching"
                      << "  prior=" << cfg.match_prior
                      << "  P_threshold=" << prob_threshold_
                      << "  effective_sim_boundary="
                      << cal_.boundary_at(prob_threshold_, log_prior_odds_) << "\n";
        } else {
            std::cerr << "[identity_matcher] threshold matching (calibration skipped)"
                      << "  threshold=" << threshold_
                      << "  ratio=" << ratio_ << "  ratio_ceil=" << ratio_ceil_ << "\n";
        }
        std::cerr << "[identity_matcher] gallery: "
                  << gallery_.actors.size() << " actors, "
                  << flat_emb_.size() << " reference embeddings\n";

        std::vector<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
        for (int i = 0; i < n_gallery_; ++i)
            std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
                        flat_emb_[i].data(), 512 * sizeof(float));

        sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
    }

    MatchedSceneFrame operator()(TrackedSceneFrame tf) {
        if (tf.source.eof) return {std::move(tf.source), {}};

        const int n_faces = static_cast<int>(tf.embeddings.size());
        std::vector<IdentifiedActor> actors;
        actors.reserve(n_faces);

        if (n_faces == 0) return {std::move(tf.source), {}};
        if (n_faces > kMaxFaces)
            throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");

        std::vector<float> host_query(static_cast<size_t>(n_faces) * 512);
        for (int fi = 0; fi < n_faces; ++fi) {
            std::memcpy(host_query.data() + static_cast<size_t>(fi) * 512,
                        tf.embeddings[fi].data(), 512 * sizeof(float));
        }

        // S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
        const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);

        for (int fi = 0; fi < n_faces; ++fi) {
            const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;

            std::vector<float> best_sim(gallery_.actors.size(),
                                        -std::numeric_limits<float>::max());
            for (int ei = 0; ei < n_gallery_; ++ei) {
                float sim = sims[ei];
                int   ai  = flat_actor_[ei];
                if (sim > best_sim[ai]) best_sim[ai] = sim;
            }

            int   best_actor   = -1;
            int   second_actor = -1;
            float best_s       = -std::numeric_limits<float>::max();
            float second_s     = -std::numeric_limits<float>::max();
            for (int ai = 0; ai < static_cast<int>(best_sim.size()); ++ai) {
                if (best_sim[ai] > best_s) {
                    second_s     = best_s;
                    second_actor = best_actor;
                    best_s       = best_sim[ai];
                    best_actor   = ai;
                } else if (best_sim[ai] > second_s) {
                    second_s     = best_sim[ai];
                    second_actor = ai;
                }
            }
            (void)second_actor;

            bool accept = false;
            if (best_actor >= 0) {
                if (cal_.valid) {
                    accept = cal_.probability(best_s, log_prior_odds_) > prob_threshold_;
                } else {
                    float best_d   = 1.f - best_s;
                    float second_d = (second_s > -std::numeric_limits<float>::max())
                                     ? 1.f - second_s
                                     : std::numeric_limits<float>::max();
                    bool absolute = best_d < threshold_;
                    bool ratio    = (best_d < ratio_ceil_) &&
                                    (second_d == std::numeric_limits<float>::max() ||
                                     best_d / second_d < ratio_);
                    accept = absolute || ratio;
                }
            }

            IdentifiedActor ia;
            ia.bbox     = tf.faces[fi].bbox;
            ia.crop     = tf.crops[fi];
            ia.track_id = tf.track_ids[fi];

            if (accept) {
                ia.actor_idx   = best_actor;
                ia.name        = gallery_.actors[best_actor].name;
                ia.imdb_id     = gallery_.actors[best_actor].imdb_id;
                ia.tmdb_id     = gallery_.actors[best_actor].tmdb_id;
                ia.jellyfin_id = gallery_.actors[best_actor].jellyfin_id;
                ia.similarity = cal_.valid
                                ? cal_.probability(best_s, log_prior_odds_)
                                : best_s;
            }

            actors.push_back(std::move(ia));
        }

        return {std::move(tf.source), std::move(actors)};
    }

private:
    ActorGallery           gallery_;
    GalleryCalibration     cal_;
    float                  prob_threshold_;
    float                  log_prior_odds_;
    float                  threshold_;
    float                  ratio_;
    float                  ratio_ceil_;
    std::vector<Embedding> flat_emb_;
    std::vector<int>       flat_actor_;
    int                    n_gallery_{0};

    std::unique_ptr<ISimilarityEngine> sim_engine_;
};