#pragma once
#include "ort_provider.hpp"
#include "types.hpp"

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

#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

// ── SCRFDDecoder ──────────────────────────────────────────────────────────────
// ONNX Runtime-based decoder for InsightFace SCRFD face detector with kps.
// Uses ORT instead of 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 (same as YuNet/ArcFace convention):
//   [0] right-eye  [1] left-eye  [2] nose  [3] right-mouth  [4] left-mouth

struct SCRFDDecoder {
    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 = OrtProvider::CPU,
                 TrtConfig 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 profile to that single shape — otherwise
        // TRT picks generic shapes and either rebuilds per-call or falls
        // back to CUDA EP.
        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_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 and so
        // passes the count check above, but is encoded entirely differently).
        // Cross-check by output channel count: SCRFD's three groups of fmc_
        // outputs encode scores (1ch), bboxes (4ch) and 5-point kps (10ch).
        // YuNet exports loc/conf/iou with 14/2/1 channels, so this trips.
        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";
    }

    // Thread-safe: ORT Run() is safe for concurrent calls on the same Session.
    std::vector<DetectedFace> detect(const cv::Mat& img) const {
        // Letterbox to 640×640: uniform scale (preserves aspect ratio) + pad
        // shorter side with constant grey. Stretching to 640×640 (the prior
        // behaviour) distorts faces non-uniformly and degrades landmark
        // localisation — matters most for portrait gallery images and 16:9
        // video frames alike. Coordinates are mapped back via inverse scale +
        // pad-offset below.
        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)));

        // BGR→RGB swap + normalize to [-1,1] → NCHW float32 blob
        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);

                        // Decode in letterboxed network space, then un-pad +
                        // un-scale to original image coordinates.
                        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_;
};