#include "engines/zed_native.h"
#include "util/env.h"

#include <iostream>
#include <string>

namespace {

sl::OBJECT_DETECTION_MODEL parse_detection_model(const std::string& model) {
    if (model == "MULTI_CLASS_BOX_FAST") {
        return sl::OBJECT_DETECTION_MODEL::MULTI_CLASS_BOX_FAST;
    }
    if (model == "MULTI_CLASS_BOX_ACCURATE") {
        return sl::OBJECT_DETECTION_MODEL::MULTI_CLASS_BOX_ACCURATE;
    }
    return sl::OBJECT_DETECTION_MODEL::MULTI_CLASS_BOX_MEDIUM;
}

} // namespace

namespace engines {

DetectionConfig DetectionConfig::from_env() {
    DetectionConfig c;
    c.model = util::get_env_str("ZED_DETECTION_MODEL", c.model);
    c.confidence_threshold = util::get_env_float("ZED_DETECTION_CONFIDENCE", c.confidence_threshold);
    c.tracking = util::get_env_bool("ZED_DETECTION_TRACKING", c.tracking);
    return c;
}

ZedNative::ZedNative(const DetectionConfig& config) : config_(config) {
    det_params_.detection_model = parse_detection_model(config_.model);
    det_params_.enable_tracking = config_.tracking;
    det_params_.enable_segmentation = false;

    det_rt_params_.detection_confidence_threshold = config_.confidence_threshold;
}

void ZedNative::enable(sl::Camera& camera) {
    if (config_.tracking) {
        sl::PositionalTrackingParameters tracking_params;
        sl::ERROR_CODE err = camera.enablePositionalTracking(tracking_params);
        if (err != sl::ERROR_CODE::SUCCESS) {
            std::cerr << "Failed to enable positional tracking: "
                      << sl::toString(err) << std::endl;
        }
    }

    std::cout << "Object Detection: loading model (" << config_.model << ")..." << std::endl;
    sl::ERROR_CODE err = camera.enableObjectDetection(det_params_);
    if (err != sl::ERROR_CODE::SUCCESS) {
        std::cerr << "Failed to enable object detection: "
                  << sl::toString(err) << std::endl;
        return;
    }
    std::cout << "Object Detection: model loaded." << std::endl;
}

void ZedNative::detect(sl::Camera& camera) {
    detected_.clear();

    sl::ERROR_CODE err = camera.retrieveObjects(raw_objects_, det_rt_params_);
    if (err != sl::ERROR_CODE::SUCCESS) {
        return;
    }

    for (const auto& obj : raw_objects_.object_list) {
        DetectedObject det{};
        det.label_id = static_cast<int>(obj.label);
        det.confidence = obj.confidence;

        det.position[0] = obj.position.x;
        det.position[1] = obj.position.y;
        det.position[2] = obj.position.z;

        for (int i = 0; i < 4 && i < static_cast<int>(obj.bounding_box_2d.size()); i++) {
            det.bounding_box_2d[i * 2 + 0] = obj.bounding_box_2d[i].x;
            det.bounding_box_2d[i * 2 + 1] = obj.bounding_box_2d[i].y;
        }

        for (int i = 0; i < 8 && i < static_cast<int>(obj.bounding_box.size()); i++) {
            det.bounding_box_3d[i * 3 + 0] = obj.bounding_box[i].x;
            det.bounding_box_3d[i * 3 + 1] = obj.bounding_box[i].y;
            det.bounding_box_3d[i * 3 + 2] = obj.bounding_box[i].z;
        }

        detected_.push_back(det);
    }
}

const std::vector<DetectedObject>& ZedNative::objects() const {
    return detected_;
}

std::string ZedNative::label_names() const {
    return "PERSON,VEHICLE,BAG,ANIMAL,ELECTRONICS,FRUIT_VEGETABLE,SPORT";
}

} // namespace engines