#include "yolo_runtime_backend.h"
#include "yolo_cpu_image_utils.h"
#include "yolo_letterbox.h"
#include "yolo_nms_iou.h"
#include "yolo_onnx_input_utils.h"

#include <onnxruntime_cxx_api.h>

#include <algorithm>
#include <cctype>
#include <chrono>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <map>
#include <memory>
#include <mutex>
#include <numeric>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace fs = std::filesystem;

namespace {

struct runtime_config {
    int width = 640;
    int height = 640;
    bool bgr = false;
    std::string layout = "nchw";
    std::vector<float> mean{0.0f, 0.0f, 0.0f};
    std::vector<float> standard_deviation{255.0f, 255.0f, 255.0f};
    std::vector<std::string> labels;
    yolo_task_type_t task = YOLO_TASK_DETECT;
};

struct candidate {
    int class_id = 0;
    float score = 0.0f;
    float x = 0.0f;
    float y = 0.0f;
    float width = 0.0f;
    float height = 0.0f;
    float angle = 0.0f;
    std::vector<float> mask_coefficients;
    std::vector<yolo_keypoint_v1_t> keypoints;
};

std::string trim(std::string value) {
    const auto first = value.find_first_not_of(" \t\r\n\"'");
    if (first == std::string::npos) return {};
    const auto last = value.find_last_not_of(" \t\r\n\"'");
    return value.substr(first, last - first + 1);
}

std::map<std::string, std::string> read_yaml(const fs::path& path) {
    std::ifstream input(path);
    if (!input) throw std::runtime_error("cannot open model config: " + path.string());
    std::map<std::string, std::string> values;
    std::map<int, std::string> sections;
    std::string line;
    while (std::getline(input, line)) {
        const auto comment = line.find('#');
        if (comment != std::string::npos) line.erase(comment);
        if (trim(line).empty()) continue;
        const int indent = static_cast<int>(line.find_first_not_of(" \t"));
        const std::string content = trim(line);
        const auto separator = content.find(':');
        if (separator == std::string::npos) continue;
        for (auto it = sections.begin(); it != sections.end();) {
            if (it->first >= indent) it = sections.erase(it); else ++it;
        }
        const std::string key = trim(content.substr(0, separator));
        const std::string value = trim(content.substr(separator + 1));
        std::string full_key;
        for (const auto& section : sections) {
            if (!full_key.empty()) full_key += '.';
            full_key += section.second;
        }
        if (!full_key.empty()) full_key += '.';
        full_key += key;
        if (value.empty()) sections[indent] = key;
        else values[full_key] = value;
    }
    return values;
}

std::vector<std::string> split_list(std::string value) {
    value = trim(value);
    if (!value.empty() && value.front() == '[') value.erase(value.begin());
    if (!value.empty() && value.back() == ']') value.pop_back();
    std::vector<std::string> items;
    std::stringstream stream(value);
    std::string item;
    while (std::getline(stream, item, ',')) {
        item = trim(item);
        if (!item.empty()) items.push_back(item);
    }
    return items;
}

std::vector<float> float_list(const std::map<std::string, std::string>& values, const char* key, std::vector<float> fallback) {
    const auto found = values.find(key);
    if (found == values.end()) return fallback;
    std::vector<float> result;
    for (const auto& item : split_list(found->second)) {
        try { result.push_back(std::stof(item)); } catch (...) { return fallback; }
    }
    return result.empty() ? fallback : result;
}

runtime_config load_config(const fs::path& config_path) {
    const auto values = read_yaml(config_path);
    runtime_config config;
    if (const auto it = values.find("task"); it != values.end()) {
        auto task = trim(it->second);
        std::transform(task.begin(), task.end(), task.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
        if (task == "cls" || task == "classify" || task == "classification") config.task = YOLO_TASK_CLS;
        else if (task == "obb") config.task = YOLO_TASK_OBB;
        else if (task == "seg" || task == "segment") config.task = YOLO_TASK_SEG;
        else if (task == "pose") config.task = YOLO_TASK_POSE;
        else if (task != "detect") throw std::runtime_error("Windows CPU Runtime supports Detect, CLS, OBB, Seg, and Pose models only");
    }
    if (const auto it = values.find("input.width"); it != values.end()) config.width = std::stoi(it->second);
    if (const auto it = values.find("input.height"); it != values.end()) config.height = std::stoi(it->second);
    if (const auto it = values.find("input.format"); it != values.end()) {
        auto format = trim(it->second);
        std::transform(format.begin(), format.end(), format.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
        config.bgr = format == "bgr";
    }
    if (const auto it = values.find("input.layout"); it != values.end()) config.layout = trim(it->second);
    config.mean = float_list(values, "input.normalization.mean", config.mean);
    config.standard_deviation = float_list(values, "input.normalization.std", config.standard_deviation);
    if (const auto it = values.find("labels"); it != values.end()) config.labels = split_list(it->second);
    if (const auto it = values.find("labels_file"); it != values.end()) {
        fs::path labels_path(trim(it->second));
        if (labels_path.is_relative()) labels_path = config_path.parent_path() / labels_path;
        std::ifstream labels(labels_path);
        std::string label;
        while (std::getline(labels, label)) {
            label = trim(label);
            if (!label.empty()) config.labels.push_back(label);
        }
    }
    if (config.width <= 0 || config.height <= 0) throw std::runtime_error("invalid configured model dimensions");
    return config;
}

float bilinear_channel(
    const yolo_image_view_v1_t& image,
    float source_x,
    float source_y,
    int channel,
    bool quantize_to_u8) {
    const int x0 = std::clamp(static_cast<int>(std::floor(source_x)), 0, image.width - 1);
    const int y0 = std::clamp(static_cast<int>(std::floor(source_y)), 0, image.height - 1);
    const int x1 = std::min(x0 + 1, image.width - 1);
    const int y1 = std::min(y0 + 1, image.height - 1);
    const float fx = source_x - x0;
    const float fy = source_y - y0;
    float r, g, b;
    yolo_cpu_image::pixel_rgb(image, x0, y0, r, g, b); const float top_left = channel == 0 ? r : channel == 1 ? g : b;
    yolo_cpu_image::pixel_rgb(image, x1, y0, r, g, b); const float top_right = channel == 0 ? r : channel == 1 ? g : b;
    yolo_cpu_image::pixel_rgb(image, x0, y1, r, g, b); const float bottom_left = channel == 0 ? r : channel == 1 ? g : b;
    yolo_cpu_image::pixel_rgb(image, x1, y1, r, g, b); const float bottom_right = channel == 0 ? r : channel == 1 ? g : b;
    if (quantize_to_u8) {
        return yolo_letterbox::bilinear_u8(
            static_cast<unsigned char>(top_left),
            static_cast<unsigned char>(top_right),
            static_cast<unsigned char>(bottom_left),
            static_cast<unsigned char>(bottom_right),
            fx,
            fy);
    }
    const float top = top_left * (1.0f - fx) + top_right * fx;
    const float bottom = bottom_left * (1.0f - fx) + bottom_right * fx;
    return top * (1.0f - fy) + bottom * fy;
}

yolo_nms::box nms_box(const candidate& value) {
    return {value.x, value.y, value.width, value.height, value.angle};
}

} // namespace

struct yolo_runtime_backend {
    std::unique_ptr<Ort::Env> environment;
    Ort::SessionOptions session_options;
    std::unique_ptr<Ort::Session> session;
    Ort::MemoryInfo memory_info{nullptr};
    std::string input_name;
    std::vector<std::string> output_names;
    std::vector<int64_t> input_shape;
    runtime_config config;
    float confidence_threshold = 0.25f;
    float nms_threshold = 0.45f;
    bool class_agnostic_nms = false;
    bool nchw = true;
    std::string last_error;
    std::mutex mutex;
    std::vector<std::vector<unsigned char>> segmentation_masks;
};

extern "C" int yolo_backend_get_abi_version(void) { return YOLO_RUNTIME_BACKEND_ABI_VERSION; }

extern "C" int yolo_backend_create(
    const char* model_path,
    const char* config_path,
    float confidence_threshold,
    float nms_threshold,
    int class_agnostic_nms,
    yolo_runtime_backend_t** output) {
    if (!model_path || !config_path || !output) return YOLO_STATUS_INVALID_ARGUMENT;
    try {
        auto backend = std::make_unique<yolo_runtime_backend>();
        backend->config = load_config(fs::u8path(config_path));
        if (backend->config.task != YOLO_TASK_DETECT && backend->config.task != YOLO_TASK_CLS &&
            backend->config.task != YOLO_TASK_OBB && backend->config.task != YOLO_TASK_SEG &&
            backend->config.task != YOLO_TASK_POSE)
            throw std::runtime_error("Windows CPU Runtime supports Detect, CLS, OBB, Seg, and Pose models only");
        backend->confidence_threshold = confidence_threshold;
        backend->nms_threshold = nms_threshold;
        backend->class_agnostic_nms = class_agnostic_nms != 0;
        backend->environment = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "rk3588_ai_windows_cpu");
        backend->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        backend->session_options.SetIntraOpNumThreads(std::max(1u, std::thread::hardware_concurrency() / 2));
        backend->session = std::make_unique<Ort::Session>(*backend->environment, fs::u8path(model_path).c_str(), backend->session_options);
        if (backend->session->GetInputCount() != 1) throw std::runtime_error("Windows CPU Runtime requires exactly one ONNX input");
        Ort::AllocatorWithDefaultOptions allocator;
        auto input_name = backend->session->GetInputNameAllocated(0, allocator);
        backend->input_name = input_name.get();
        backend->input_shape = backend->session->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
        if (backend->input_shape.size() != 4) throw std::runtime_error("Windows CPU Runtime requires a four-dimensional image input");
        backend->nchw = yolo_onnx_input::resolve_nchw(backend->input_shape, backend->config.layout);
        yolo_onnx_input::validate_three_channels(backend->input_shape, backend->nchw);
        const size_t outputs = backend->session->GetOutputCount();
        if (outputs == 0) throw std::runtime_error("ONNX model has no outputs");
        for (size_t index = 0; index < outputs; ++index) {
            auto output_name = backend->session->GetOutputNameAllocated(index, allocator);
            backend->output_names.emplace_back(output_name.get());
        }
        backend->memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        *output = backend.release();
        return YOLO_STATUS_OK;
    } catch (const std::exception& exception) {
        return YOLO_STATUS_INITIALIZATION_FAILED;
    }
}

extern "C" int yolo_backend_infer(yolo_runtime_backend_t* backend, const yolo_image_view_v1_t* image, yolo_backend_result_list_t* results) {
    if (!backend || !yolo_cpu_image::valid(image) || !results) return YOLO_STATUS_INVALID_ARGUMENT;
    try {
        std::lock_guard<std::mutex> lock(backend->mutex);
        std::memset(results, 0, sizeof(*results));
        results->task = backend->config.task;
        const auto total_start = std::chrono::steady_clock::now();

        const bool nchw = backend->nchw;
        const int model_height = nchw && backend->input_shape[2] > 0 ? static_cast<int>(backend->input_shape[2])
            : !nchw && backend->input_shape[1] > 0 ? static_cast<int>(backend->input_shape[1]) : backend->config.height;
        const int model_width = nchw && backend->input_shape[3] > 0 ? static_cast<int>(backend->input_shape[3])
            : !nchw && backend->input_shape[2] > 0 ? static_cast<int>(backend->input_shape[2]) : backend->config.width;
        if (model_width <= 0 || model_height <= 0) throw std::runtime_error("invalid ONNX input dimensions");
        const auto preprocess_start = std::chrono::steady_clock::now();
        const bool classification = backend->config.task == YOLO_TASK_CLS;
        const auto letterbox = classification
            ? yolo_letterbox::plan{1.0f, model_width, model_height, 0, 0}
            : yolo_letterbox::make_plan(image->width, image->height, model_width, model_height);
        const float scale = letterbox.scale;
        const int resized_width = letterbox.resized_width;
        const int resized_height = letterbox.resized_height;
        const int pad_x = letterbox.pad_x;
        const int pad_y = letterbox.pad_y;
        const size_t values = static_cast<size_t>(model_width) * model_height * 3;
        std::vector<float> tensor(values, 0.0f);
        auto write = [&](int channel, int x, int y, float value) {
            const float mean = backend->config.mean.size() > static_cast<size_t>(channel) ? backend->config.mean[channel] : 0.0f;
            const float standard_deviation = backend->config.standard_deviation.size() > static_cast<size_t>(channel) && std::abs(backend->config.standard_deviation[channel]) > 0.000001f
                ? backend->config.standard_deviation[channel] : 1.0f;
            const size_t index = nchw ? (static_cast<size_t>(channel) * model_height + y) * model_width + x
                : (static_cast<size_t>(y) * model_width + x) * 3 + channel;
            tensor[index] = (value - mean) / standard_deviation;
        };
        for (int y = 0; y < model_height; ++y)
            for (int x = 0; x < model_width; ++x)
                for (int channel = 0; channel < 3; ++channel) write(channel, x, y, 114.0f);
        for (int y = 0; y < resized_height; ++y) {
            const float source_y = yolo_letterbox::source_coordinate(y, image->height, resized_height);
            for (int x = 0; x < resized_width; ++x) {
                const float source_x = yolo_letterbox::source_coordinate(x, image->width, resized_width);
                const float red = bilinear_channel(*image, source_x, source_y, 0, !classification);
                const float green = bilinear_channel(*image, source_x, source_y, 1, !classification);
                const float blue = bilinear_channel(*image, source_x, source_y, 2, !classification);
                write(0, x + pad_x, y + pad_y, backend->config.bgr ? blue : red);
                write(1, x + pad_x, y + pad_y, green);
                write(2, x + pad_x, y + pad_y, backend->config.bgr ? red : blue);
            }
        }
        results->performance.preprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - preprocess_start).count();

        std::vector<int64_t> input_shape = backend->input_shape;
        for (auto& dimension : input_shape) if (dimension <= 0) dimension = 1;
        if (nchw) { input_shape[1] = 3; input_shape[2] = model_height; input_shape[3] = model_width; }
        else { input_shape[1] = model_height; input_shape[2] = model_width; input_shape[3] = 3; }
        auto input_tensor = Ort::Value::CreateTensor<float>(backend->memory_info, tensor.data(), tensor.size(), input_shape.data(), input_shape.size());
        std::vector<const char*> output_names;
        for (const auto& name : backend->output_names) output_names.push_back(name.c_str());
        const char* input_names[] = { backend->input_name.c_str() };
        const auto inference_start = std::chrono::steady_clock::now();
        auto output_values = backend->session->Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names.data(), output_names.size());
        results->performance.inference_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - inference_start).count();

        const auto postprocess_start = std::chrono::steady_clock::now();
        if (classification) {
            Ort::Value* classification_value = nullptr;
            for (auto& output : output_values) {
                if (!output.IsTensor()) continue;
                const auto info = output.GetTensorTypeAndShapeInfo();
                if (info.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && info.GetElementCount() > 0) {
                    classification_value = &output;
                    break;
                }
            }
            if (!classification_value) throw std::runtime_error("cannot find FP32 classification output");
            const auto output_info = classification_value->GetTensorTypeAndShapeInfo();
            const size_t class_count = output_info.GetElementCount();
            const float* raw_scores = classification_value->GetTensorData<float>();
            std::vector<float> scores(raw_scores, raw_scores + class_count);
            for (const float score : scores)
                if (!std::isfinite(score)) throw std::runtime_error("classification output contains a non-finite score");

            const double raw_sum = std::accumulate(scores.begin(), scores.end(), 0.0);
            const bool probabilities = std::all_of(scores.begin(), scores.end(), [](float score) {
                return score >= 0.0f && score <= 1.0f;
            }) && std::abs(raw_sum - 1.0) < 0.01;
            if (!probabilities) {
                const float maximum = *std::max_element(scores.begin(), scores.end());
                double total = 0.0;
                for (float& score : scores) {
                    score = std::exp(score - maximum);
                    total += score;
                }
                if (total <= 0.0) throw std::runtime_error("classification Softmax total is zero");
                for (float& score : scores) score = static_cast<float>(score / total);
            }

            std::vector<size_t> order(class_count);
            std::iota(order.begin(), order.end(), 0);
            std::sort(order.begin(), order.end(), [&](size_t first, size_t second) {
                return scores[first] > scores[second];
            });
            const size_t result_count = std::min<size_t>({5, class_count, YOLO_BACKEND_MAX_RESULTS});
            for (size_t index = 0; index < result_count; ++index) {
                const size_t class_id = order[index];
                auto& result = results->results[results->count++];
                result.task = YOLO_TASK_CLS;
                result.class_id = static_cast<int>(class_id);
                result.score = scores[class_id];
                const std::string label = class_id < backend->config.labels.size()
                    ? backend->config.labels[class_id] : "class_" + std::to_string(class_id);
                std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
            }
            results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
            results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
            return YOLO_STATUS_OK;
        }

        if (backend->config.task == YOLO_TASK_POSE) {
            Ort::Value* pose_keypoints = nullptr;
            std::vector<Ort::Value*> pose_heads;
            for (auto& output_value : output_values) {
                if (!output_value.IsTensor()) continue;
                const auto info = output_value.GetTensorTypeAndShapeInfo();
                const auto shape = info.GetShape();
                if (shape.size() != 4 || info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) continue;
                if (shape[2] == 3 && shape[1] > 0 && shape[1] <= YOLO_BACKEND_MAX_KEYPOINTS)
                    pose_keypoints = &output_value;
                else if (shape[1] > 64)
                    pose_heads.push_back(&output_value);
            }
            if (pose_keypoints && !pose_heads.empty()) {
                const auto keypoint_shape = pose_keypoints->GetTensorTypeAndShapeInfo().GetShape();
                const int pose_keypoint_count = static_cast<int>(keypoint_shape[1]);
                const int total_positions = static_cast<int>(keypoint_shape[3]);
                const float* keypoint_data = pose_keypoints->GetTensorData<float>();
                std::sort(pose_heads.begin(), pose_heads.end(), [](Ort::Value* first, Ort::Value* second) {
                    return first->GetTensorTypeAndShapeInfo().GetShape()[2] >
                        second->GetTensorTypeAndShapeInfo().GetShape()[2];
                });
                std::vector<candidate> pose_candidates;
                int position_offset = 0;
                for (Ort::Value* head : pose_heads) {
                    const auto shape = head->GetTensorTypeAndShapeInfo().GetShape();
                    const int head_channels = static_cast<int>(shape[1]);
                    const int grid_height = static_cast<int>(shape[2]);
                    const int grid_width = static_cast<int>(shape[3]);
                    const int head_positions = grid_height * grid_width;
                    const int pose_class_count = head_channels - 64;
                    if (pose_class_count <= 0 || position_offset + head_positions > total_positions)
                        throw std::runtime_error("invalid split pose output shape");
                    const int stride = backend->config.height / grid_height;
                    const float* head_data = head->GetTensorData<float>();
                    const auto head_at = [&](int channel, int position) -> float {
                        return head_data[static_cast<size_t>(channel) * head_positions + position];
                    };
                    for (int position = 0; position < head_positions; ++position) {
                        int class_id = 0;
                        float score = 1.0f / (1.0f + std::exp(-head_at(64, position)));
                        for (int class_index = 1; class_index < pose_class_count; ++class_index) {
                            const float value = 1.0f / (1.0f + std::exp(-head_at(64 + class_index, position)));
                            if (value > score) { score = value; class_id = class_index; }
                        }
                        if (score < backend->confidence_threshold) continue;
                        float distances[4]{};
                        for (int side = 0; side < 4; ++side) {
                            float maximum = head_at(side * 16, position);
                            for (int bin = 1; bin < 16; ++bin)
                                maximum = std::max(maximum, head_at(side * 16 + bin, position));
                            float denominator = 0.0f;
                            float numerator = 0.0f;
                            for (int bin = 0; bin < 16; ++bin) {
                                const float probability = std::exp(head_at(side * 16 + bin, position) - maximum);
                                denominator += probability;
                                numerator += probability * bin;
                            }
                            distances[side] = numerator / denominator;
                        }
                        const int grid_x = position % grid_width;
                        const int grid_y = position / grid_width;
                        const float left = (grid_x + 0.5f - distances[0]) * stride;
                        const float top = (grid_y + 0.5f - distances[1]) * stride;
                        candidate item{
                            class_id, score,
                            std::clamp((left - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
                            std::clamp((top - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
                            std::clamp((distances[0] + distances[2]) * stride / scale, 0.0f, static_cast<float>(image->width)),
                            std::clamp((distances[1] + distances[3]) * stride / scale, 0.0f, static_cast<float>(image->height))
                        };
                        item.keypoints.resize(static_cast<size_t>(pose_keypoint_count));
                        const int global_position = position_offset + position;
                        for (int keypoint = 0; keypoint < pose_keypoint_count; ++keypoint) {
                            const size_t base = static_cast<size_t>(keypoint) * 3 * total_positions;
                            item.keypoints[static_cast<size_t>(keypoint)] = {
                                std::clamp((keypoint_data[base + global_position] - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
                                std::clamp((keypoint_data[base + total_positions + global_position] - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
                                keypoint_data[base + 2 * total_positions + global_position]
                            };
                        }
                        pose_candidates.push_back(std::move(item));
                    }
                    position_offset += head_positions;
                }
                std::sort(pose_candidates.begin(), pose_candidates.end(), [](const candidate& first, const candidate& second) {
                    return first.score > second.score;
                });
                std::vector<candidate> selected_pose;
                for (const auto& item : pose_candidates) {
                    if (selected_pose.size() >= YOLO_BACKEND_MAX_RESULTS) break;
                    const bool suppressed = std::any_of(selected_pose.begin(), selected_pose.end(), [&](const candidate& existing) {
                        return (backend->class_agnostic_nms || existing.class_id == item.class_id) &&
                            yolo_nms::overlap_iou(
                                backend->config.task, nms_box(existing), nms_box(item)) > backend->nms_threshold;
                    });
                    if (!suppressed) selected_pose.push_back(item);
                }
                for (const auto& item : selected_pose) {
                    auto& result = results->results[results->count++];
                    result.task = YOLO_TASK_POSE;
                    result.class_id = item.class_id;
                    result.score = item.score;
                    result.bbox = {item.x, item.y, item.width, item.height, 0.0f, 0};
                    result.keypoint_count = static_cast<int>(item.keypoints.size());
                    std::copy(item.keypoints.begin(), item.keypoints.end(), result.keypoints);
                    const std::string label = item.class_id >= 0 && item.class_id < static_cast<int>(backend->config.labels.size())
                        ? backend->config.labels[static_cast<size_t>(item.class_id)] : "class_" + std::to_string(item.class_id);
                    std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
                }
                results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
                results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
                return YOLO_STATUS_OK;
            }
        }

        Ort::Value* detection_value = nullptr;
        Ort::Value* prototype_value = nullptr;
        for (auto& output : output_values) {
            if (!output.IsTensor()) continue;
            const auto shape = output.GetTensorTypeAndShapeInfo().GetShape();
            if (shape.size() == 3 && !detection_value) detection_value = &output;
            if (shape.size() == 4 && !prototype_value) prototype_value = &output;
        }
        if (!detection_value) throw std::runtime_error("cannot find Ultralytics detection output");
        const auto output_info = detection_value->GetTensorTypeAndShapeInfo();
        const auto output_shape = output_info.GetShape();
        if (output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || output_shape.size() != 3)
            throw std::runtime_error("Windows CPU Runtime requires FP32 [1,C,N] or [1,N,C] output");
        const int64_t first_dimension = output_shape[1];
        const int64_t second_dimension = output_shape[2];
        // Usual Ultralytics output is [1, C, N] with N much larger than C.
        // A valid one-candidate export has N == 1, however, so do not mistake
        // [1, C, 1] for [1, N, C] merely because C is greater than N.
        const bool channels_first = first_dimension > 4 &&
            (second_dimension <= 4 || first_dimension <= second_dimension);
        const int channels = static_cast<int>(channels_first ? first_dimension : second_dimension);
        const int positions = static_cast<int>(channels_first ? second_dimension : first_dimension);
        if (channels <= 4 || positions <= 0) throw std::runtime_error("invalid detection output shape");

        int mask_channels = 0;
        int prototype_height = 0;
        int prototype_width = 0;
        bool prototype_channels_first = true;
        const float* prototypes = nullptr;
        if (backend->config.task == YOLO_TASK_SEG) {
            if (!prototype_value) throw std::runtime_error("cannot find Ultralytics segmentation prototype output");
            const auto prototype_info = prototype_value->GetTensorTypeAndShapeInfo();
            const auto prototype_shape = prototype_info.GetShape();
            if (prototype_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || prototype_shape.size() != 4)
                throw std::runtime_error("Windows CPU Runtime requires FP32 segmentation prototypes");
            prototype_channels_first = prototype_shape[1] == 32 ||
                (prototype_shape[1] > 0 && prototype_shape[1] <= 256 && prototype_shape[1] < prototype_shape[3]);
            mask_channels = static_cast<int>(prototype_channels_first ? prototype_shape[1] : prototype_shape[3]);
            prototype_height = static_cast<int>(prototype_channels_first ? prototype_shape[2] : prototype_shape[1]);
            prototype_width = static_cast<int>(prototype_channels_first ? prototype_shape[3] : prototype_shape[2]);
            if (mask_channels <= 0 || prototype_height <= 0 || prototype_width <= 0)
                throw std::runtime_error("invalid segmentation prototype output shape");
            prototypes = prototype_value->GetTensorData<float>();
        }

        const int geometry_channels = backend->config.task == YOLO_TASK_OBB ? 5 : 4;
        int keypoint_count = 0;
        int class_count = channels - geometry_channels - mask_channels;
        if (backend->config.task == YOLO_TASK_POSE) {
            class_count = backend->config.labels.empty() ? 1 : static_cast<int>(backend->config.labels.size());
            const int keypoint_values = channels - 4 - class_count;
            if (keypoint_values <= 0 || keypoint_values % 3 != 0)
                throw std::runtime_error("pose output must contain x/y/score triplets after class scores");
            keypoint_count = keypoint_values / 3;
            if (keypoint_count > YOLO_BACKEND_MAX_KEYPOINTS)
                throw std::runtime_error("pose output exceeds ABI keypoint capacity");
        }
        if (class_count <= 0) throw std::runtime_error("model output does not contain any class scores");
        const float* output = detection_value->GetTensorData<float>();
        const auto at = [&](int channel, int position) -> float {
            return channels_first ? output[static_cast<size_t>(channel) * positions + position]
                : output[static_cast<size_t>(position) * channels + channel];
        };
        std::vector<candidate> candidates;
        for (int position = 0; position < positions; ++position) {
            int class_id = 0;
            float score = at(4, position);
            for (int index = 1; index < class_count; ++index) {
                const float value = at(4 + index, position);
                if (value > score) { score = value; class_id = index; }
            }
            if (score < backend->confidence_threshold) continue;
            const float width = std::clamp(at(2, position) / scale, 0.0f, static_cast<float>(image->width));
            const float height = std::clamp(at(3, position) / scale, 0.0f, static_cast<float>(image->height));
            candidate item{
                class_id, score,
                std::clamp((at(0, position) - pad_x) / scale - width * 0.5f, 0.0f, static_cast<float>(image->width)),
                std::clamp((at(1, position) - pad_y) / scale - height * 0.5f, 0.0f, static_cast<float>(image->height)),
                width, height };
            if (backend->config.task == YOLO_TASK_OBB) item.angle = at(channels - 1, position);
            if (backend->config.task == YOLO_TASK_SEG) {
                item.mask_coefficients.resize(static_cast<size_t>(mask_channels));
                for (int index = 0; index < mask_channels; ++index)
                    item.mask_coefficients[static_cast<size_t>(index)] = at(4 + class_count + index, position);
            }
            if (backend->config.task == YOLO_TASK_POSE) {
                item.keypoints.resize(static_cast<size_t>(keypoint_count));
                for (int index = 0; index < keypoint_count; ++index) {
                    const int offset = 4 + class_count + index * 3;
                    item.keypoints[static_cast<size_t>(index)] = {
                        std::clamp((at(offset, position) - pad_x) / scale, 0.0f, static_cast<float>(image->width)),
                        std::clamp((at(offset + 1, position) - pad_y) / scale, 0.0f, static_cast<float>(image->height)),
                        at(offset + 2, position)
                    };
                }
            }
            candidates.push_back(std::move(item));
        }
        std::sort(candidates.begin(), candidates.end(), [](const candidate& a, const candidate& b) { return a.score > b.score; });
        std::vector<candidate> selected_candidates;
        for (const auto& item : candidates) {
            if (selected_candidates.size() >= YOLO_BACKEND_MAX_RESULTS) break;
            bool suppressed = false;
            for (const auto& existing : selected_candidates) {
                if ((backend->class_agnostic_nms || existing.class_id == item.class_id) &&
                    yolo_nms::overlap_iou(
                        backend->config.task, nms_box(existing), nms_box(item)) > backend->nms_threshold) {
                    suppressed = true;
                    break;
                }
            }
            if (suppressed) continue;
            selected_candidates.push_back(item);
        }

        backend->segmentation_masks.clear();
        backend->segmentation_masks.reserve(selected_candidates.size());
        const auto prototype_at = [&](int channel, int y, int x) -> float {
            if (prototype_channels_first)
                return prototypes[(static_cast<size_t>(channel) * prototype_height + y) * prototype_width + x];
            return prototypes[(static_cast<size_t>(y) * prototype_width + x) * mask_channels + channel];
        };
        const auto mask_logit = [&](const candidate& item, float source_x, float source_y) -> float {
            const float model_x = source_x * scale + pad_x;
            const float model_y = source_y * scale + pad_y;
            const float px = std::clamp((model_x + 0.5f) * prototype_width / backend->config.width - 0.5f,
                0.0f, static_cast<float>(prototype_width - 1));
            const float py = std::clamp((model_y + 0.5f) * prototype_height / backend->config.height - 0.5f,
                0.0f, static_cast<float>(prototype_height - 1));
            const int x0 = static_cast<int>(std::floor(px));
            const int y0 = static_cast<int>(std::floor(py));
            const int x1 = std::min(x0 + 1, prototype_width - 1);
            const int y1 = std::min(y0 + 1, prototype_height - 1);
            const float dx = px - x0;
            const float dy = py - y0;
            float value = 0.0f;
            for (int channel = 0; channel < mask_channels; ++channel) {
                const float top = prototype_at(channel, y0, x0) * (1.0f - dx) + prototype_at(channel, y0, x1) * dx;
                const float bottom = prototype_at(channel, y1, x0) * (1.0f - dx) + prototype_at(channel, y1, x1) * dx;
                value += item.mask_coefficients[static_cast<size_t>(channel)] * (top * (1.0f - dy) + bottom * dy);
            }
            return value;
        };

        for (const auto& item : selected_candidates) {
            if (results->count >= YOLO_BACKEND_MAX_RESULTS) break;
            int mask_min_x = image->width;
            int mask_min_y = image->height;
            int mask_max_x = -1;
            int mask_max_y = -1;
            std::vector<unsigned char> full_mask;
            if (backend->config.task == YOLO_TASK_SEG) {
                full_mask.assign(static_cast<size_t>(image->width) * image->height, 0);
                const int start_x = std::clamp(static_cast<int>(std::floor(item.x)), 0, image->width);
                const int start_y = std::clamp(static_cast<int>(std::floor(item.y)), 0, image->height);
                const int end_x = std::clamp(static_cast<int>(std::ceil(item.x + item.width)), 0, image->width);
                const int end_y = std::clamp(static_cast<int>(std::ceil(item.y + item.height)), 0, image->height);
                for (int y = start_y; y < end_y; ++y) {
                    for (int x = start_x; x < end_x; ++x) {
                        if (mask_logit(item, x + 0.5f, y + 0.5f) <= 0.0f) continue;
                        full_mask[static_cast<size_t>(y) * image->width + x] = 255;
                        mask_min_x = std::min(mask_min_x, x);
                        mask_min_y = std::min(mask_min_y, y);
                        mask_max_x = std::max(mask_max_x, x);
                        mask_max_y = std::max(mask_max_y, y);
                    }
                }
                if (mask_max_x < mask_min_x || mask_max_y < mask_min_y) continue;
            }
            auto& result = results->results[results->count++];
            result.task = backend->config.task;
            result.class_id = item.class_id;
            result.score = item.score;
            result.bbox.x = item.x;
            result.bbox.y = item.y;
            result.bbox.width = item.width;
            result.bbox.height = item.height;
            if (backend->config.task == YOLO_TASK_OBB) {
                result.bbox.angle = item.angle;
                result.bbox.has_angle = 1;
            }
            if (backend->config.task == YOLO_TASK_POSE) {
                result.keypoint_count = static_cast<int>(item.keypoints.size());
                std::copy(item.keypoints.begin(), item.keypoints.end(), result.keypoints);
            }
            if (backend->config.task == YOLO_TASK_SEG) {
                const int mask_width = mask_max_x - mask_min_x + 1;
                const int mask_height = mask_max_y - mask_min_y + 1;
                backend->segmentation_masks.emplace_back(static_cast<size_t>(mask_width) * mask_height);
                auto& mask = backend->segmentation_masks.back();
                for (int y = 0; y < mask_height; ++y)
                    std::copy_n(full_mask.data() + static_cast<size_t>(mask_min_y + y) * image->width + mask_min_x,
                        mask_width, mask.data() + static_cast<size_t>(y) * mask_width);
                result.mask_bbox.x = static_cast<float>(mask_min_x);
                result.mask_bbox.y = static_cast<float>(mask_min_y);
                result.mask_bbox.width = static_cast<float>(mask_width);
                result.mask_bbox.height = static_cast<float>(mask_height);
                result.mask_data = mask.data();
                result.mask_width = mask_width;
                result.mask_height = mask_height;
                result.mask_stride_bytes = mask_width;
                result.mask_origin_x = mask_min_x;
                result.mask_origin_y = mask_min_y;
            }
            const std::string label = item.class_id >= 0 && item.class_id < static_cast<int>(backend->config.labels.size())
                ? backend->config.labels[static_cast<size_t>(item.class_id)] : "class_" + std::to_string(item.class_id);
            std::snprintf(result.label, sizeof(result.label), "%s", label.c_str());
        }
        results->performance.postprocess_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - postprocess_start).count();
        results->performance.total_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - total_start).count();
        return YOLO_STATUS_OK;
    } catch (const std::exception& exception) {
        backend->last_error = exception.what();
        return YOLO_STATUS_BACKEND_FAILED;
    }
}

extern "C" const char* yolo_backend_last_error(const yolo_runtime_backend_t* backend) {
    return backend ? backend->last_error.c_str() : "backend is null";
}

extern "C" void yolo_backend_destroy(yolo_runtime_backend_t* backend) { delete backend; }