// ── ORT inference backend ─────────────────────────────────────────────────────
// ONNX Runtime implementations of IFaceDetector (SCRFD) and IFaceEmbedder
// (ArcFace), plus the make_* factories the core links against. Selected at
// compile time by CMake when SAE_INFERENCE_BACKEND=ORT.
//
// This is the ONLY translation unit that includes onnxruntime headers; the core
// application never sees them.

#include "inference/face_detector.hpp"
#include "inference/face_embedder.hpp"
#include "inference/scene_detector.hpp"
#include "backends/ort_provider.hpp"
#include "config.hpp"
#include "face_utils.hpp"
#include "types.hpp"

#include <onnxruntime/onnxruntime_cxx_api.h>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

// ── SCRFDDecoder ──────────────────────────────────────────────────────────────
// ONNX Runtime SCRFD face detector with kps. Uses ORT (not cv::dnn) because
// OpenCV 4.x cannot load SCRFD's dynamic Shape nodes. ORT handles dynamic shapes
// natively and is thread-safe for concurrent Run() calls.
//
// Model output layout (9 tensors, InsightFace export order):
//   [0-2]  score_s8 / score_s16 / score_s32  — flat (N,)
//   [3-5]  bbox_s8  / bbox_s16  / bbox_s32   — flat (N*4,)  distance format
//   [6-8]  kps_s8   / kps_s16   / kps_s32    — flat (N*10,) distance format
// Landmark order: right-eye, left-eye, nose, right-mouth, left-mouth.
class SCRFDDecoder final : public IFaceDetector {
public:
    static constexpr int kInputW        = 640;
    static constexpr int kInputH        = 640;
    static constexpr int kAllStrides[4] = {8, 16, 32, 64};
    static constexpr int kAnchors       = 2;

    SCRFDDecoder(const std::string& model_path,
                 float conf_threshold, float nms_threshold,
                 OrtProvider provider, BackendConfig trt_cfg)
        : conf_threshold_(conf_threshold)
        , nms_threshold_(nms_threshold)
    {
        Ort::SessionOptions opts;
        opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        opts.SetIntraOpNumThreads(1);

        // SCRFD ONNX has a dynamic H/W input; we letterbox to 640×640 at
        // runtime, so pin the TRT-EP profile to that single shape.
        if (provider == OrtProvider::TensorRT) {
            if (trt_cfg.input_name.empty()) {
                Ort::SessionOptions probe_opts;
                probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
                Ort::Session probe(env_, model_path.c_str(), probe_opts);
                Ort::AllocatorWithDefaultOptions alloc;
                trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
            }
            if (trt_cfg.profile_min.empty()) {
                const std::string shape =
                    "1x3x" + std::to_string(kInputH) + "x" + std::to_string(kInputW);
                trt_cfg.profile_min = shape;
                trt_cfg.profile_opt = shape;
                trt_cfg.profile_max = shape;
            }
        }

        apply_ort_model_cache(opts, model_path, trt_cfg);
        apply_ort_provider(opts, provider, "SCRFDDecoder", trt_cfg);

        session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);

        Ort::AllocatorWithDefaultOptions alloc;
        auto in_name = session_->GetInputNameAllocated(0, alloc);
        input_name_  = in_name.get();

        const size_t n_out = session_->GetOutputCount();
        if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
            throw std::runtime_error(
                "[SCRFDDecoder] expected 9 or 12 outputs (kps-variant model), got "
                + std::to_string(n_out));
        fmc_ = static_cast<int>(n_out / 3);

        for (size_t i = 0; i < n_out; ++i) {
            auto name = session_->GetOutputNameAllocated(i, alloc);
            out_name_storage_.emplace_back(name.get());
        }
        for (auto& s : out_name_storage_)
            out_name_ptrs_.push_back(s.c_str());

