#include <algorithm>
#include <atomic>
#include <array>
#include <cctype>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <memory>
#include <numeric>
#include <string>
#include <vector>

#include "ai_scheduler.h"
#include "face/face_result.h"
#include "node.h"
#include "utils/logger.h"

namespace rk3588 {

namespace {

inline int ClampInt(int v, int lo, int hi) {
    return v < lo ? lo : (v > hi ? hi : v);
}

struct Prior {
    float cx = 0.0f;
    float cy = 0.0f;
    float w = 0.0f;
    float h = 0.0f;
};

float IoU(const Rect& a, const Rect& b) {
    const float ax1 = a.x;
    const float ay1 = a.y;
    const float ax2 = a.x + a.w;
    const float ay2 = a.y + a.h;
    const float bx1 = b.x;
    const float by1 = b.y;
    const float bx2 = b.x + b.w;
    const float by2 = b.y + b.h;

    const float ix1 = std::max(ax1, bx1);
    const float iy1 = std::max(ay1, by1);
    const float ix2 = std::min(ax2, bx2);
    const float iy2 = std::min(ay2, by2);

    const float iw = std::max(0.0f, ix2 - ix1);
    const float ih = std::max(0.0f, iy2 - iy1);
    const float inter = iw * ih;
    const float ua = a.w * a.h + b.w * b.h - inter;
    return ua <= 0.0f ? 0.0f : (inter / ua);
}

void NmsSorted(const std::vector<Rect>& boxes, const std::vector<float>& scores,
              float nms_thresh, std::vector<int>& keep) {
    keep.clear();
    std::vector<int> order(scores.size());
    std::iota(order.begin(), order.end(), 0);
    std::sort(order.begin(), order.end(), [&](int a, int b) { return scores[a] > scores[b]; });

    for (int idx : order) {
        bool suppressed = false;
        for (int kept : keep) {
            if (IoU(boxes[idx], boxes[kept]) > nms_thresh) {
                suppressed = true;
                break;
            }
        }
        if (!suppressed) keep.push_back(idx);
    }
}

inline float Sigmoid(float x) {
    return 1.0f / (1.0f + std::exp(-x));
}

inline float Softmax2(float a, float b) {
    const float m = std::max(a, b);
    const float ea = std::exp(a - m);
    const float eb = std::exp(b - m);
    return eb / (ea + eb);
}

inline float HalfToFloat(uint16_t h) {
    const uint32_t sign = (static_cast<uint32_t>(h & 0x8000u)) << 16;
    uint32_t exp = (h & 0x7C00u) >> 10;
    uint32_t mant = (h & 0x03FFu);

    uint32_t f = 0;
    if (exp == 0) {
        if (mant == 0) {
            f = sign;
        } else {
            // Subnormal
            exp = 1;
            while ((mant & 0x0400u) == 0) {
                mant <<= 1;
                --exp;
            }
            mant &= 0x03FFu;
            exp = exp + (127 - 15);
            f = sign | (exp << 23) | (mant << 13);
        }
    } else if (exp == 31) {
        // Inf/NaN
        f = sign | 0x7F800000u | (mant << 13);
    } else {
        exp = exp + (127 - 15);
        f = sign | (exp << 23) | (mant << 13);
    }

    float out;
    memcpy(&out, &f, sizeof(out));
    return out;
}

template <typename T>
inline float Dequant(T q, int32_t zp, float scale) {
    return (static_cast<float>(q) - static_cast<float>(zp)) * scale;
}

struct Tensor {
    const uint8_t* data = nullptr;
    size_t size = 0;
    int32_t zp = 0;
    float scale = 1.0f;
    std::vector<uint32_t> dims;
#if defined(RK3588_ENABLE_RKNN)
    rknn_tensor_type type = RKNN_TENSOR_UINT8;
#endif
};

struct NcTensor {
    int n = 0;
    int c = 0;
    std::vector<float> data;  // N*C row-major
};

bool ExtractNc(const Tensor& t, int c, NcTensor& out) {
    out = {};
    out.c = c;
    if (!t.data || t.size == 0) return false;

    size_t elem_size = 1;
    bool is_float32 = false;
    bool is_float16 = false;
#if defined(RK3588_ENABLE_RKNN)
    if (t.type == RKNN_TENSOR_FLOAT16) {
        elem_size = 2;
        is_float16 = true;
    }
    if (t.type == RKNN_TENSOR_FLOAT32) {
        elem_size = 4;
        is_float32 = true;
    }
#endif
    const size_t elem_cnt = elem_size > 0 ? (t.size / elem_size) : 0;
    if (elem_cnt == 0) return false;

    int n = 0;
    enum class Layout { FlatNc, CxN, NCHW, NHWC } layout = Layout::FlatNc;
    int dN = 1, dH = 1, dW = 1;
    const bool tried_shape = (t.dims.size() == 2 || t.dims.size() == 3 || t.dims.size() == 4);

    if (t.dims.size() == 3) {
        // Common: [1, C, N] or [1, N, C]
        const uint32_t d1 = t.dims[1];
        const uint32_t d2 = t.dims[2];
        if (static_cast<int>(d1) == c) {
            n = static_cast<int>(d2);
            layout = Layout::CxN;  // [1, C, N]
        } else if (static_cast<int>(d2) == c) {
            n = static_cast<int>(d1);
            layout = Layout::FlatNc;  // treat as contiguous [1, N, C]
        }
    } else if (t.dims.size() == 2) {
        // [N, C] or [C, N]
        const uint32_t d0 = t.dims[0];
        const uint32_t d1 = t.dims[1];
        if (static_cast<int>(d1) == c) {
            n = static_cast<int>(d0);
            layout = Layout::FlatNc;
        } else if (static_cast<int>(d0) == c) {
            n = static_cast<int>(d1);
            layout = Layout::CxN;
        }
    } else if (t.dims.size() == 4) {
        // Common: [N, C, H, W] (NCHW) or [N, H, W, C] (NHWC)
        const int dn = static_cast<int>(t.dims[0]);
        const int d1 = static_cast<int>(t.dims[1]);
        const int d2 = static_cast<int>(t.dims[2]);
        const int d3 = static_cast<int>(t.dims[3]);
        if (d3 == c) {
            // NHWC
            layout = Layout::NHWC;
            dN = std::max(1, dn);
            dH = std::max(1, d1);
            dW = std::max(1, d2);
            n = dN * dH * dW;
        } else if (d1 == c) {
            // NCHW
            layout = Layout::NCHW;
            dN = std::max(1, dn);
            dH = std::max(1, d2);
            dW = std::max(1, d3);
            n = dN * dH * dW;
        }
    }

    // IMPORTANT: when dims are present but don't match the requested channel count,
    // do NOT fallback to element-count heuristics (it will mis-classify tensors).
    if (tried_shape && n <= 0) {
        return false;
    }

    if (n <= 0) {
        if (elem_cnt % static_cast<size_t>(c) != 0) return false;
        n = static_cast<int>(elem_cnt / static_cast<size_t>(c));
        layout = Layout::FlatNc;
    }

    if (static_cast<size_t>(n) * static_cast<size_t>(c) != elem_cnt) {
        return false;
    }

    out.n = n;
    out.data.resize(static_cast<size_t>(n) * static_cast<size_t>(c));

    auto ReadElem = [&](size_t idx) -> float {
        if (is_float32) {
            const float* fp = reinterpret_cast<const float*>(t.data);
            return fp[idx];
        }
#if defined(RK3588_ENABLE_RKNN)
        if (is_float16) {
            const uint16_t* hp = reinterpret_cast<const uint16_t*>(t.data);
            return HalfToFloat(hp[idx]);
        }
        if (t.type == RKNN_TENSOR_INT8) {
            const int8_t* p = reinterpret_cast<const int8_t*>(t.data);
            return Dequant(p[idx], t.zp, t.scale);
        }
#endif
        const uint8_t* p = reinterpret_cast<const uint8_t*>(t.data);
        return Dequant(p[idx], t.zp, t.scale);
    };

    if (layout == Layout::FlatNc) {
        for (size_t i = 0; i < out.data.size(); ++i) out.data[i] = ReadElem(i);
    } else if (layout == Layout::CxN) {
        // Input is [C, N] contiguous. Transpose to [N, C].
        for (int ci = 0; ci < c; ++ci) {
            for (int ni = 0; ni < n; ++ni) {
                const size_t src_idx = static_cast<size_t>(ci) * static_cast<size_t>(n) + static_cast<size_t>(ni);
                const size_t dst_idx = static_cast<size_t>(ni) * static_cast<size_t>(c) + static_cast<size_t>(ci);
                out.data[dst_idx] = ReadElem(src_idx);
            }
        }
    } else if (layout == Layout::NHWC) {
        // [N, H, W, C] contiguous; pack to [N*H*W, C]
        size_t dst = 0;
        for (int n0 = 0; n0 < dN; ++n0) {
            for (int y = 0; y < dH; ++y) {
                for (int x = 0; x < dW; ++x) {
                    const size_t base = ((static_cast<size_t>(n0) * static_cast<size_t>(dH) + static_cast<size_t>(y)) *
                                             static_cast<size_t>(dW) + static_cast<size_t>(x)) *
                                        static_cast<size_t>(c);
                    for (int ci = 0; ci < c; ++ci) {
                        out.data[dst++] = ReadElem(base + static_cast<size_t>(ci));
                    }
                }
            }
        }
    } else if (layout == Layout::NCHW) {
        // [N, C, H, W] channel-first; interleave channels into [N*H*W, C]
        size_t dst = 0;
        for (int n0 = 0; n0 < dN; ++n0) {
            for (int y = 0; y < dH; ++y) {
                for (int x = 0; x < dW; ++x) {
                    for (int ci = 0; ci < c; ++ci) {
                        const size_t src_idx =
                            ((static_cast<size_t>(n0) * static_cast<size_t>(c) + static_cast<size_t>(ci)) *
                                 static_cast<size_t>(dH) + static_cast<size_t>(y)) *
                                static_cast<size_t>(dW) + static_cast<size_t>(x);
                        out.data[dst++] = ReadElem(src_idx);
                    }
                }
            }
        }
    }

    return true;
}

std::vector<Prior> GenerateRetinaFacePriors(int in_w, int in_h,
                                           const std::vector<int>& steps,
                                           const std::vector<std::vector<int>>& min_sizes) {
    std::vector<Prior> priors;
    if (steps.empty() || steps.size() != min_sizes.size()) return priors;
    priors.reserve(5000);

    for (size_t s = 0; s < steps.size(); ++s) {
        const int step = steps[s];
        const int fm_w = in_w / step;
        const int fm_h = in_h / step;
        for (int i = 0; i < fm_h; ++i) {
            for (int j = 0; j < fm_w; ++j) {
                for (int ms : min_sizes[s]) {
                    const float s_kx = static_cast<float>(ms) / static_cast<float>(in_w);
                    const float s_ky = static_cast<float>(ms) / static_cast<float>(in_h);
                    const float cx = (static_cast<float>(j) + 0.5f) * static_cast<float>(step) / static_cast<float>(in_w);
                    const float cy = (static_cast<float>(i) + 0.5f) * static_cast<float>(step) / static_cast<float>(in_h);
                    priors.push_back(Prior{cx, cy, s_kx, s_ky});
                }
            }
        }
    }
    return priors;
}

void ResizeRgbBilinear(const uint8_t* src, int src_w, int src_h, int src_stride,
                       uint8_t* dst, int dst_w, int dst_h, bool swap_rb) {
    const float scale_x = static_cast<float>(src_w) / static_cast<float>(dst_w);
    const float scale_y = static_cast<float>(src_h) / static_cast<float>(dst_h);

    for (int y = 0; y < dst_h; ++y) {
        const float fy = (static_cast<float>(y) + 0.5f) * scale_y - 0.5f;
        int y0 = static_cast<int>(std::floor(fy));
        int y1 = y0 + 1;
        const float wy1 = fy - static_cast<float>(y0);
        const float wy0 = 1.0f - wy1;
        y0 = ClampInt(y0, 0, src_h - 1);
        y1 = ClampInt(y1, 0, src_h - 1);

        const uint8_t* row0 = src + static_cast<size_t>(y0) * static_cast<size_t>(src_stride);
        const uint8_t* row1 = src + static_cast<size_t>(y1) * static_cast<size_t>(src_stride);
        uint8_t* out = dst + static_cast<size_t>(y) * static_cast<size_t>(dst_w) * 3;

        for (int x = 0; x < dst_w; ++x) {
            const float fx = (static_cast<float>(x) + 0.5f) * scale_x - 0.5f;
            int x0 = static_cast<int>(std::floor(fx));
            int x1 = x0 + 1;
            const float wx1 = fx - static_cast<float>(x0);
            const float wx0 = 1.0f - wx1;
            x0 = ClampInt(x0, 0, src_w - 1);
            x1 = ClampInt(x1, 0, src_w - 1);

            const uint8_t* p00 = row0 + x0 * 3;
            const uint8_t* p01 = row0 + x1 * 3;
            const uint8_t* p10 = row1 + x0 * 3;
            const uint8_t* p11 = row1 + x1 * 3;

            for (int c = 0; c < 3; ++c) {
                const float v =
                    (static_cast<float>(p00[c]) * wx0 + static_cast<float>(p01[c]) * wx1) * wy0 +
                    (static_cast<float>(p10[c]) * wx0 + static_cast<float>(p11[c]) * wx1) * wy1;
                out[c] = static_cast<uint8_t>(ClampInt(static_cast<int>(v + 0.5f), 0, 255));
            }

            if (swap_rb) {
                std::swap(out[0], out[2]);
            }
            out += 3;
        }
    }
}

struct LetterboxInfo {
    float scale = 1.0f;
    int pad_x = 0;
    int pad_y = 0;
};

void LetterboxRgb(const uint8_t* src, int src_w, int src_h, int src_stride,
                  uint8_t* dst, int dst_w, int dst_h, bool swap_rb,
                  uint8_t pad_value, LetterboxInfo& info) {
    info = {};
    if (!src || !dst || src_w <= 0 || src_h <= 0 || dst_w <= 0 || dst_h <= 0 || src_stride <= 0) return;

    const float s = std::min(static_cast<float>(dst_w) / static_cast<float>(src_w),
                             static_cast<float>(dst_h) / static_cast<float>(src_h));
    const int rw = std::max(1, static_cast<int>(std::round(static_cast<float>(src_w) * s)));
    const int rh = std::max(1, static_cast<int>(std::round(static_cast<float>(src_h) * s)));
    const int px = (dst_w - rw) / 2;
    const int py = (dst_h - rh) / 2;

    info.scale = s;
    info.pad_x = px;
    info.pad_y = py;

    // fill background
    std::fill(dst, dst + static_cast<size_t>(dst_w) * static_cast<size_t>(dst_h) * 3, pad_value);

    std::vector<uint8_t> tmp(static_cast<size_t>(rw) * static_cast<size_t>(rh) * 3);
    ResizeRgbBilinear(src, src_w, src_h, src_stride, tmp.data(), rw, rh, swap_rb);

    for (int y = 0; y < rh; ++y) {
        uint8_t* drow = dst + (static_cast<size_t>(py + y) * static_cast<size_t>(dst_w) + static_cast<size_t>(px)) * 3;
        const uint8_t* srow = tmp.data() + static_cast<size_t>(y) * static_cast<size_t>(rw) * 3;
        memcpy(drow, srow, static_cast<size_t>(rw) * 3);
    }
}

}  // namespace

class AiFaceDetNode : public INode {
public:
    std::string Id() const override { return id_; }
    std::string Type() const override { return "ai_face_det"; }

