#pragma once
#include "config.hpp"
#include "inference/face_embedder.hpp"

#include <algorithm>
#include <memory>
#include <stdexcept>
#include <string>

// ── EmbedderFunc ──────────────────────────────────────────────────────────────
// KPN node: runs ArcFace on every 112×112 crop in an AlignedSceneFrame,
// producing one L2-normalised 512-dim embedding per face.
//
// The inference backend (ONNX Runtime or raw TensorRT) is selected at compile
// time; this node talks only to IFaceEmbedder via make_face_embedder(cfg).
//
// All crops in one frame are batched into a single forward pass (capped at
// embed_batch_size). The backend serialises itself; we only call it from the
// single embedder thread.
//
/// TRACES: AR-006 | SR-002
struct EmbedderFunc {
    static constexpr std::string_view label() { return "embedder"; }

    explicit EmbedderFunc(const Config& cfg)
        : embedder_(make_face_embedder(cfg))
        , batch_size_(std::max(1, cfg.embed_batch_size))
    {
        if (embedder_->max_batch() < static_cast<int>(batch_size_))
            throw std::runtime_error(
                "embed_batch_size " + std::to_string(batch_size_) +
                " exceeds backend max_batch " + std::to_string(embedder_->max_batch()) +
                " — rebuild the engine with EMBED_BATCH=" + std::to_string(batch_size_));
    }

    EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
        if (af.source.eof || af.crops.empty())
            return {std::move(af.source), {}, {}, {}};

        const auto& crops = af.crops;
        std::vector<Embedding> embeddings;
        embeddings.reserve(crops.size());
        for (size_t i = 0; i < crops.size(); i += batch_size_) {
            const size_t end = std::min(i + batch_size_, crops.size());
            std::vector<cv::Mat> chunk_crops(crops.begin() + i, crops.begin() + end);
            auto chunk = embedder_->embed(chunk_crops);
            embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
        }

        return {std::move(af.source),
                std::move(af.faces),
                std::move(af.crops),
                std::move(embeddings)};
    }

private:
    std::unique_ptr<IFaceEmbedder> embedder_;
    size_t batch_size_;
};