// sae_kpn — run the real downstream pipeline nodes (face_tracker, identity_matcher,
// scene_tracker) inside a Python-assembled KPN network, fed by a Python HDF5 replay
// source. Lets a parameter sweep re-run the exact C++ matching/tracking logic over
// dumped embeddings — no video decode, no GPU — with different Config knobs each run.
//
// Boundary types (cross the Python seam):
//   EmbeddedSceneFrame  IN   (built by the Python replay source from HDF5 arrays)
//   SceneAnnotation     OUT  (read by the Python sink → presence JSON)
// Intermediate types (TrackedSceneFrame, MatchedSceneFrame) flow C++→C++ only, but
// still need channel factories + converters registered so PyNetwork can wire them.

#define KPN_BUILD_PYTHON
#include <kpn/python/bindings.hpp>
#include <kpn/python/object_variant_node.hpp>

#include "types.hpp"
#include "config.hpp"
#include "gallery/embedder_stamp.hpp"
#include "gallery/gallery_store.hpp"
#include "nodes/face_tracker_node.hpp"
#include "nodes/identity_matcher_node.hpp"
#include "nodes/scene_tracker_node.hpp"

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>
#include <nanobind/stl/map.h>

#include <memory>
#include <variant>

namespace nb = nanobind;
using namespace nb::literals;

// The variant spanning every type that flows on a channel in the replay chain.
using SaeVariant = std::variant<EmbeddedSceneFrame, TrackedSceneFrame,
                                MatchedSceneFrame, SceneAnnotation>;

// ── Converters ─────────────────────────────────────────────────────────────────
// Only EmbeddedSceneFrame (in) and SceneAnnotation (out) actually cross the seam;
// the two intermediates get identity-ish stubs (never converted in practice) so the
// variant's converter map is complete.

namespace kpn {

// EmbeddedSceneFrame: built FROM Python (a dict of numpy arrays). to_python is a
// stub (the replay source only produces it; nothing reads it back).
template<> struct PythonConverter<EmbeddedSceneFrame> {
    static constexpr const char* type_name = "EmbeddedSceneFrame";

    static nb::object to_python(const EmbeddedSceneFrame&) {
        // Not needed downstream; return None. (Kept total for map completeness.)
        return nb::none();
    }

    static EmbeddedSceneFrame from_python(nb::object o) {
        nb::dict d = nb::cast<nb::dict>(o);
        EmbeddedSceneFrame ef;
        ef.source.timestamp_sec = nb::cast<double>(d["timestamp_sec"]);
        ef.source.frame_idx     = d.contains("frame_idx") ? nb::cast<int64_t>(d["frame_idx"]) : -1;
        ef.source.eof           = d.contains("eof") ? nb::cast<bool>(d["eof"]) : false;
        ef.source.is_cut        = d.contains("is_cut") ? nb::cast<bool>(d["is_cut"]) : false;
        if (ef.source.eof) return ef;

        // faces: (N,4) bbox, (N,10) landmarks, (N,) confidence, (N,512) embeddings
        auto bbox = nb::cast<nb::ndarray<float, nb::shape<-1, 4>, nb::c_contig>>(d["bbox"]);
        auto lmk  = nb::cast<nb::ndarray<float, nb::shape<-1, 10>, nb::c_contig>>(d["landmarks"]);
        auto conf = nb::cast<nb::ndarray<float, nb::shape<-1>, nb::c_contig>>(d["confidence"]);
        auto emb  = nb::cast<nb::ndarray<float, nb::shape<-1, 512>, nb::c_contig>>(d["embeddings"]);

        const size_t n = bbox.shape(0);
        ef.faces.reserve(n);
        ef.embeddings.reserve(n);
        const float* bp = bbox.data();
        const float* lp = lmk.data();
        const float* cp = conf.data();
        const float* ep = emb.data();
        for (size_t i = 0; i < n; ++i) {
            DetectedFace f;
            f.bbox = cv::Rect2f(bp[i*4+0], bp[i*4+1], bp[i*4+2], bp[i*4+3]);
            for (int k = 0; k < 5; ++k)
                f.landmarks[k] = cv::Point2f(lp[i*10 + k*2], lp[i*10 + k*2 + 1]);
            f.confidence = cp[i];
            ef.faces.push_back(f);

            Embedding e;
            for (int k = 0; k < 512; ++k) e[k] = ep[i*512 + k];
            ef.embeddings.push_back(e);
        }
        // crops left empty: tracker/matcher only forward them for debug rendering.
        ef.crops.resize(n);
        return ef;
    }
};

// SceneAnnotation: read INTO Python. from_python is a stub (Python never builds one).
template<> struct PythonConverter<SceneAnnotation> {
    static constexpr const char* type_name = "SceneAnnotation";

    static nb::object to_python(const SceneAnnotation& sa) {
        nb::dict d;
        d["timestamp_sec"] = sa.timestamp_sec;
        d["eof"]           = sa.eof;
        nb::list actors;
        for (const auto& a : sa.visible_actors) {
            nb::dict ad;
            ad["actor_idx"]   = a.actor_idx;
            ad["track_id"]    = a.track_id;
            ad["name"]        = a.name;
            ad["imdb_id"]     = a.imdb_id;
            ad["tmdb_id"]     = a.tmdb_id;
            ad["jellyfin_id"] = a.jellyfin_id;
            ad["similarity"]  = a.similarity;
            ad["bbox"] = nb::make_tuple(a.bbox.x, a.bbox.y, a.bbox.width, a.bbox.height);
            actors.append(ad);
        }
        d["visible_actors"] = actors;
        return d;
    }

    static SceneAnnotation from_python(nb::object) {
        return {};  // never called
    }
};

// Intermediates: never cross the seam. Provide stubs so register_full_type compiles.
template<> struct PythonConverter<TrackedSceneFrame> {
    static constexpr const char* type_name = "TrackedSceneFrame";
    static nb::object to_python(const TrackedSceneFrame&) { return nb::none(); }
    static TrackedSceneFrame from_python(nb::object) { return {}; }
};
template<> struct PythonConverter<MatchedSceneFrame> {
    static constexpr const char* type_name = "MatchedSceneFrame";
    static nb::object to_python(const MatchedSceneFrame&) { return nb::none(); }
    static MatchedSceneFrame from_python(nb::object) { return {}; }
};

} // namespace kpn