        // Reject non-SCRFD models (e.g. YuNet, which also has 12 outputs).
        const int expected_last[3] = {1, 4, 10};
        for (size_t gi = 0; gi < 3; ++gi) {
            for (int si = 0; si < fmc_; ++si) {
                const size_t oi = gi * fmc_ + si;
                auto shape = session_->GetOutputTypeInfo(oi)
                                     .GetTensorTypeAndShapeInfo().GetShape();
                if (shape.empty() || shape.back() != expected_last[gi]) {
                    throw std::runtime_error(
                        "[SCRFDDecoder] model does not look like InsightFace SCRFD: "
                        "output '" + out_name_storage_[oi] + "' last-dim is "
                        + std::to_string(shape.empty() ? -1 : shape.back())
                        + ", expected " + std::to_string(expected_last[gi])
                        + ". Hint: pass scrfd_500m_bnkps.onnx, not yunet/*.onnx.");
                }
            }
        }

        std::cerr << "[SCRFDDecoder] loaded: " << model_path << "\n";
    }

    std::vector<DetectedFace> detect(const cv::Mat& img) override {
        const float scale = std::min(static_cast<float>(kInputW) / img.cols,
                                      static_cast<float>(kInputH) / img.rows);
        const int new_w = static_cast<int>(std::round(img.cols * scale));
        const int new_h = static_cast<int>(std::round(img.rows * scale));
        const int pad_x = (kInputW - new_w) / 2;
        const int pad_y = (kInputH - new_h) / 2;

        cv::Mat resized;
        cv::resize(img, resized, {new_w, new_h}, 0, 0, cv::INTER_LINEAR);
        cv::Mat letterboxed(kInputH, kInputW, img.type(),
                             cv::Scalar(114, 114, 114));
        resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));

        cv::Mat blob = cv::dnn::blobFromImage(
            letterboxed, 1.0 / 128.0, {kInputW, kInputH},
            cv::Scalar(127.5f, 127.5f, 127.5f),
            /*swapRB=*/true, /*crop=*/false, CV_32F);

        const std::array<int64_t, 4> in_shape = {1, 3, kInputH, kInputW};
        auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        auto in_tensor = Ort::Value::CreateTensor<float>(
            mem, blob.ptr<float>(), blob.total(),
            in_shape.data(), in_shape.size());

        const char* in_name_c = input_name_.c_str();
        auto outs = session_->Run(
            Ort::RunOptions{nullptr},
            &in_name_c, &in_tensor, 1,
            out_name_ptrs_.data(), out_name_ptrs_.size());

        std::vector<cv::Rect2d>                 raw_boxes;
        std::vector<float>                      raw_scores;
        std::vector<std::array<cv::Point2f, 5>> raw_kps;

        for (int si = 0; si < fmc_; ++si) {
            const int stride = kAllStrides[si];
            const int fh     = kInputH / stride;
            const int fw     = kInputW / stride;

            const float* s = outs[si].GetTensorData<float>();
            const float* b = outs[fmc_     + si].GetTensorData<float>();
            const float* k = outs[fmc_ * 2 + si].GetTensorData<float>();

            for (int r = 0; r < fh; ++r) {
                for (int c = 0; c < fw; ++c) {
                    for (int a = 0; a < kAnchors; ++a) {
                        const int   idx   = (r * fw + c) * kAnchors + a;
                        const float score = s[idx];
                        if (score < conf_threshold_) continue;

                        const float cx = static_cast<float>(c * stride);
                        const float cy = static_cast<float>(r * stride);

                        const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
                        const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };

                        const float x1 = to_img_x(cx - b[idx*4+0] * stride);
                        const float y1 = to_img_y(cy - b[idx*4+1] * stride);
                        const float x2 = to_img_x(cx + b[idx*4+2] * stride);
                        const float y2 = to_img_y(cy + b[idx*4+3] * stride);
                        raw_boxes.push_back({(double)x1, (double)y1,
                                             (double)(x2-x1), (double)(y2-y1)});
                        raw_scores.push_back(score);

                        std::array<cv::Point2f, 5> lms;
                        for (int p = 0; p < 5; ++p)
                            lms[p] = {to_img_x(cx + k[idx*10+p*2  ] * stride),
                                      to_img_y(cy + k[idx*10+p*2+1] * stride)};
                        raw_kps.push_back(lms);
                    }
                }
            }
        }

        std::vector<int> keep;
        cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep);

        const float img_w = static_cast<float>(img.cols);
        const float img_h = static_cast<float>(img.rows);

        std::vector<DetectedFace> faces;
        faces.reserve(keep.size());
        for (int i : keep) {
            const auto& rb = raw_boxes[i];
            DetectedFace f;
            const float x = std::max(0.f, (float)rb.x);
            const float y = std::max(0.f, (float)rb.y);
            f.bbox       = {x, y,
                            std::min((float)rb.width,  img_w - x),
                            std::min((float)rb.height, img_h - y)};
            f.confidence = raw_scores[i];
            f.landmarks  = raw_kps[i];
            faces.push_back(f);
        }
        return faces;
    }

