// ── TensorRT inference backend ────────────────────────────────────────────────
// Pure-TensorRT 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=TRT.
//
// Loads serialised engines built by scripts/build_trt_engines.sh (or any
// trtexec-produced .engine matching the I/O contract). Skips ONNX Runtime
// entirely — useful where ORT was built without the TensorRT EP.
//
// This is the ONLY translation unit that includes NvInfer.h / cuda_runtime; the
// core application never sees them.

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

#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <cstdlib>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>

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

namespace {

struct CudaError : std::runtime_error {
    using std::runtime_error::runtime_error;
};

inline void check_cuda(cudaError_t e, const char* what) {
    if (e != cudaSuccess)
        throw CudaError(std::string(what) + ": " + cudaGetErrorString(e));
}

/// TRACES: VR-015 | PR-004
/// Select how a thread waits for the GPU. Must run before the CUDA context is
/// created, so every engine constructor calls it and the first one wins.
///
/// The default (`cudaDeviceScheduleAuto`) spin-waits: `cudaStreamSynchronize`
/// burns the calling thread's CPU for the whole of the device's work. Measured
/// here, the embedder thread sat at 99.7% *user* time with 0.5 s of system time
/// across 183 s — i.e. no blocking syscalls at all — while the GPU ran flat out.
///
/// On this laptop that is not merely wasted CPU. `nvidia-powerd` arbitrates one
/// power budget across CPU and GPU, and the GPU's ceiling was observed dropping
/// from 20 W idle to 15 W under our load, with the SM clock *falling* from
/// 1005 MHz to 210 MHz once work started. Spinning may therefore be buying
/// watts away from the device the pipeline is actually waiting on.
///
/// SAE_CUDA_BLOCKING_SYNC=1 switches to a blocking wait so the A/B needs no
/// rebuild. Default is unchanged until the measurement says otherwise.
inline void configure_cuda_sync_once() {
    static const bool done = [] {
        const char* env = std::getenv("SAE_CUDA_BLOCKING_SYNC");
        if (env && env[0] == '1') {
            cudaError_t e = cudaSetDeviceFlags(cudaDeviceScheduleBlockingSync);
            std::cerr << "[cuda] sync policy: BlockingSync"
                      << (e == cudaSuccess ? "" : " (FAILED — context already created)")
                      << "\n";
        } else {
            std::cerr << "[cuda] sync policy: default (spin) — "
                         "set SAE_CUDA_BLOCKING_SYNC=1 to compare\n";
        }
        return true;
    }();
    (void)done;
}

class TrtLogger : public nvinfer1::ILogger {
public:
    void log(Severity sev, const char* msg) noexcept override {
        if (sev <= Severity::kWARNING)
            std::cerr << "[TRT] " << msg << "\n";
    }
};
inline TrtLogger& logger() { static TrtLogger g; return g; }

struct TrtDeleter { template<class T> void operator()(T* p) const { delete p; } };

inline std::vector<char> read_file(const std::string& path, const char* who) {
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) throw std::runtime_error(std::string(who) + ": cannot open " + path);
    const std::streamsize sz = f.tellg();
    f.seekg(0);
    std::vector<char> blob(sz);
    f.read(blob.data(), sz);
    return blob;
}