    bool Init(const SimpleJson& config, const NodeContext& ctx) override {
        id_ = config.ValueOr<std::string>("id", "face_det");
        model_path_ = config.ValueOr<std::string>("model_path", "");
        conf_thresh_ = config.ValueOr<float>("conf", 0.6f);
        nms_thresh_ = config.ValueOr<float>("nms", 0.4f);
        max_faces_ = std::max(1, config.ValueOr<int>("max_faces", 10));
        output_landmarks_ = config.ValueOr<bool>("output_landmarks", true);
        letterbox_ = config.ValueOr<bool>("letterbox", letterbox_);
        face_class_index_ = config.ValueOr<int>("face_class_index", face_class_index_);
        if (face_class_index_ != 0 && face_class_index_ != 1) {
            face_class_index_ = -1;  // auto
        }

        const std::string fmt = config.ValueOr<std::string>("input_format", "rgb");
        input_format_ = fmt;
        for (auto& c : input_format_) c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));

        // RetinaFace priors defaults for 320 input (MobileNet0.25).
        steps_ = {8, 16, 32};
        min_sizes_ = {{16, 32}, {64, 128}, {256, 512}};
        if (const SimpleJson* pri = config.Find("prior"); pri && pri->IsObject()) {
            if (const SimpleJson* steps = pri->Find("steps"); steps && steps->IsArray()) {
                steps_.clear();
                for (const auto& v : steps->AsArray()) {
                    steps_.push_back(std::max(1, v.AsInt(1)));
                }
            }
            if (const SimpleJson* mins = pri->Find("min_sizes"); mins && mins->IsArray()) {
                min_sizes_.clear();
                for (const auto& grp : mins->AsArray()) {
                    std::vector<int> g;
                    for (const auto& v : grp.AsArray()) {
                        g.push_back(std::max(1, v.AsInt(1)));
                    }
                    if (!g.empty()) min_sizes_.push_back(std::move(g));
                }
            }
        }