// ── Config from Python dict ─────────────────────────────────────────────────────
// Only the knobs relevant to the replayed chain; everything else keeps its default.

static Config config_from_dict(nb::dict d) {
    Config cfg;
    auto getf = [&](const char* k, float& dst) { if (d.contains(k)) dst = nb::cast<float>(d[k]); };
    auto geti = [&](const char* k, int& dst)   { if (d.contains(k)) dst = nb::cast<int>(d[k]); };
    auto getd = [&](const char* k, double& dst){ if (d.contains(k)) dst = nb::cast<double>(d[k]); };
    // identity matcher
    getf("match_prior",       cfg.match_prior);
    getf("prob_threshold",    cfg.prob_threshold);
    getf("match_threshold",   cfg.match_threshold);
    getf("match_ratio",       cfg.match_ratio);
    getf("match_ratio_ceil",  cfg.match_ratio_ceil);
    // face tracker
    getf("track_alpha",           cfg.track_alpha);
    getf("track_min_iou",         cfg.track_min_iou);
    getf("track_max_embed_dist",  cfg.track_max_embed_dist);
    geti("track_max_frames_missing", cfg.track_max_frames_missing);
    getf("cut_revive_sim",        cfg.cut_revive_sim);
    geti("cut_inactive_max_frames", cfg.cut_inactive_max_frames);
    // scene tracker
    getd("extinction_sec", cfg.extinction_sec);
    getd("anneal_sec",     cfg.anneal_sec);
    // gallery expansion (usually off for sweeps; expose so it can be toggled)
    if (d.contains("expand_gallery")) cfg.expand_gallery = nb::cast<bool>(d["expand_gallery"]);
    /// TRACES: GR-004 | SR-001
    if (d.contains("require_gallery_stamp"))
        cfg.require_gallery_stamp = nb::cast<bool>(d["require_gallery_stamp"]);
    return cfg;
}

using Net = kpn::python::PyNetwork<SaeVariant>;