// ── TrtArcFaceEmbedder ────────────────────────────────────────────────────────
// Engine I/O contract: input Nx3x112x112 float32/float16, output Nx512.
class TrtArcFaceEmbedder final : public IFaceEmbedder {
public:
    explicit TrtArcFaceEmbedder(const std::string& engine_path) {
        std::vector<char> blob = read_file(engine_path, "TrtArcFaceEmbedder");

        runtime_.reset(nvinfer1::createInferRuntime(logger()));
        if (!runtime_) throw std::runtime_error("createInferRuntime failed");
        engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
        if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
        context_.reset(engine_->createExecutionContext());
        if (!context_) throw std::runtime_error("createExecutionContext failed");

        const int n_io = engine_->getNbIOTensors();
        for (int i = 0; i < n_io; ++i) {
            const char* name = engine_->getIOTensorName(i);
            if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
                input_name_ = name;
            else
                output_name_ = name;
        }
        if (input_name_.empty() || output_name_.empty())
            throw std::runtime_error("TrtArcFaceEmbedder: engine missing input/output tensor");

        auto in_dtype  = engine_->getTensorDataType(input_name_.c_str());
        auto out_dtype = engine_->getTensorDataType(output_name_.c_str());
        input_is_fp16_  = (in_dtype  == nvinfer1::DataType::kHALF);
        output_is_fp16_ = (out_dtype == nvinfer1::DataType::kHALF);

        auto max_dims = engine_->getProfileShape(input_name_.c_str(), 0,
                                                 nvinfer1::OptProfileSelector::kMAX);
        if (max_dims.nbDims != 4 || max_dims.d[1] != 3 ||
            max_dims.d[2] != 112 || max_dims.d[3] != 112)
            throw std::runtime_error("TrtArcFaceEmbedder: unexpected input shape in engine");
        max_batch_ = max_dims.d[0];

        const std::size_t in_bytes  = static_cast<std::size_t>(max_batch_) * 3 * 112 * 112 *
                                      (input_is_fp16_ ? 2 : 4);
        const std::size_t out_bytes = static_cast<std::size_t>(max_batch_) * 512 *
                                      (output_is_fp16_ ? 2 : 4);
        check_cuda(cudaMalloc(&d_input_,  in_bytes),  "cudaMalloc input");
        check_cuda(cudaMalloc(&d_output_, out_bytes), "cudaMalloc output");
        configure_cuda_sync_once();
        check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");

        context_->setTensorAddress(input_name_.c_str(),  d_input_);
        context_->setTensorAddress(output_name_.c_str(), d_output_);

        std::cerr << "[TrtArcFace] loaded: " << engine_path
                  << "  max_batch=" << max_batch_
                  << (input_is_fp16_ ? "  fp16-in" : "")
                  << (output_is_fp16_ ? "  fp16-out" : "")
                  << "\n";
    }

    ~TrtArcFaceEmbedder() override {
        if (stream_)   cudaStreamDestroy(stream_);
        if (d_input_)  cudaFree(d_input_);
        if (d_output_) cudaFree(d_output_);
    }

    TrtArcFaceEmbedder(const TrtArcFaceEmbedder&) = delete;
    TrtArcFaceEmbedder& operator=(const TrtArcFaceEmbedder&) = delete;

    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());
        if (n > max_batch_)
            throw std::runtime_error("TrtArcFaceEmbedder: batch " + std::to_string(n) +
                                      " exceeds engine max " + std::to_string(max_batch_));

        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);

        std::lock_guard<std::mutex> lk(mu_);
        context_->setInputShape(input_name_.c_str(),
                                nvinfer1::Dims4{n, 3, 112, 112});

        const std::size_t in_count = static_cast<std::size_t>(n) * 3 * 112 * 112;
        if (input_is_fp16_) {
            cv::Mat blob16;
            blob.convertTo(blob16, CV_16F);
            check_cuda(cudaMemcpyAsync(d_input_, blob16.ptr(), in_count * 2,
                                       cudaMemcpyHostToDevice, stream_),
                       "H2D input fp16");
        } else {
            check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
                                       cudaMemcpyHostToDevice, stream_),
                       "H2D input fp32");
        }

        if (!context_->enqueueV3(stream_))
            throw std::runtime_error("TrtArcFaceEmbedder: enqueueV3 failed");

        const std::size_t out_count = static_cast<std::size_t>(n) * 512;
        std::vector<float> host_f32(out_count);
        if (output_is_fp16_) {
            std::vector<uint16_t> host_f16(out_count);
            check_cuda(cudaMemcpyAsync(host_f16.data(), d_output_, out_count * 2,
                                       cudaMemcpyDeviceToHost, stream_),
                       "D2H output fp16");
            check_cuda(cudaStreamSynchronize(stream_), "stream sync");
            cv::Mat src16(1, static_cast<int>(out_count), CV_16F, host_f16.data());
            cv::Mat dst32(1, static_cast<int>(out_count), CV_32F, host_f32.data());
            src16.convertTo(dst32, CV_32F);
        } else {
            check_cuda(cudaMemcpyAsync(host_f32.data(), d_output_, out_count * 4,
                                       cudaMemcpyDeviceToHost, stream_),
                       "D2H output fp32");
            check_cuda(cudaStreamSynchronize(stream_), "stream sync");
        }

        std::vector<Embedding> out(n);
        for (int i = 0; i < n; ++i)
            out[i] = l2_normalise(host_f32.data() + i * 512);
        return out;
    }