        input_queue_ = ctx.input_queue;
        output_queues_ = ctx.output_queues;
        if (!input_queue_) {
            std::cerr << "[ai_face_det] no input queue for node " << id_ << "\n";
            return false;
        }
        if (output_queues_.empty()) {
            std::cerr << "[ai_face_det] no output queue for node " << id_ << "\n";
            return false;
        }

#if defined(RK3588_ENABLE_RKNN)
        if (model_path_.empty()) {
            std::cerr << "[ai_face_det] model_path is required\n";
            return false;
        }
        std::string err;
        model_handle_ = AiScheduler::Instance().LoadModel(model_path_, err);
        if (model_handle_ == kInvalidModelHandle) {
            std::cerr << "[ai_face_det] failed to load model: " << err << "\n";
            return false;
        }
        ModelInfo info;
        if (AiScheduler::Instance().GetModelInfo(model_handle_, info)) {
            model_w_ = info.input_width;
            model_h_ = info.input_height;
            n_output_ = info.n_output;
        }
        LogInfo("[ai_face_det] model loaded: " + model_path_ +
                " (" + std::to_string(model_w_) + "x" + std::to_string(model_h_) +
                ", outputs=" + std::to_string(n_output_) + ")");
#else
        LogWarn("[ai_face_det] RKNN disabled, will passthrough frames");
#endif
        return true;
    }

    bool Start() override {
        LogInfo("[ai_face_det] start id=" + id_ + " conf=" + std::to_string(conf_thresh_) +
                " nms=" + std::to_string(nms_thresh_) + " max_faces=" + std::to_string(max_faces_));
        return true;
    }

    bool UpdateConfig(const SimpleJson& new_config) override {
        const std::string new_id = new_config.ValueOr<std::string>("id", id_);
        if (!new_id.empty() && new_id != id_) return false;

        const std::string new_model = new_config.ValueOr<std::string>("model_path", model_path_);
        if (new_model != model_path_) {
            // Changing model requires graph rebuild.
            return false;
        }

        conf_thresh_ = new_config.ValueOr<float>("conf", conf_thresh_);
        nms_thresh_ = new_config.ValueOr<float>("nms", nms_thresh_);
        max_faces_ = std::max(1, new_config.ValueOr<int>("max_faces", max_faces_));
        output_landmarks_ = new_config.ValueOr<bool>("output_landmarks", output_landmarks_);

        if (const SimpleJson* pri = new_config.Find("prior"); pri && pri->IsObject()) {
            if (const SimpleJson* steps = pri->Find("steps"); steps && steps->IsArray()) {
                std::vector<int> new_steps;
                for (const auto& v : steps->AsArray()) new_steps.push_back(std::max(1, v.AsInt(1)));
                if (!new_steps.empty()) steps_ = std::move(new_steps);
            }
            if (const SimpleJson* mins = pri->Find("min_sizes"); mins && mins->IsArray()) {
                std::vector<std::vector<int>> new_mins;
                for (const auto& grp : mins->AsArray()) {
                    std::vector<int> g;
                    for (const auto& v : grp.AsArray()) g.push_back(std::max(1, v.AsInt(1)));
                    if (!g.empty()) new_mins.push_back(std::move(g));
                }
                if (!new_mins.empty()) min_sizes_ = std::move(new_mins);
            }
        }

        return true;
    }

    void Stop() override {
#if defined(RK3588_ENABLE_RKNN)
        if (model_handle_ != kInvalidModelHandle) {
            AiScheduler::Instance().UnloadModel(model_handle_);
            model_handle_ = kInvalidModelHandle;
        }
#endif
        LogInfo("[ai_face_det] stop id=" + id_);
    }

    NodeStatus Process(FramePtr frame) override {
        if (!frame) return NodeStatus::DROP;

#if defined(RK3588_ENABLE_RKNN)
        Run(frame);
#endif
        Push(frame);
        return NodeStatus::OK;
    }

