#include "yolo_runtime_backend.h"
#include "yolo_mask_utils.h"

#if defined(YOLO_BACKEND_SEG)
#include "yolov8_seg.h"
#elif defined(YOLO_BACKEND_CLS)
#include "mobilenet.h"
#elif defined(YOLO_BACKEND_POSE)
#include "yolov8-pose.h"
#elif defined(YOLO_BACKEND_OBB)
#include "yolov8-obb.h"
#else
#include "yolov8.h"
#endif

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <mutex>
#include <new>
#include <vector>

extern "C" void yolo_adapter_set_thresholds(float confidence, float nms);
#if defined(YOLO_BACKEND_SEG)
extern "C" void yolo_adapter_set_class_agnostic_nms(int enabled);
#endif
extern "C" void yolo_adapter_get_performance(
    double* preprocess_ms,
    double* inference_ms,
    double* postprocess_ms);

struct yolo_runtime_backend {
    rknn_app_context_t context{};
    std::mutex context_mutex;
    bool initialized = false;
    float confidence_threshold = 0.25f;
    float nms_threshold = 0.45f;
    bool class_agnostic_nms = false;
    std::vector<std::vector<unsigned char>> segmentation_masks;
};

namespace {

int packed_row_bytes(const yolo_image_view_v1_t& image) {
    switch (image.pixel_format) {
        case YOLO_PIXEL_FORMAT_GRAY8:
        case YOLO_PIXEL_FORMAT_NV12:
        case YOLO_PIXEL_FORMAT_NV21: return image.width;
        case YOLO_PIXEL_FORMAT_RGB888:
        case YOLO_PIXEL_FORMAT_BGR888: return image.width * 3;
        case YOLO_PIXEL_FORMAT_RGBA8888:
        case YOLO_PIXEL_FORMAT_BGRA8888: return image.width * 4;
        default: return 0;
    }
}

bool prepare_image(const yolo_image_view_v1_t& source, image_buffer_t& destination, std::vector<unsigned char>& owned) {
    const int packed_stride = packed_row_bytes(source);
    if (packed_stride <= 0) return false;
    const int source_stride = source.row_stride_bytes > 0 ? source.row_stride_bytes : packed_stride;
    if (source_stride < packed_stride) return false;
    int rows = source.height;
    if (source.pixel_format == YOLO_PIXEL_FORMAT_NV12 || source.pixel_format == YOLO_PIXEL_FORMAT_NV21) rows = source.height * 3 / 2;
    const bool swap_red_blue = source.pixel_format == YOLO_PIXEL_FORMAT_BGR888 || source.pixel_format == YOLO_PIXEL_FORMAT_BGRA8888;
    const bool copy_rows = source_stride != packed_stride;
    const unsigned char* pixels = source.data;
    if (swap_red_blue || copy_rows) {
        owned.resize(static_cast<size_t>(packed_stride) * rows);
        for (int row = 0; row < rows; ++row) {
            const unsigned char* input = source.data + static_cast<size_t>(row) * source_stride;
            unsigned char* output = owned.data() + static_cast<size_t>(row) * packed_stride;
            std::memcpy(output, input, static_cast<size_t>(packed_stride));
            if (swap_red_blue && row < source.height) {
                const int channels = source.pixel_format == YOLO_PIXEL_FORMAT_BGR888 ? 3 : 4;
                for (int column = 0; column < source.width; ++column) std::swap(output[column * channels], output[column * channels + 2]);
            }
        }
        pixels = owned.data();
    }
    destination.width = source.width;
    destination.height = source.height;
    destination.width_stride = source.width;
    destination.height_stride = source.height;
    destination.virt_addr = const_cast<unsigned char*>(pixels);
    destination.size = packed_stride * rows;
    destination.fd = 0;
    switch (source.pixel_format) {
        case YOLO_PIXEL_FORMAT_GRAY8: destination.format = IMAGE_FORMAT_GRAY8; break;
        case YOLO_PIXEL_FORMAT_RGB888:
        case YOLO_PIXEL_FORMAT_BGR888: destination.format = IMAGE_FORMAT_RGB888; break;
        case YOLO_PIXEL_FORMAT_RGBA8888:
        case YOLO_PIXEL_FORMAT_BGRA8888: destination.format = IMAGE_FORMAT_RGBA8888; break;
        case YOLO_PIXEL_FORMAT_NV12: destination.format = IMAGE_FORMAT_YUV420SP_NV12; break;
        case YOLO_PIXEL_FORMAT_NV21: destination.format = IMAGE_FORMAT_YUV420SP_NV21; break;
        default: return false;
    }
    return true;
}

bool valid_image(const yolo_image_view_v1_t* image) {
    return image && image->struct_size >= sizeof(yolo_image_view_v1_t) && image->version == YOLO_IMAGE_VIEW_ABI_VERSION && image->data && image->width > 0 && image->height > 0;
}

#if defined(YOLO_BACKEND_OBB)
void set_box(yolo_backend_result_t& target, const image_obb_box_t& box) {
    target.bbox.x = static_cast<float>(box.x);
    target.bbox.y = static_cast<float>(box.y);
    target.bbox.width = static_cast<float>(box.w);
    target.bbox.height = static_cast<float>(box.h);
    target.bbox.angle = box.angle;
    target.bbox.has_angle = 1;
}
#else
void set_box(yolo_backend_result_t& target, const image_rect_t& box) {
    target.bbox.x = static_cast<float>(box.left);
    target.bbox.y = static_cast<float>(box.top);
    target.bbox.width = static_cast<float>(box.right - box.left);
    target.bbox.height = static_cast<float>(box.bottom - box.top);
}
#endif

} // namespace

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 || !output) return YOLO_STATUS_INVALID_ARGUMENT;
    (void)config_path;
    auto* backend = new (std::nothrow) yolo_runtime_backend();
    if (!backend) return YOLO_STATUS_INITIALIZATION_FAILED;
    backend->confidence_threshold = confidence_threshold;
    backend->nms_threshold = nms_threshold;
    backend->class_agnostic_nms = class_agnostic_nms != 0;