private:
    std::unique_ptr<nvinfer1::IRuntime,          TrtDeleter> runtime_;
    std::unique_ptr<nvinfer1::ICudaEngine,       TrtDeleter> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;

    std::string  input_name_;
    std::string  output_name_;
    bool         input_is_fp16_  = false;
    bool         output_is_fp16_ = false;
    int          max_batch_ = 1;

    void*        d_input_  = nullptr;
    void*        d_output_ = nullptr;
    cudaStream_t stream_   = nullptr;

    mutable std::mutex mu_;
};

// ── TrtScrfdDecoder ───────────────────────────────────────────────────────────
// Pure-TensorRT SCRFD detector (1x3x640x640 input pinned). Post-processing
// matches the ORT decoder byte-for-byte — only inference is swapped.
class TrtScrfdDecoder 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;

    TrtScrfdDecoder(const std::string& engine_path,
                    float conf_threshold, float nms_threshold)
        : conf_threshold_(conf_threshold)
        , nms_threshold_(nms_threshold)
    {
        std::vector<char> blob = read_file(engine_path, "TrtScrfdDecoder");

        runtime_.reset(nvinfer1::createInferRuntime(logger()));
        if (!runtime_) throw std::runtime_error("createInferRuntime failed");
        engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
        if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
        context_.reset(engine_->createExecutionContext());
        if (!context_) throw std::runtime_error("createExecutionContext failed");

        const int n_io = engine_->getNbIOTensors();
        for (int i = 0; i < n_io; ++i) {
            const char* name = engine_->getIOTensorName(i);
            if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) {
                if (!input_name_.empty())
                    throw std::runtime_error("TrtScrfdDecoder: multiple inputs not supported");
                input_name_ = name;
            } else {
                output_names_.emplace_back(name);
            }
        }
        if (input_name_.empty())
            throw std::runtime_error("TrtScrfdDecoder: no input tensor");
        const int n_out = static_cast<int>(output_names_.size());
        if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
            throw std::runtime_error(
                "TrtScrfdDecoder: expected 9 or 12 outputs (kps-variant SCRFD), got "
                + std::to_string(n_out));
        fmc_ = n_out / 3;

        auto in_dims = engine_->getProfileShape(input_name_.c_str(), 0,
                                                nvinfer1::OptProfileSelector::kOPT);
        if (in_dims.nbDims != 4 || in_dims.d[0] != 1 || in_dims.d[1] != 3 ||
            in_dims.d[2] != kInputH || in_dims.d[3] != kInputW)
            throw std::runtime_error(
                "TrtScrfdDecoder: engine input must be 1x3x" +
                std::to_string(kInputH) + "x" + std::to_string(kInputW));

        const std::size_t in_bytes = static_cast<std::size_t>(3) * kInputH * kInputW * 4;
        check_cuda(cudaMalloc(&d_input_, in_bytes), "cudaMalloc input");
        context_->setTensorAddress(input_name_.c_str(), d_input_);
        context_->setInputShape(input_name_.c_str(),
                                nvinfer1::Dims4{1, 3, kInputH, kInputW});

        d_outputs_.resize(n_out, nullptr);
        host_outputs_.resize(n_out);
        out_elem_counts_.resize(n_out, 0);

        const int expected_last[3] = {1, 4, 10};
        for (int oi = 0; oi < n_out; ++oi) {
            auto dims = context_->getTensorShape(output_names_[oi].c_str());
            if (dims.nbDims < 1)
                throw std::runtime_error("TrtScrfdDecoder: bad shape for output " +
                                          output_names_[oi]);
            std::size_t count = 1;
            for (int d = 0; d < dims.nbDims; ++d) count *= static_cast<std::size_t>(dims.d[d]);
            const int last = dims.d[dims.nbDims - 1];
            const int group = oi / fmc_;        // 0=scores, 1=bboxes, 2=kps
            if (last != expected_last[group])
                throw std::runtime_error(
                    "TrtScrfdDecoder: output '" + output_names_[oi] + "' last-dim is " +
                    std::to_string(last) + ", expected " + std::to_string(expected_last[group]) +
                    ". Engine does not match SCRFD-bnkps layout.");

            check_cuda(cudaMalloc(&d_outputs_[oi], count * 4), "cudaMalloc output");
            context_->setTensorAddress(output_names_[oi].c_str(), d_outputs_[oi]);
            host_outputs_[oi].resize(count);
            out_elem_counts_[oi] = count;
        }

        configure_cuda_sync_once();
        check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");

        std::cerr << "[TrtScrfd] loaded: " << engine_path
                  << "  fmc=" << fmc_ << "  outputs=" << n_out << "\n";
    }

    ~TrtScrfdDecoder() override {
        if (stream_) cudaStreamDestroy(stream_);
        if (d_input_) cudaFree(d_input_);
        for (void* p : d_outputs_) if (p) cudaFree(p);
    }

    TrtScrfdDecoder(const TrtScrfdDecoder&) = delete;
    TrtScrfdDecoder& operator=(const TrtScrfdDecoder&) = delete;

    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);

        std::lock_guard<std::mutex> lk(mu_);
        const std::size_t in_count = static_cast<std::size_t>(3) * kInputH * kInputW;
        check_cuda(cudaMemcpyAsync(d_input_, blob.ptr<float>(), in_count * 4,
                                   cudaMemcpyHostToDevice, stream_),
                   "H2D input");

        if (!context_->enqueueV3(stream_))
            throw std::runtime_error("TrtScrfdDecoder: enqueueV3 failed");

        for (std::size_t oi = 0; oi < d_outputs_.size(); ++oi) {
            check_cuda(cudaMemcpyAsync(host_outputs_[oi].data(), d_outputs_[oi],
                                        out_elem_counts_[oi] * 4,
                                        cudaMemcpyDeviceToHost, stream_),
                       "D2H output");
        }
        check_cuda(cudaStreamSynchronize(stream_), "stream sync");

        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 = host_outputs_[si].data();
            const float* b = host_outputs_[fmc_     + si].data();
            const float* k = host_outputs_[fmc_ * 2 + si].data();

            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:
    std::unique_ptr<nvinfer1::IRuntime,          TrtDeleter> runtime_;
    std::unique_ptr<nvinfer1::ICudaEngine,       TrtDeleter> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;

    float conf_threshold_;
    float nms_threshold_;
    int   fmc_{3};

    std::string                          input_name_;
    std::vector<std::string>             output_names_;
    void*                                d_input_  = nullptr;
    std::vector<void*>                   d_outputs_;
    mutable std::vector<std::vector<float>> host_outputs_;
    std::vector<std::size_t>             out_elem_counts_;

    cudaStream_t       stream_ = nullptr;
    mutable std::mutex mu_;
};