private:
    void Push(FramePtr frame) {
        for (auto& q : output_queues_) q->Push(frame);
    }

#if defined(RK3588_ENABLE_RKNN)
    void Run(FramePtr frame) {
        if (!frame->data || frame->data_size == 0) return;
        if (frame->format != PixelFormat::RGB && frame->format != PixelFormat::BGR) {
            std::cerr << "[ai_face_det] input must be RGB/BGR\n";
            return;
        }

        const int src_w = frame->width;
        const int src_h = frame->height;
        const size_t src_row = static_cast<size_t>(src_w) * 3;
        const uint8_t* src = frame->planes[0].data ? frame->planes[0].data : frame->data;
        const int src_stride = frame->planes[0].stride > 0 ? frame->planes[0].stride
                              : (frame->stride > 0 ? frame->stride : static_cast<int>(src_row));
        if (!src || src_stride <= 0) return;

        const bool need_swap = (frame->format == PixelFormat::BGR && input_format_ == "rgb") ||
                               (frame->format == PixelFormat::RGB && input_format_ == "bgr");

        const int in_w = model_w_ > 0 ? model_w_ : src_w;
        const int in_h = model_h_ > 0 ? model_h_ : src_h;
        const size_t in_size = static_cast<size_t>(in_w) * static_cast<size_t>(in_h) * 3;

        const uint8_t* input_ptr = nullptr;

        LetterboxInfo lb;
        if (!letterbox_) {
            // Fast path: already packed, correct size, no channel swap.
            if (!need_swap && src_w == in_w && src_h == in_h &&
                static_cast<size_t>(src_stride) == src_row && frame->data_size >= src_row * static_cast<size_t>(src_h)) {
                input_ptr = src;
            } else {
                input_buf_.resize(in_size);
                if (src_w == in_w && src_h == in_h && static_cast<size_t>(src_stride) == src_row) {
                    memcpy(input_buf_.data(), src, in_size);
                    if (need_swap) {
                        for (size_t i = 0; i < in_size; i += 3) {
                            std::swap(input_buf_[i], input_buf_[i + 2]);
                        }
                    }
                } else {
                    ResizeRgbBilinear(src, src_w, src_h, src_stride, input_buf_.data(), in_w, in_h, need_swap);
                }
                input_ptr = input_buf_.data();
            }
        } else {
            input_buf_.resize(in_size);
            LetterboxRgb(src, src_w, src_h, src_stride, input_buf_.data(), in_w, in_h,
                         need_swap, 0, lb);
            input_ptr = input_buf_.data();
        }

        InferInput input;
        input.data = input_ptr;
        input.size = in_size;
        input.width = in_w;
        input.height = in_h;
        input.is_nhwc = true;

        auto r = AiScheduler::Instance().InferBorrowed(model_handle_, input);
        if (!r.success) {
            std::cerr << "[ai_face_det] inference failed: " << r.error << "\n";
            return;
        }

        if (!logged_io_.exchange(true)) {
            std::string s = "[ai_face_det] outputs:";
            for (const auto& o : r.outputs) {
                s += " idx=" + std::to_string(o.index) + " size=" + std::to_string(o.size);
#if defined(RK3588_ENABLE_RKNN)
                s += " type=" + std::to_string(static_cast<int>(o.type));
                s += " zp=" + std::to_string(o.zp) + " scale=" + std::to_string(o.scale);
                s += " dims=[";
                for (size_t i = 0; i < o.dims.size(); ++i) {
                    s += std::to_string(o.dims[i]);
                    if (i + 1 < o.dims.size()) s += ",";
                }
                s += "]";
#endif
            }
            LogInfo(s);
        }

        std::vector<Tensor> tensors;
        tensors.reserve(r.outputs.size());
        for (const auto& o : r.outputs) {
            Tensor t;
            t.data = o.data;
            t.size = o.size;
            t.zp = o.zp;
            t.scale = o.scale;
            t.dims = o.dims;
            t.type = o.type;
            tensors.push_back(std::move(t));
        }

        FaceDetResult det;
        det.img_w = src_w;
        det.img_h = src_h;
        det.model_name = "retinaface";

        DecodeRetinaFace(tensors, src_w, src_h, in_w, in_h, lb, det);
        if (!logged_decode_.exchange(true)) {
            LogInfo("[ai_face_det] decoded faces=" + std::to_string(det.faces.size()));
        }
        frame->face_det = std::make_shared<FaceDetResult>(std::move(det));
    }

    void DecodeRetinaFace(const std::vector<Tensor>& outs,
                          int orig_w, int orig_h,
                          int in_w, int in_h,
                          const LetterboxInfo& lb,
                          FaceDetResult& out) {
        // Find loc/conf/landms tensors.
        std::vector<NcTensor> locs;
        std::vector<NcTensor> confs;
        std::vector<NcTensor> landms;
        locs.reserve(4);
        confs.reserve(4);
        landms.reserve(4);

        for (const auto& t : outs) {
            NcTensor tmp;
            if (ExtractNc(t, 4, tmp)) {
                locs.push_back(std::move(tmp));
                continue;
            }
            if (ExtractNc(t, 2, tmp)) {
                confs.push_back(std::move(tmp));
                continue;
            }
            if (ExtractNc(t, 10, tmp)) {
                landms.push_back(std::move(tmp));
                continue;
            }
        }
        if (!logged_decode_detail_.exchange(true)) {
            LogInfo("[ai_face_det] decode tensors: loc_parts=" + std::to_string(locs.size()) +
                    " conf_parts=" + std::to_string(confs.size()) +
                    " lmk_parts=" + std::to_string(landms.size()));
        }
        if (locs.empty() || confs.empty()) return;

        // Concatenate along N.
        auto Concat = [](const std::vector<NcTensor>& parts) -> NcTensor {
            NcTensor all;
            if (parts.empty()) return all;
            all.c = parts[0].c;
            int total_n = 0;
            for (const auto& p : parts) total_n += p.n;
            all.n = total_n;
            all.data.resize(static_cast<size_t>(all.n) * static_cast<size_t>(all.c));
            size_t off = 0;
            for (const auto& p : parts) {
                if (p.c != all.c) continue;
                memcpy(all.data.data() + off, p.data.data(), p.data.size() * sizeof(float));
                off += p.data.size();
            }
            return all;
        };

        NcTensor loc = Concat(locs);
        NcTensor conf = Concat(confs);
        NcTensor lmk;
        if (output_landmarks_ && !landms.empty()) lmk = Concat(landms);

        if (loc.n <= 0 || conf.n != loc.n) return;
        const int n = loc.n;

        // Determine which class index represents "face".
        // Most RetinaFace uses [bg, face], i.e., face_idx=1, but some exports swap it.
        int face_idx = face_class_index_;
        float max_p0 = 0.0f;
        float max_p1 = 0.0f;
        double sum_p0 = 0.0;
        double sum_p1 = 0.0;
        if (face_idx < 0) {
            for (int i = 0; i < n; ++i) {
                const float s0 = conf.data[static_cast<size_t>(i) * 2 + 0];
                const float s1 = conf.data[static_cast<size_t>(i) * 2 + 1];
                float p1;
                if (s0 >= 0.0f && s0 <= 1.0f && s1 >= 0.0f && s1 <= 1.0f && std::fabs((s0 + s1) - 1.0f) < 0.1f) {
                    p1 = s1;
                } else {
                    p1 = Softmax2(s0, s1);
                }
                const float p0 = 1.0f - p1;
                if (p0 > max_p0) max_p0 = p0;
                if (p1 > max_p1) max_p1 = p1;
                sum_p0 += p0;
                sum_p1 += p1;
            }
            const double mean_p0 = sum_p0 / static_cast<double>(n);
            const double mean_p1 = sum_p1 / static_cast<double>(n);
            // Face channel is typically sparse (low mean). Background channel has high mean.
            face_idx = (mean_p1 <= mean_p0) ? 1 : 0;
        }

        if (!logged_conf_probe_.exchange(true)) {
            LogInfo("[ai_face_det] conf probe: max_p0=" + std::to_string(max_p0) +
                    " max_p1=" + std::to_string(max_p1) +
                    " face_idx=" + std::to_string(face_idx) +
                    (face_class_index_ < 0 ? " (auto)" : " (cfg)"));
        }

        const std::vector<Prior> priors = GenerateRetinaFacePriors(in_w, in_h, steps_, min_sizes_);
        if (!priors.empty() && static_cast<int>(priors.size()) != n) {
            // Mismatch: can't reliably decode.
            std::cerr << "[ai_face_det] prior mismatch: priors=" << priors.size() << " n=" << n << "\n";
            return;
        }

        const bool use_lb = letterbox_ && lb.scale > 0.0f;
        const float inv_s = use_lb ? (1.0f / lb.scale) : 1.0f;

        std::vector<Rect> boxes;
        std::vector<float> scores;
        std::vector<std::array<Point2f, 5>> lmks;
        boxes.reserve(static_cast<size_t>(n));
        scores.reserve(static_cast<size_t>(n));
        if (output_landmarks_) lmks.reserve(static_cast<size_t>(n));

        constexpr float var0 = 0.1f;
        constexpr float var1 = 0.2f;

        float max_face_prob = 0.0f;
        for (int i = 0; i < n; ++i) {
            const float s0 = conf.data[static_cast<size_t>(i) * 2 + 0];
            const float s1 = conf.data[static_cast<size_t>(i) * 2 + 1];
            float p1;
            if (s0 >= 0.0f && s0 <= 1.0f && s1 >= 0.0f && s1 <= 1.0f && std::fabs((s0 + s1) - 1.0f) < 0.1f) {
                p1 = s1;
            } else {
                p1 = Softmax2(s0, s1);
            }
            const float score = (face_idx == 1) ? p1 : (1.0f - p1);
            if (score > max_face_prob) max_face_prob = score;
            if (score < conf_thresh_) continue;

            const Prior p = priors.empty() ? Prior{0, 0, 0, 0} : priors[static_cast<size_t>(i)];

            const float dx = loc.data[static_cast<size_t>(i) * 4 + 0];
            const float dy = loc.data[static_cast<size_t>(i) * 4 + 1];
            const float dw = loc.data[static_cast<size_t>(i) * 4 + 2];
            const float dh = loc.data[static_cast<size_t>(i) * 4 + 3];

            const float cx = p.cx + dx * var0 * p.w;
            const float cy = p.cy + dy * var0 * p.h;
            const float ww = p.w * std::exp(dw * var1);
            const float hh = p.h * std::exp(dh * var1);

            float x1 = (cx - ww * 0.5f) * static_cast<float>(in_w);
            float y1 = (cy - hh * 0.5f) * static_cast<float>(in_h);
            float x2 = (cx + ww * 0.5f) * static_cast<float>(in_w);
            float y2 = (cy + hh * 0.5f) * static_cast<float>(in_h);

            if (use_lb) {
                x1 = (x1 - static_cast<float>(lb.pad_x)) * inv_s;
                x2 = (x2 - static_cast<float>(lb.pad_x)) * inv_s;
                y1 = (y1 - static_cast<float>(lb.pad_y)) * inv_s;
                y2 = (y2 - static_cast<float>(lb.pad_y)) * inv_s;
            } else {
                const float sx = static_cast<float>(orig_w) / static_cast<float>(in_w);
                const float sy = static_cast<float>(orig_h) / static_cast<float>(in_h);
                x1 *= sx;
                x2 *= sx;
                y1 *= sy;
                y2 *= sy;
            }

            Rect bb;
            bb.x = static_cast<float>(ClampInt(static_cast<int>(x1), 0, orig_w - 1));
            bb.y = static_cast<float>(ClampInt(static_cast<int>(y1), 0, orig_h - 1));
            const float rx2 = static_cast<float>(ClampInt(static_cast<int>(x2), 0, orig_w - 1));
            const float ry2 = static_cast<float>(ClampInt(static_cast<int>(y2), 0, orig_h - 1));
            bb.w = std::max(0.0f, rx2 - bb.x);
            bb.h = std::max(0.0f, ry2 - bb.y);
            if (bb.w <= 1.0f || bb.h <= 1.0f) continue;

            boxes.push_back(bb);
            scores.push_back(score);

            if (output_landmarks_ && !lmk.data.empty() && lmk.n == n) {
                std::array<Point2f, 5> pts{};
                for (int k = 0; k < 5; ++k) {
                    const float lx = lmk.data[static_cast<size_t>(i) * 10 + k * 2 + 0];
                    const float ly = lmk.data[static_cast<size_t>(i) * 10 + k * 2 + 1];
                    float px = (p.cx + lx * var0 * p.w) * static_cast<float>(in_w);
                    float py = (p.cy + ly * var0 * p.h) * static_cast<float>(in_h);
                    if (use_lb) {
                        px = (px - static_cast<float>(lb.pad_x)) * inv_s;
                        py = (py - static_cast<float>(lb.pad_y)) * inv_s;
                    } else {
                        const float sx = static_cast<float>(orig_w) / static_cast<float>(in_w);
                        const float sy = static_cast<float>(orig_h) / static_cast<float>(in_h);
                        px *= sx;
                        py *= sy;
                    }
                    pts[k].x = static_cast<float>(ClampInt(static_cast<int>(px), 0, orig_w - 1));
                    pts[k].y = static_cast<float>(ClampInt(static_cast<int>(py), 0, orig_h - 1));
                }
                lmks.push_back(pts);
            }
        }

        if (boxes.empty()) {
            if (!logged_no_face_detail_.exchange(true)) {
                LogInfo("[ai_face_det] no face passed conf=" + std::to_string(conf_thresh_) +
                        ", max_face_prob=" + std::to_string(max_face_prob));
            }
            return;
        }

        std::vector<int> keep;
        NmsSorted(boxes, scores, nms_thresh_, keep);
        if (keep.empty()) return;

        const int out_n = std::min<int>(max_faces_, static_cast<int>(keep.size()));
        out.faces.reserve(static_cast<size_t>(out_n));
        for (int i = 0; i < out_n; ++i) {
            const int k = keep[static_cast<size_t>(i)];
            FaceDetItem item;
            item.bbox = boxes[static_cast<size_t>(k)];
            item.score = scores[static_cast<size_t>(k)];
            item.track_id = -1;
            if (output_landmarks_ && k < static_cast<int>(lmks.size())) {
                item.has_landmarks = true;
                item.landmarks = lmks[static_cast<size_t>(k)];
            }
            out.faces.push_back(std::move(item));
        }
    }

#endif

    std::string id_;
    std::string model_path_;

    float conf_thresh_ = 0.6f;
    float nms_thresh_ = 0.4f;
    int max_faces_ = 10;
    bool output_landmarks_ = true;
    int face_class_index_ = -1;  // -1=auto, 0/1=explicit
    bool letterbox_ = false;

    std::string input_format_ = "rgb";

    std::vector<int> steps_;
    std::vector<std::vector<int>> min_sizes_;

    std::shared_ptr<SpscQueue<FramePtr>> input_queue_;
    std::vector<std::shared_ptr<SpscQueue<FramePtr>>> output_queues_;

    std::vector<uint8_t> input_buf_;

    ModelHandle model_handle_ = kInvalidModelHandle;
    int model_w_ = 320;
    int model_h_ = 320;
    uint32_t n_output_ = 0;

    std::atomic<bool> logged_io_{false};
    std::atomic<bool> logged_decode_{false};
    std::atomic<bool> logged_decode_detail_{false};
    std::atomic<bool> logged_conf_probe_{false};
    std::atomic<bool> logged_no_face_detail_{false};
};

REGISTER_NODE(AiFaceDetNode, "ai_face_det");

}  // namespace rk3588