private:
    float                           conf_threshold_;
    float                           nms_threshold_;
    int                             fmc_{3};
    Ort::Env                        env_{ORT_LOGGING_LEVEL_WARNING, "scrfd"};
    std::unique_ptr<Ort::Session>   session_;
    std::string                     input_name_;
    std::vector<std::string>        out_name_storage_;
    std::vector<const char*>        out_name_ptrs_;
};

// ── ArcFaceEmbedder ───────────────────────────────────────────────────────────
// ONNX Runtime ArcFace embedder (w600k_r50, mbf, r18).
//   Input:  [N, 3, 112, 112] float32, BGR→RGB, normalised to [-1, 1]
//   Output: [N, 512] float32 → L2-normalised per row
class ArcFaceEmbedder final : public IFaceEmbedder {
public:
    ArcFaceEmbedder(const std::string& model_path,
                    OrtProvider provider, BackendConfig trt_cfg, int max_batch)
        : max_batch_(std::max(1, max_batch))
    {
        Ort::SessionOptions opts;
        opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        opts.SetIntraOpNumThreads(1);

        if (provider == OrtProvider::TensorRT) {
            if (trt_cfg.input_name.empty()) {
                Ort::SessionOptions probe_opts;
                probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
                Ort::Session probe(env_, model_path.c_str(), probe_opts);
                Ort::AllocatorWithDefaultOptions alloc;
                trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
            }
            if (trt_cfg.profile_min.empty()) {
                const std::string tail = "x3x112x112";
                trt_cfg.profile_min = "1" + tail;
                trt_cfg.profile_opt = std::to_string(max_batch_) + tail;
                trt_cfg.profile_max = std::to_string(max_batch_) + tail;
            }
        }

        apply_ort_model_cache(opts, model_path, trt_cfg);
        apply_ort_provider(opts, provider, "ArcFace", trt_cfg);

        session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);

        Ort::AllocatorWithDefaultOptions alloc;
        auto in_name  = session_->GetInputNameAllocated(0, alloc);
        auto out_name = session_->GetOutputNameAllocated(0, alloc);
        input_name_  = in_name.get();
        output_name_ = out_name.get();

        auto in_type  = session_->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
        auto out_type = session_->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetElementType();
        input_is_fp16_  = (in_type  == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
        output_is_fp16_ = (out_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);

        std::cerr << "[ArcFace] loaded: " << model_path << "\n";
    }

    int max_batch() const override { return max_batch_; }