#if defined(YOLO_BACKEND_SEG)
    const int status = init_yolov8_seg_model(model_path, &backend->context);
#elif defined(YOLO_BACKEND_CLS)
    const int status = init_mobilenet_model(model_path, &backend->context);
#elif defined(YOLO_BACKEND_POSE)
    const int status = init_yolov8_pose_model(model_path, &backend->context);
#elif defined(YOLO_BACKEND_OBB)
    const int status = init_yolov8_obb_model(model_path, &backend->context);
#else
    const int status = init_yolov8_model(model_path, &backend->context);
#endif
    if (status != 0) { delete backend; return YOLO_STATUS_INITIALIZATION_FAILED; }
    backend->initialized = true;
    *output = backend;
    return YOLO_STATUS_OK;
}

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 || !backend->initialized || !valid_image(image) || !results) return YOLO_STATUS_INVALID_ARGUMENT;
    std::lock_guard<std::mutex> lock(backend->context_mutex);
    std::memset(results, 0, sizeof(*results));
    image_buffer_t input{};
    std::vector<unsigned char> owned_pixels;
    if (!prepare_image(*image, input, owned_pixels)) return YOLO_STATUS_INVALID_ARGUMENT;
    const auto start = std::chrono::steady_clock::now();
#if defined(YOLO_BACKEND_CLS)
    mobilenet_result classifications[5]{};
    const int status = inference_mobilenet_model(&backend->context, &input, classifications, 5);
    results->task = YOLO_TASK_CLS;
    if (status == 0) {
        for (const auto& classification : classifications) {
            if (classification.cls < 0 || results->count >= YOLO_BACKEND_MAX_RESULTS) continue;
            auto& result = results->results[results->count++];
            result.task = YOLO_TASK_CLS;
            result.class_id = classification.cls;
            result.score = classification.score;
            std::snprintf(result.label, sizeof(result.label), "class_%d", result.class_id);
        }
    }
#else
    object_detect_result_list detections{};
#if defined(YOLO_BACKEND_SEG)
    yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
    yolo_adapter_set_class_agnostic_nms(backend->class_agnostic_nms ? 1 : 0);
    const int status = inference_yolov8_seg_model(&backend->context, &input, &detections);
    results->task = YOLO_TASK_SEG;
#elif defined(YOLO_BACKEND_POSE)
    yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
    const int status = inference_yolov8_pose_model(&backend->context, &input, &detections);
    results->task = YOLO_TASK_POSE;
