#pragma once
#include "scrfd_decoder.hpp"
#include "trt_scrfd_decoder.hpp"
#include "config.hpp"
#include "ort_provider.hpp"

#include <memory>
#include <string>

// ── FaceDetectorFunc ──────────────────────────────────────────────────────────
// KPN node: runs SCRFD-500MF to detect ALL faces in a frame.
//
// Backend selection:
//   --detector-engine <path>  → TrtScrfdDecoder (raw TensorRT, no ORT)
//   otherwise                 → SCRFDDecoder (ONNX Runtime)

struct FaceDetectorFunc {
    static constexpr std::string_view label() { return "face_detector"; }

    explicit FaceDetectorFunc(const Config& cfg, OrtProvider provider)
        : max_faces_(cfg.max_faces)
        , min_face_px_(cfg.min_face_px)
    {
        if (!cfg.detector_engine.empty()) {
            trt_ = std::make_unique<TrtScrfdDecoder>(
                cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
        } else {
            ort_ = std::make_unique<SCRFDDecoder>(
                cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
        }
    }

    SceneFrame operator()(Frame f) {
        if (f.eof) return {std::move(f), {}};

        auto faces = trt_ ? trt_->detect(f.image) : ort_->detect(f.image);

        // Drop faces below minimum pixel size (too small for reliable ArcFace alignment)
        faces.erase(
            std::remove_if(faces.begin(), faces.end(), [&](const DetectedFace& d) {
                return d.bbox.width < min_face_px_ || d.bbox.height < min_face_px_;
            }),
            faces.end());

        // Sort largest-first so max_faces_ keeps the most informative detections
        std::sort(faces.begin(), faces.end(),
                  [](const DetectedFace& a, const DetectedFace& b) {
                      return a.bbox.area() > b.bbox.area();
                  });
        if (static_cast<int>(faces.size()) > max_faces_)
            faces.resize(max_faces_);

        return {std::move(f), std::move(faces)};
    }

private:
    std::unique_ptr<SCRFDDecoder>    ort_;
    std::unique_ptr<TrtScrfdDecoder> trt_;
    int                              max_faces_{10};
    float                            min_face_px_{40.f};
};