    std::vector<Embedding> embed(const std::vector<cv::Mat>& crops) override {
        if (crops.empty()) return {};
        const int n = static_cast<int>(crops.size());

        std::vector<cv::Mat> rgbs(n);
        for (int i = 0; i < n; ++i)
            cv::cvtColor(crops[i], rgbs[i], cv::COLOR_BGR2RGB);

        cv::Mat blob = cv::dnn::blobFromImages(
            rgbs, 1.0 / 128.0, {112, 112},
            cv::Scalar(127.5, 127.5, 127.5),
            /*swapRB=*/false, /*crop=*/false, CV_32F);

        const std::array<int64_t, 4> in_shape = {n, 3, 112, 112};
        auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

        const char* in_name  = input_name_.c_str();
        const char* out_name = output_name_.c_str();

        cv::Mat blob16;
        if (input_is_fp16_) blob.convertTo(blob16, CV_16F);

        Ort::Value in_tensor = input_is_fp16_
            ? Ort::Value::CreateTensor<Ort::Float16_t>(
                mem, reinterpret_cast<Ort::Float16_t*>(blob16.ptr<uint16_t>()), blob16.total(),
                in_shape.data(), in_shape.size())
            : Ort::Value::CreateTensor<float>(
                mem, blob.ptr<float>(), blob.total(),
                in_shape.data(), in_shape.size());

        auto outs = session_->Run(Ort::RunOptions{nullptr}, &in_name, &in_tensor, 1, &out_name, 1);

        std::vector<Embedding> result(n);
        if (output_is_fp16_) {
            const auto* data16 = outs[0].GetTensorData<Ort::Float16_t>();
            std::vector<float> buf(n * 512);
            for (int j = 0; j < n * 512; ++j)
                buf[j] = data16[j].ToFloat();
            for (int i = 0; i < n; ++i)
                result[i] = l2_normalise(buf.data() + i * 512);
        } else {
            const float* data = outs[0].GetTensorData<float>();
            for (int i = 0; i < n; ++i)
                result[i] = l2_normalise(data + i * 512);
        }
        return result;
    }

private:
    int                             max_batch_;
    Ort::Env                        env_{ORT_LOGGING_LEVEL_ERROR, "arcface"};
    std::unique_ptr<Ort::Session>   session_;
    std::string                     input_name_;
    std::string                     output_name_;
    bool                            input_is_fp16_  = false;
    bool                            output_is_fp16_ = false;
};

// ── TransNetV2SceneDetector ───────────────────────────────────────────────────
// ONNX Runtime TransNetV2 shot-boundary detector.
//   Input  "input" : float32 [1, 100, 27, 48, 3]  RGB, channels-last, 0-255
//   Output "534"   : float32 [1, 100, 1]  single-frame boundary logits (used)
//   Output "535"   : float32 [1, 100, 1]  "many-hot" head (ignored)
// Returns per-frame sigmoid boundary probabilities. Fixed input shape, so the
// TRT-EP profile is pinned to the single 1×100×27×48×3 shape.
class TransNetV2SceneDetector final : public ISceneDetector {
public:
    TransNetV2SceneDetector(const std::string& model_path,
                            OrtProvider provider, BackendConfig trt_cfg)
    {
        Ort::SessionOptions opts;
        opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        opts.SetIntraOpNumThreads(1);

        if (provider == OrtProvider::TensorRT) {
            if (trt_cfg.input_name.empty()) {
                Ort::SessionOptions probe_opts;
                probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
                Ort::Session probe(env_, model_path.c_str(), probe_opts);
                Ort::AllocatorWithDefaultOptions alloc;
                trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
            }
            if (trt_cfg.profile_min.empty()) {
                const std::string shape =
                    "1x" + std::to_string(kWindow) + "x" +
                    std::to_string(kFrameH) + "x" + std::to_string(kFrameW) + "x3";
                trt_cfg.profile_min = shape;
                trt_cfg.profile_opt = shape;
                trt_cfg.profile_max = shape;
            }
        }

        apply_ort_model_cache(opts, model_path, trt_cfg);
        apply_ort_provider(opts, provider, "TransNetV2", trt_cfg);

        session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);

        Ort::AllocatorWithDefaultOptions alloc;
        input_name_ = session_->GetInputNameAllocated(0, alloc).get();

        // The primary boundary head is the first graph output ("534"); the
        // second ("535") is the many-hot auxiliary head we ignore. Bind both by
        // name so ORT returns them in a known order.
        const size_t n_out = session_->GetOutputCount();
        if (n_out < 1)
            throw std::runtime_error("[TransNetV2] model has no outputs");
        for (size_t i = 0; i < n_out; ++i)
            out_name_storage_.emplace_back(
                session_->GetOutputNameAllocated(i, alloc).get());
        for (auto& s : out_name_storage_) out_name_ptrs_.push_back(s.c_str());