// ── TrtTransNetV2SceneDetector ────────────────────────────────────────────────
// Pure-TensorRT TransNetV2. Engine input pinned to 1×100×27×48×3 (RGB, 0-255);
// primary boundary head is the engine's first output tensor, sigmoided here to
// match the ORT path byte-for-byte.
class TrtTransNetV2SceneDetector final : public ISceneDetector {
public:
    explicit TrtTransNetV2SceneDetector(const std::string& engine_path) {
        std::vector<char> blob = read_file(engine_path, "TrtTransNetV2");

        runtime_.reset(nvinfer1::createInferRuntime(logger()));
        if (!runtime_) throw std::runtime_error("createInferRuntime failed");
        engine_.reset(runtime_->deserializeCudaEngine(blob.data(), blob.size()));
        if (!engine_) throw std::runtime_error("deserializeCudaEngine failed: " + engine_path);
        context_.reset(engine_->createExecutionContext());
        if (!context_) throw std::runtime_error("createExecutionContext failed");

        const int n_io = engine_->getNbIOTensors();
        for (int i = 0; i < n_io; ++i) {
            const char* name = engine_->getIOTensorName(i);
            if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT)
                input_name_ = name;
            else if (output_name_.empty())
                output_name_ = name;   // first output = primary boundary head
        }
        if (input_name_.empty() || output_name_.empty())
            throw std::runtime_error("TrtTransNetV2: engine missing input/output tensor");