#elif defined(YOLO_BACKEND_OBB)
    yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
    const int status = inference_yolov8_obb_model(&backend->context, &input, &detections);
    results->task = YOLO_TASK_OBB;
#else
    yolo_adapter_set_thresholds(backend->confidence_threshold, backend->nms_threshold);
    const int status = inference_yolov8_model(&backend->context, &input, &detections);
    results->task = YOLO_TASK_DETECT;
#endif
    if (status == 0) {
#if defined(YOLO_BACKEND_SEG)
        backend->segmentation_masks.clear();
        backend->segmentation_masks.reserve(static_cast<size_t>(std::min(detections.count, YOLO_BACKEND_MAX_RESULTS)));
#else
        results->count = std::min(detections.count, YOLO_BACKEND_MAX_RESULTS);
#endif
        const int source_count = std::min(detections.count, YOLO_BACKEND_MAX_RESULTS);
        for (int index = 0; index < source_count; ++index) {
            const auto& source = detections.results[index];
#if defined(YOLO_BACKEND_SEG)
            unsigned char* full_mask = detections.results_seg[index].seg_mask;
            const auto bounds = yolo_find_mask_bounds(full_mask, image->width, image->height, image->width);
            auto mask_roi = yolo_copy_binary_mask_roi(full_mask, image->width, bounds);
            std::free(full_mask);
            detections.results_seg[index].seg_mask = nullptr;
            if (bounds.empty() || mask_roi.empty()) continue;
            backend->segmentation_masks.push_back(std::move(mask_roi));
            auto& result = results->results[results->count++];
#else
            auto& result = results->results[index];
#endif
            result.task = results->task;
            result.class_id = source.cls_id;
            result.score = source.prop;
            std::snprintf(result.label, sizeof(result.label), "class_%d", source.cls_id);
            set_box(result, source.box);
#if defined(YOLO_BACKEND_SEG)
            auto& mask = backend->segmentation_masks.back();
            result.mask_bbox.x = static_cast<float>(bounds.x);
            result.mask_bbox.y = static_cast<float>(bounds.y);
            result.mask_bbox.width = static_cast<float>(bounds.width);
            result.mask_bbox.height = static_cast<float>(bounds.height);
            result.mask_data = mask.data();
            result.mask_width = bounds.width;
            result.mask_height = bounds.height;
            result.mask_stride_bytes = bounds.width;
            result.mask_origin_x = bounds.x;
            result.mask_origin_y = bounds.y;
#elif defined(YOLO_BACKEND_POSE)
            result.keypoint_count = 17;
            for (int keypoint = 0; keypoint < result.keypoint_count; ++keypoint) {
                result.keypoints[keypoint] = { source.keypoints[keypoint][0], source.keypoints[keypoint][1], source.keypoints[keypoint][2] };
            }
#endif
        }
#if defined(YOLO_BACKEND_SEG)
        for (int index = source_count; index < detections.count; ++index) {
            std::free(detections.results_seg[index].seg_mask);
            detections.results_seg[index].seg_mask = nullptr;
        }
#endif
    }
#endif
    const double elapsed = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
    results->performance.total_ms = elapsed;
    yolo_adapter_get_performance(&results->performance.preprocess_ms, &results->performance.inference_ms, &results->performance.postprocess_ms);
    return status == 0 ? YOLO_STATUS_OK : YOLO_STATUS_BACKEND_FAILED;
}

extern "C" const char* yolo_backend_last_error(const yolo_runtime_backend_t* backend) {
    (void)backend;
    return "RKNN backend did not provide detailed error information";
}

extern "C" void yolo_backend_destroy(yolo_runtime_backend_t* backend) {
    if (!backend) return;
    if (backend->initialized) {
#if defined(YOLO_BACKEND_SEG)
        release_yolov8_seg_model(&backend->context);
#elif defined(YOLO_BACKEND_CLS)
        release_mobilenet_model(&backend->context);
#elif defined(YOLO_BACKEND_POSE)
        release_yolov8_pose_model(&backend->context);
#elif defined(YOLO_BACKEND_OBB)
        release_yolov8_obb_model(&backend->context);
#else
        release_yolov8_model(&backend->context);
#endif
    }
    delete backend;
}