        std::cerr << "[TransNetV2] loaded: " << model_path
                  << "  (window=" << kWindow << ")\n";
    }

    std::vector<float> detect_window(const std::vector<cv::Mat>& window) override {
        if (static_cast<int>(window.size()) != kWindow)
            throw std::runtime_error(
                "[TransNetV2] detect_window expects exactly " +
                std::to_string(kWindow) + " frames, got " +
                std::to_string(window.size()));

        // Pack into NHWC-per-frame contiguous buffer [100][27][48][3], RGB 0-255.
        std::vector<float> buf(
            static_cast<size_t>(kWindow) * kFrameH * kFrameW * 3);
        size_t o = 0;
        for (int f = 0; f < kWindow; ++f) {
            const cv::Mat& m = window[f];
            // Expect 48×27 BGR CV_8UC3; guard against mis-sized input.
            const cv::Mat* src = &m;
            cv::Mat resized;
            if (m.cols != kFrameW || m.rows != kFrameH || m.type() != CV_8UC3) {
                cv::Mat tmp;
                if (m.type() != CV_8UC3) m.convertTo(tmp, CV_8UC3); else tmp = m;
                cv::resize(tmp, resized, {kFrameW, kFrameH}, 0, 0, cv::INTER_AREA);
                src = &resized;
            }
            for (int y = 0; y < kFrameH; ++y) {
                const cv::Vec3b* row = src->ptr<cv::Vec3b>(y);
                for (int x = 0; x < kFrameW; ++x) {
                    const cv::Vec3b& px = row[x];   // BGR
                    buf[o++] = static_cast<float>(px[2]);  // R
                    buf[o++] = static_cast<float>(px[1]);  // G
                    buf[o++] = static_cast<float>(px[0]);  // B
                }
            }
        }

        const std::array<int64_t, 5> in_shape = {1, kWindow, kFrameH, kFrameW, 3};
        auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        auto in_tensor = Ort::Value::CreateTensor<float>(
            mem, buf.data(), buf.size(), in_shape.data(), in_shape.size());

        const char* in_name_c = input_name_.c_str();
        auto outs = session_->Run(
            Ort::RunOptions{nullptr},
            &in_name_c, &in_tensor, 1,
            out_name_ptrs_.data(), out_name_ptrs_.size());

        // Primary head: logits [1,100,1] → sigmoid.
        const float* logits = outs[0].GetTensorData<float>();
        std::vector<float> probs(kWindow);
        for (int i = 0; i < kWindow; ++i)
            probs[i] = 1.f / (1.f + std::exp(-logits[i]));
        return probs;
    }

private:
    Ort::Env                        env_{ORT_LOGGING_LEVEL_WARNING, "transnetv2"};
    std::unique_ptr<Ort::Session>   session_;
    std::string                     input_name_;
    std::vector<std::string>        out_name_storage_;
    std::vector<const char*>        out_name_ptrs_;
};

}  // namespace

// ── Factories ─────────────────────────────────────────────────────────────────

std::unique_ptr<IFaceDetector> make_face_detector(const Config& cfg) {
    if (!cfg.detector_engine.empty())
        throw std::runtime_error(
            "detector_engine set but this build uses the ORT inference backend; "
            "rebuild with -DSAE_INFERENCE_BACKEND=TRT to use raw TensorRT engines.");
    const OrtProvider provider = detect_ort_provider();
    std::cerr << "[face_detector] ORT backend, provider: "
              << provider_name(provider) << "\n";
    return std::make_unique<SCRFDDecoder>(
        cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
}

std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
    if (!cfg.arcface_engine.empty())
        throw std::runtime_error(
            "arcface_engine set but this build uses the ORT inference backend; "
            "rebuild with -DSAE_INFERENCE_BACKEND=TRT to use raw TensorRT engines.");
    const OrtProvider provider = detect_ort_provider();
    std::cerr << "[face_embedder] ORT backend, provider: "
              << provider_name(provider) << "\n";
    return std::make_unique<ArcFaceEmbedder>(
        cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
}

std::unique_ptr<ISceneDetector> make_scene_detector(const Config& cfg) {
    if (!cfg.scene_engine.empty())
        throw std::runtime_error(
            "scene_engine set but this build uses the ORT inference backend; "
            "rebuild with -DSAE_INFERENCE_BACKEND=TRT to use raw TensorRT engines.");
    const OrtProvider provider = detect_ort_provider();
    std::cerr << "[scene_detector] ORT backend, provider: "
              << provider_name(provider) << "\n";
    return std::make_unique<TransNetV2SceneDetector>(
        cfg.scene_model, provider, cfg.trt);
}