        const std::size_t in_count =
            static_cast<std::size_t>(kWindow) * kFrameH * kFrameW * 3;
        const std::size_t out_count = static_cast<std::size_t>(kWindow);
        check_cuda(cudaMalloc(&d_input_,  in_count  * 4), "cudaMalloc input");
        check_cuda(cudaMalloc(&d_output_, out_count * 4), "cudaMalloc output");
        configure_cuda_sync_once();
        check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
        context_->setTensorAddress(input_name_.c_str(),  d_input_);
        context_->setTensorAddress(output_name_.c_str(), d_output_);

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

    ~TrtTransNetV2SceneDetector() override {
        if (stream_)   cudaStreamDestroy(stream_);
        if (d_input_)  cudaFree(d_input_);
        if (d_output_) cudaFree(d_output_);
    }

    TrtTransNetV2SceneDetector(const TrtTransNetV2SceneDetector&) = delete;
    TrtTransNetV2SceneDetector& operator=(const TrtTransNetV2SceneDetector&) = delete;

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

        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];
            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
                }
            }
        }

        std::lock_guard<std::mutex> lk(mu_);
        // TensorRT 10 removed the fixed-rank Dims5 helper (Dims2..Dims4 remain
        // in NvInferLegacyDims.h). Build the rank-5 shape via the generic Dims,
        // which works on both 8.x and 10.x.
        nvinfer1::Dims shape{};
        shape.nbDims = 5;
        shape.d[0] = 1;
        shape.d[1] = kWindow;
        shape.d[2] = kFrameH;
        shape.d[3] = kFrameW;
        shape.d[4] = 3;
        context_->setInputShape(input_name_.c_str(), shape);

        check_cuda(cudaMemcpyAsync(d_input_, buf.data(), buf.size() * 4,
                                   cudaMemcpyHostToDevice, stream_), "H2D input");
        if (!context_->enqueueV3(stream_))
            throw std::runtime_error("TrtTransNetV2: enqueueV3 failed");

        std::vector<float> logits(kWindow);
        check_cuda(cudaMemcpyAsync(logits.data(), d_output_, kWindow * 4,
                                   cudaMemcpyDeviceToHost, stream_), "D2H output");
        check_cuda(cudaStreamSynchronize(stream_), "stream sync");

        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:
    std::unique_ptr<nvinfer1::IRuntime,          TrtDeleter> runtime_;
    std::unique_ptr<nvinfer1::ICudaEngine,       TrtDeleter> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext, TrtDeleter> context_;
    std::string  input_name_;
    std::string  output_name_;
    void*        d_input_  = nullptr;
    void*        d_output_ = nullptr;
    cudaStream_t stream_   = nullptr;
    mutable std::mutex mu_;
};

}  // namespace

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

std::unique_ptr<IFaceDetector> make_face_detector(const Config& cfg) {
    if (cfg.detector_engine.empty())
        throw std::runtime_error(
            "TRT inference backend requires a pre-built detector engine "
            "(--detector-engine / cfg.detector_engine). Build one with "
            "scripts/build_trt_engines.sh, or rebuild with "
            "-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
    return std::make_unique<TrtScrfdDecoder>(
        cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
}

std::unique_ptr<IFaceEmbedder> make_face_embedder(const Config& cfg) {
    if (cfg.arcface_engine.empty())
        throw std::runtime_error(
            "TRT inference backend requires a pre-built ArcFace engine "
            "(--arcface-engine / cfg.arcface_engine). Build one with "
            "scripts/build_trt_engines.sh, or rebuild with "
            "-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
    return std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
}

std::unique_ptr<ISceneDetector> make_scene_detector(const Config& cfg) {
    if (cfg.scene_engine.empty())
        throw std::runtime_error(
            "TRT inference backend requires a pre-built TransNetV2 engine "
            "(--scene-detector-engine / cfg.scene_engine). Build one with "
            "scripts/convert_transnetv2.py --trt, or rebuild with "
            "-DSAE_INFERENCE_BACKEND=ORT to load the .onnx model directly.");
    return std::make_unique<TrtTransNetV2SceneDetector>(cfg.scene_engine);
}