NB_MODULE(sae_kpn, m) {
    m.doc() = "Real KPN downstream nodes (tracker/matcher/scene_tracker) for Python replay sweeps";

    kpn::python::register_py_network<SaeVariant>(m, "Network");

    // Register converters + channel factories for all four channel types on a net.
    // register_py_network doesn't do this (auto_bind does); we patch __init__ to.
    // Simpler: expose a free helper the Python side calls right after construction.
    m.def("_register_types", [](Net& net) {
        net.register_full_type<EmbeddedSceneFrame>(
            [](const EmbeddedSceneFrame& v){ return kpn::PythonConverter<EmbeddedSceneFrame>::to_python(v); },
            [](nb::object o){ return kpn::PythonConverter<EmbeddedSceneFrame>::from_python(std::move(o)); },
            "EmbeddedSceneFrame");
        net.register_full_type<TrackedSceneFrame>(
            [](const TrackedSceneFrame& v){ return kpn::PythonConverter<TrackedSceneFrame>::to_python(v); },
            [](nb::object o){ return kpn::PythonConverter<TrackedSceneFrame>::from_python(std::move(o)); },
            "TrackedSceneFrame");
        net.register_full_type<MatchedSceneFrame>(
            [](const MatchedSceneFrame& v){ return kpn::PythonConverter<MatchedSceneFrame>::to_python(v); },
            [](nb::object o){ return kpn::PythonConverter<MatchedSceneFrame>::from_python(std::move(o)); },
            "MatchedSceneFrame");
        net.register_full_type<SceneAnnotation>(
            [](const SceneAnnotation& v){ return kpn::PythonConverter<SceneAnnotation>::to_python(v); },
            [](nb::object o){ return kpn::PythonConverter<SceneAnnotation>::from_python(std::move(o)); },
            "SceneAnnotation");
    });

    m.def("add_node_python", [](Net& net, std::string name, nb::object callable,
                                std::vector<std::string> ins, std::vector<std::string> outs,
                                std::size_t cap) {
        net.add_node_python(std::move(name), std::move(callable), std::move(ins),
                            std::move(outs), cap);
    }, "net"_a, "name"_a, "callable"_a, "inputs"_a, "outputs"_a, "capacity"_a = 5);

    // ── Real node factories ─────────────────────────────────────────────────────
    m.def("add_face_tracker", [](Net& net, std::string name, nb::dict cfg_dict, std::size_t cap) {
        Config cfg = config_from_dict(cfg_dict);
        auto node = std::make_shared<kpn::ObjectVariantNodeWrapper<
            FaceTrackerFunc, SaeVariant, kpn::in<"embedded">, kpn::out<"tracked">>>(cap, cfg);
        net.add(std::move(name), std::move(node));
    }, "net"_a, "name"_a, "config"_a, "capacity"_a = 16);

    /// TRACES: GR-004 | SR-001
    // embedder_model / embedder_sha256 identify whatever produced the embeddings
    // that will be fed in. In a replay those come from the dump's own stamp (see
    // scripts/optimizer/SCHEMA.md), because there is no live embedder in the
    // network — the dump *is* the embedder as far as this gallery is concerned.
    // Passing neither leaves the binding unverifiable, which warns loudly and is
    // fatal under SAE_REQUIRE_GALLERY_STAMP.
    m.def("add_identity_matcher", [](Net& net, std::string name, std::string gallery_path,
                                     nb::dict cfg_dict, std::size_t cap,
                                     std::string embedder_model,
                                     std::string embedder_sha256) {
        Config cfg = config_from_dict(cfg_dict);
        cfg.gallery_path = gallery_path;  // needed to persist refreshed calibration back
        // Cache loaded galleries by path so a threshold sweep (many networks, same
        // gallery) pays the ~24s JSON parse only once. The matcher holds a const
        // ref; the cache keeps the gallery alive for the process lifetime.
        static std::map<std::string, std::shared_ptr<ActorGallery>> cache;
        auto it = cache.find(gallery_path);
        if (it == cache.end())
            it = cache.emplace(gallery_path,
                     std::make_shared<ActorGallery>(load_gallery(gallery_path))).first;

        // Checked on every construction, not only on the cache miss: the same
        // process may replay several dumps against one cached gallery.
        EmbedderStamp feeding;
        feeding.model_name   = std::move(embedder_model);
        feeding.model_sha256 = std::move(embedder_sha256);
        enforce_embedder_stamp(it->second->embedder, feeding, gallery_path,
                               feeding.model_name.empty()
                                   ? "embeddings fed into this network"
                                   : feeding.model_name,
                               cfg.require_gallery_stamp);

        auto node = std::make_shared<kpn::ObjectVariantNodeWrapper<
            IdentityMatcherFunc, SaeVariant, kpn::in<"tracked">, kpn::out<"matched">>>(
            cap, *it->second, cfg);
        net.add(std::move(name), std::move(node));
    }, "net"_a, "name"_a, "gallery"_a, "config"_a, "capacity"_a = 16,
       "embedder_model"_a = "", "embedder_sha256"_a = "");

    m.def("add_scene_tracker", [](Net& net, std::string name, nb::dict cfg_dict, std::size_t cap) {
        Config cfg = config_from_dict(cfg_dict);
        auto node = std::make_shared<kpn::ObjectVariantNodeWrapper<
            SceneTrackerFunc, SaeVariant, kpn::in<"matched">, kpn::out<"annotation">>>(cap, cfg);
        net.add(std::move(name), std::move(node));
    }, "net"_a, "name"_a, "config"_a, "capacity"_a = 16);

    // ── Runtime setters (persistent-pipeline reuse across a threshold sweep) ─────
    // Build the network once, then change thresholds between replays — no rebuild,
    // no teardown (which is where the ROCm deadlock lives), no gallery reload.
    using MatcherWrap = kpn::ObjectVariantNodeWrapper<
        IdentityMatcherFunc, SaeVariant, kpn::in<"tracked">, kpn::out<"matched">>;
    using SceneWrap = kpn::ObjectVariantNodeWrapper<
        SceneTrackerFunc, SaeVariant, kpn::in<"matched">, kpn::out<"annotation">>;

    m.def("set_prob_threshold", [](Net& net, std::string name, float t) {
        auto* w = dynamic_cast<MatcherWrap*>(net.node_ptr(name));
        if (!w) throw std::runtime_error("set_prob_threshold: '" + name + "' is not an identity_matcher");
        w->functor().set_prob_threshold(t);
    }, "net"_a, "name"_a, "value"_a);

    m.def("set_extinction_sec", [](Net& net, std::string name, double s) {
        auto* w = dynamic_cast<SceneWrap*>(net.node_ptr(name));
        if (!w) throw std::runtime_error("set_extinction_sec: '" + name + "' is not a scene_tracker");
        w->functor().set_extinction_sec(s);
    }, "net"_a, "name"_a, "value"_a);
}