#include "engines/open_vocab.h"
#include "util/clip_tokenizer.h"
#include "util/env.h"

#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <sl/Camera.hpp>

#include <algorithm>
#include <cmath>
#include <iostream>
#include <numeric>
#include <sstream>
#include <string>
#include <vector>

namespace {

// --- NMS helpers ---

struct RawDetection {
    float x1, y1, x2, y2;
    float score;
    int class_id;
};

float iou(const RawDetection& a, const RawDetection& b) {
    float inter_x1 = std::max(a.x1, b.x1);
    float inter_y1 = std::max(a.y1, b.y1);
    float inter_x2 = std::min(a.x2, b.x2);
    float inter_y2 = std::min(a.y2, b.y2);
    float inter_w = std::max(0.0f, inter_x2 - inter_x1);
    float inter_h = std::max(0.0f, inter_y2 - inter_y1);
    float inter_area = inter_w * inter_h;
    float area_a = (a.x2 - a.x1) * (a.y2 - a.y1);
    float area_b = (b.x2 - b.x1) * (b.y2 - b.y1);
    float union_area = area_a + area_b - inter_area;
    return union_area > 0.0f ? inter_area / union_area : 0.0f;
}

std::vector<RawDetection> nms(std::vector<RawDetection> dets, float threshold) {
    std::sort(dets.begin(), dets.end(),
              [](const RawDetection& a, const RawDetection& b) {
                  return a.score > b.score;
              });
    std::vector<bool> suppressed(dets.size(), false);
    std::vector<RawDetection> result;
    for (size_t i = 0; i < dets.size(); i++) {
        if (suppressed[i]) continue;
        result.push_back(dets[i]);
        for (size_t j = i + 1; j < dets.size(); j++) {
            if (!suppressed[j] && iou(dets[i], dets[j]) > threshold) {
                suppressed[j] = true;
            }
        }
    }
    return result;
}

// --- Letterbox resize ---
cv::Mat letterbox(const cv::Mat& src, int size, float& scale, float& pad_x, float& pad_y) {
    float scale_w = static_cast<float>(size) / src.cols;
    float scale_h = static_cast<float>(size) / src.rows;
    scale = std::min(scale_w, scale_h);
    int new_w = static_cast<int>(std::round(src.cols * scale));
    int new_h = static_cast<int>(std::round(src.rows * scale));
    pad_x = (size - new_w) / 2.0f;
    pad_y = (size - new_h) / 2.0f;

    cv::Mat resized;
    cv::resize(src, resized, {new_w, new_h});

    cv::Mat canvas(size, size, CV_8UC3, cv::Scalar(114, 114, 114));
    resized.copyTo(canvas(cv::Rect(
        static_cast<int>(pad_x), static_cast<int>(pad_y), new_w, new_h)));
    return canvas;
}

// --- Point cloud fusion ---
struct FusionResult {
    float cx, cy, cz;
    float depth_near, depth_far;
    bool valid = false;
};

FusionResult fuse_bbox(const sl::Mat& pc, int x1, int y1, int x2, int y2) {
    FusionResult result{};
    int w = static_cast<int>(pc.getWidth());
    int h = static_cast<int>(pc.getHeight());
    x1 = std::max(x1, 0);
    y1 = std::max(y1, 0);
    x2 = std::min(x2, w - 1);
    y2 = std::min(y2, h - 1);

    std::vector<float> xs, ys, zs;
    xs.reserve((x2 - x1) * (y2 - y1));
    ys.reserve(xs.capacity());
    zs.reserve(xs.capacity());

    for (int y = y1; y <= y2; y++) {
        for (int x = x1; x <= x2; x++) {
            sl::float4 pt;
            pc.getValue(x, y, &pt);
            if (!std::isnan(pt.x) && !std::isnan(pt.y) && !std::isnan(pt.z) &&
                !std::isinf(pt.x) && !std::isinf(pt.y) && !std::isinf(pt.z)) {
                xs.push_back(pt.x);
                ys.push_back(pt.y);
                zs.push_back(pt.z);
            }
        }
    }

    if (xs.empty()) return result;

    std::sort(xs.begin(), xs.end());
    std::sort(ys.begin(), ys.end());
    std::sort(zs.begin(), zs.end());

    size_t n   = xs.size();
    size_t mid = n / 2;
    size_t p25 = n / 4;
    size_t p75 = 3 * n / 4;

    result.cx         = xs[mid];
    result.cy         = ys[mid];
    result.cz         = zs[mid];
    result.depth_near = zs[p25];
    result.depth_far  = zs[p75];
    result.valid      = true;
    return result;
}

void backproject(float u, float v, float z, float fx, float fy, float cx, float cy,
                 float& X, float& Y, float& Z) {
    Z = z;
    X = (u - cx) * z / fx;
    Y = (v - cy) * z / fy;
}

// Split a comma-separated string into a vector of trimmed tokens
std::vector<std::string> split_labels(const std::string& s) {
    std::vector<std::string> out;
    std::stringstream ss(s);
    std::string tok;
    while (std::getline(ss, tok, ',')) {
        size_t l = tok.find_first_not_of(" \t");
        size_t r = tok.find_last_not_of(" \t");
        if (l != std::string::npos)
            out.push_back(tok.substr(l, r - l + 1));
    }
    return out;
}

} // namespace

namespace engines {

// --- Impl ---

struct OpenVocab::Impl {
    // Visual model
    Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "open_vocab"};
    Ort::SessionOptions session_options;
    std::unique_ptr<Ort::Session> session;
    Ort::AllocatorWithDefaultOptions allocator;

    std::string input_name;       // visual model: image input
    std::string output_name;      // visual model: detection output
    std::string text_input_name;  // visual model: text_feats input (empty if single-input model)
    int num_classes = 0;
    bool has_text_input = false;

    // Text encoder (optional; used when has_text_input == true)
    Ort::SessionOptions text_session_options;
    std::unique_ptr<Ort::Session> text_session;
    std::string text_enc_input_name;   // "input_ids"
    std::string text_enc_output_name;  // "text_embeds"

    // Tokenizer + embedding cache
    util::ClipTokenizer tokenizer;
    std::vector<float> cached_text_feats;  // (num_classes * 512) flat, row-major
    std::string cached_labels;

    // Preprocessing + ONNX inference + NMS.
    // Takes a BGR image (cv::imread format); returns 2D detections in original image coords.
    std::vector<RawDetection> run_inference(const OpenVocabConfig& config,
                                             const cv::Mat& image_bgr);
};

// --- Impl::run_inference ---

std::vector<RawDetection> OpenVocab::Impl::run_inference(
    const OpenVocabConfig& config, const cv::Mat& image_bgr)
{
    if (!session) return {};
    if (has_text_input && cached_text_feats.empty()) return {};

    int orig_w = image_bgr.cols;
    int orig_h = image_bgr.rows;

    // BGR → RGB (YOLO expects RGB)
    cv::Mat rgb;
    cv::cvtColor(image_bgr, rgb, cv::COLOR_BGR2RGB);

    float scale, pad_x, pad_y;
    cv::Mat lb = letterbox(rgb, config.input_size, scale, pad_x, pad_y);

    cv::Mat lb_float;
    lb.convertTo(lb_float, CV_32F, 1.0f / 255.0f);

    int sz = config.input_size;
    std::vector<float> input_data(3 * sz * sz);
    for (int c = 0; c < 3; c++) {
        for (int h = 0; h < sz; h++) {
            for (int w = 0; w < sz; w++) {
                input_data[c * sz * sz + h * sz + w] =
                    lb_float.at<cv::Vec3f>(h, w)[c];
            }
        }
    }

    auto mem_info = Ort::MemoryInfo::CreateCpu(
        OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);

    std::vector<int64_t> image_shape{1, 3, sz, sz};
    Ort::Value image_tensor = Ort::Value::CreateTensor<float>(
        mem_info, input_data.data(), input_data.size(),
        image_shape.data(), image_shape.size());

    const char* output_name_cstr = output_name.c_str();
    std::vector<Ort::Value> outputs;

    try {
        if (has_text_input) {
            std::vector<int64_t> text_shape{1, num_classes, 512};
            Ort::Value text_tensor = Ort::Value::CreateTensor<float>(
                mem_info,
                cached_text_feats.data(),
                cached_text_feats.size(),
                text_shape.data(), text_shape.size());

            const char* input_names[2] = {
                input_name.c_str(),
                text_input_name.c_str()
            };
            Ort::Value input_tensors[2] = {
                std::move(image_tensor),
                std::move(text_tensor)
            };
            outputs = session->Run(
                Ort::RunOptions{nullptr},
                input_names, input_tensors, 2,
                &output_name_cstr, 1);
        } else {
            const char* input_name_cstr = input_name.c_str();
            outputs = session->Run(
                Ort::RunOptions{nullptr},
                &input_name_cstr, &image_tensor, 1,
                &output_name_cstr, 1);
        }
    } catch (const Ort::Exception& e) {
        std::cerr << "OpenVocab: inference failed: " << e.what() << std::endl;
        return {};
    }

    // Post-process: output shape (1, 4+num_classes, num_anchors)
    auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
    if (out_shape.size() < 3) return {};

    int num_anchors = static_cast<int>(out_shape[2]);
    int nc = num_classes;
    if (nc <= 0) return {};

    const float* out = outputs[0].GetTensorData<float>();

    std::vector<RawDetection> candidates;
    for (int a = 0; a < num_anchors; a++) {
        float cx_lb = out[0 * num_anchors + a];
        float cy_lb = out[1 * num_anchors + a];
        float w_lb  = out[2 * num_anchors + a];
        float h_lb  = out[3 * num_anchors + a];

        float best_score = 0.0f;
        int   best_class = 0;
        for (int c = 0; c < nc; c++) {
            float s = out[(4 + c) * num_anchors + a];
            if (s > best_score) { best_score = s; best_class = c; }
        }
        if (best_score < config.confidence_threshold) continue;

        float x1_lb = cx_lb - w_lb / 2.0f;
        float y1_lb = cy_lb - h_lb / 2.0f;
        float x2_lb = cx_lb + w_lb / 2.0f;
        float y2_lb = cy_lb + h_lb / 2.0f;

        float x1 = std::max(0.0f, std::min((x1_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
        float y1 = std::max(0.0f, std::min((y1_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));
        float x2 = std::max(0.0f, std::min((x2_lb - pad_x) / scale, static_cast<float>(orig_w - 1)));
        float y2 = std::max(0.0f, std::min((y2_lb - pad_y) / scale, static_cast<float>(orig_h - 1)));

        candidates.push_back({x1, y1, x2, y2, best_score, best_class});
    }

    return nms(std::move(candidates), config.nms_threshold);
}

// --- Config ---

OpenVocabConfig OpenVocabConfig::from_env() {
    OpenVocabConfig c;
    c.model_path           = util::get_env_str("ZED_ONNX_MODEL_PATH",        c.model_path);
    c.confidence_threshold = util::get_env_float("ZED_ONNX_CONFIDENCE",      c.confidence_threshold);
    c.nms_threshold        = util::get_env_float("ZED_ONNX_NMS_THRESHOLD",   c.nms_threshold);
    c.input_size           = util::get_env_int("ZED_ONNX_INPUT_SIZE",        c.input_size);
    c.labels               = util::get_env_str("ZED_ONNX_LABELS",            c.labels);
    c.text_encoder_path    = util::get_env_str("ZED_ONNX_TEXT_ENCODER_PATH", c.text_encoder_path);
    c.clip_vocab_path      = util::get_env_str("ZED_ONNX_CLIP_VOCAB",        c.clip_vocab_path);
    c.clip_merges_path     = util::get_env_str("ZED_ONNX_CLIP_MERGES",       c.clip_merges_path);
    return c;
}

// --- OpenVocab ---

OpenVocab::OpenVocab(const OpenVocabConfig& config)
    : config_(config), impl_(std::make_unique<Impl>()) {}

OpenVocab::~OpenVocab() = default;

void OpenVocab::enable(sl::Camera& /* camera */) {
    if (config_.model_path.empty()) {
        std::cerr << "OpenVocab: ZED_ONNX_MODEL_PATH is not set." << std::endl;
        return;
    }

    // --- Load visual model ---
    OrtCUDAProviderOptions cuda_opts{};
    cuda_opts.device_id = 0;
    try {
        impl_->session_options.AppendExecutionProvider_CUDA(cuda_opts);
    } catch (const Ort::Exception& e) {
        std::cerr << "OpenVocab: CUDA unavailable (" << e.what() << "), using CPU." << std::endl;
    }
    impl_->session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    std::cout << "OpenVocab: loading visual model (" << config_.model_path << ")..." << std::endl;
    try {
        impl_->session = std::make_unique<Ort::Session>(
            impl_->env, config_.model_path.c_str(), impl_->session_options);
    } catch (const Ort::Exception& e) {
        std::cerr << "OpenVocab: failed to load visual model: " << e.what() << std::endl;
        return;
    }

    impl_->input_name  = impl_->session->GetInputNameAllocated(0, impl_->allocator).get();
    impl_->output_name = impl_->session->GetOutputNameAllocated(0, impl_->allocator).get();

    size_t num_inputs = impl_->session->GetInputCount();
    impl_->has_text_input = (num_inputs >= 2);
    if (impl_->has_text_input) {
        impl_->text_input_name = impl_->session->GetInputNameAllocated(1, impl_->allocator).get();
    }

    auto out_info = impl_->session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo();
    auto shape    = out_info.GetShape();
    if (shape.size() == 3 && shape[1] > 4) {
        impl_->num_classes = static_cast<int>(shape[1]) - 4;
    }

    std::cout << "OpenVocab: visual model loaded."
              << " Inputs=" << num_inputs
              << " Classes=" << impl_->num_classes
              << (impl_->has_text_input ? " (open-vocab)" : " (fixed-vocab)")
              << std::endl;

    // --- Load text encoder + tokenizer (only needed for open-vocab model) ---
    if (impl_->has_text_input) {
        if (config_.text_encoder_path.empty()) {
            std::cerr << "OpenVocab: ZED_ONNX_TEXT_ENCODER_PATH is not set "
                         "(required for open-vocab model)." << std::endl;
            return;
        }

        impl_->text_session_options.SetGraphOptimizationLevel(
            GraphOptimizationLevel::ORT_ENABLE_ALL);

        std::cout << "OpenVocab: loading text encoder (" << config_.text_encoder_path
                  << ")..." << std::endl;
        try {
            impl_->text_session = std::make_unique<Ort::Session>(
                impl_->env, config_.text_encoder_path.c_str(), impl_->text_session_options);
        } catch (const Ort::Exception& e) {
            std::cerr << "OpenVocab: failed to load text encoder: " << e.what() << std::endl;
            return;
        }

        impl_->text_enc_input_name  =
            impl_->text_session->GetInputNameAllocated(0, impl_->allocator).get();
        impl_->text_enc_output_name =
            impl_->text_session->GetOutputNameAllocated(0, impl_->allocator).get();

        if (!config_.clip_vocab_path.empty() && !config_.clip_merges_path.empty()) {
            try {
                impl_->tokenizer.load(config_.clip_vocab_path, config_.clip_merges_path);
                std::cout << "OpenVocab: tokenizer loaded." << std::endl;
            } catch (const std::exception& e) {
                std::cerr << "OpenVocab: failed to load tokenizer: " << e.what() << std::endl;
                return;
            }
        } else {
            std::cerr << "OpenVocab: ZED_ONNX_CLIP_VOCAB / ZED_ONNX_CLIP_MERGES not set." << std::endl;
            return;
        }

        if (!config_.labels.empty()) {
            encode_labels(config_.labels);
        }
    }
}

// Encode a comma-separated label string into cached_text_feats.
// Always fills exactly num_classes * 512 floats (pads with zeros if fewer labels).
void OpenVocab::encode_labels(const std::string& labels_str) {
    if (!impl_->text_session || !impl_->tokenizer.is_loaded()) return;

    auto label_vec = split_labels(labels_str);
    int N = impl_->num_classes;
    impl_->cached_text_feats.assign(static_cast<size_t>(N) * 512, 0.0f);

    auto mem_info = Ort::MemoryInfo::CreateCpu(
        OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault);

    int encode_n = std::min(static_cast<int>(label_vec.size()), N);
    for (int i = 0; i < encode_n; i++) {
        auto tokens = impl_->tokenizer.tokenize(label_vec[i]);

        std::vector<int64_t> input_shape{1, util::ClipTokenizer::SEQ_LEN};
        Ort::Value token_tensor = Ort::Value::CreateTensor<int64_t>(
            mem_info, tokens.data(), tokens.size(),
            input_shape.data(), input_shape.size());

        const char* in  = impl_->text_enc_input_name.c_str();
        const char* out = impl_->text_enc_output_name.c_str();
        std::vector<Ort::Value> outputs;
        try {
            outputs = impl_->text_session->Run(
                Ort::RunOptions{nullptr}, &in, &token_tensor, 1, &out, 1);
        } catch (const Ort::Exception& e) {
            std::cerr << "OpenVocab: text encoder failed for \"" << label_vec[i]
                      << "\": " << e.what() << std::endl;
            continue;
        }

        const float* emb = outputs[0].GetTensorData<float>();

        // L2 normalize
        float norm = 0.0f;
        for (int j = 0; j < 512; j++) norm += emb[j] * emb[j];
        norm = std::sqrt(norm);
        float inv_norm = (norm > 1e-8f) ? 1.0f / norm : 0.0f;
        for (int j = 0; j < 512; j++)
            impl_->cached_text_feats[i * 512 + j] = emb[j] * inv_norm;
    }

    impl_->cached_labels = labels_str;
    std::cout << "OpenVocab: encoded " << encode_n << " label(s): " << labels_str << std::endl;
}

// --- detect(sl::Camera&) — production path with 3D fusion ---

void OpenVocab::detect(sl::Camera& camera) {
    detected_.clear();
    if (!impl_->session) return;

    // Re-encode if labels changed (on-the-fly vocabulary switching)
    if (impl_->has_text_input) {
        const char* env_labels = std::getenv("ZED_ONNX_LABELS");
        std::string current_labels = env_labels ? env_labels : config_.labels;
        if (current_labels != impl_->cached_labels) {
            encode_labels(current_labels);
        }
        if (impl_->cached_text_feats.empty()) return;
    }

    // Retrieve image and point cloud from camera
    sl::Mat img_mat;
    camera.retrieveImage(img_mat, sl::VIEW::LEFT);

    sl::Mat pc_mat;
    camera.retrieveMeasure(pc_mat, sl::MEASURE::XYZRGBA);

    sl::CameraInformation info = camera.getCameraInformation();
    auto calib = info.camera_configuration.calibration_parameters.left_cam;
    float fx = calib.fx, fy = calib.fy, cx = calib.cx, cy = calib.cy;

    int orig_w = static_cast<int>(img_mat.getWidth());
    int orig_h = static_cast<int>(img_mat.getHeight());

    // RGBA → BGR for run_inference
    cv::Mat rgba(orig_h, orig_w, CV_8UC4, img_mat.getPtr<sl::uchar1>());
    cv::Mat bgr;
    cv::cvtColor(rgba, bgr, cv::COLOR_RGBA2BGR);

    auto kept = impl_->run_inference(config_, bgr);

    // 3D fusion
    for (const auto& det : kept) {
        FusionResult fusion = fuse_bbox(pc_mat,
            static_cast<int>(det.x1), static_cast<int>(det.y1),
            static_cast<int>(det.x2), static_cast<int>(det.y2));

        DetectedObject obj{};
        obj.label_id   = det.class_id;
        obj.confidence = det.score * 100.0f;

        if (fusion.valid) {
            obj.position[0] = fusion.cx;
            obj.position[1] = fusion.cy;
            obj.position[2] = fusion.cz;
        }

        obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
        obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
        obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
        obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;

        if (fusion.valid) {
            float corners_u[4] = {det.x1, det.x2, det.x2, det.x1};
            float corners_v[4] = {det.y1, det.y1, det.y2, det.y2};
            float depths[2]    = {fusion.depth_near, fusion.depth_far};
            for (int face = 0; face < 2; face++) {
                for (int c = 0; c < 4; c++) {
                    float X, Y, Z;
                    backproject(corners_u[c], corners_v[c], depths[face],
                                fx, fy, cx, cy, X, Y, Z);
                    int idx = (face * 4 + c) * 3;
                    obj.bounding_box_3d[idx + 0] = X;
                    obj.bounding_box_3d[idx + 1] = Y;
                    obj.bounding_box_3d[idx + 2] = Z;
                }
            }
        }

        detected_.push_back(obj);
    }
}

// --- detect(const cv::Mat&) — image-only path, no 3D fusion ---

void OpenVocab::detect(const cv::Mat& image_bgr) {
    detected_.clear();
    if (!impl_->session) return;

    // Re-encode if labels changed
    if (impl_->has_text_input) {
        const char* env_labels = std::getenv("ZED_ONNX_LABELS");
        std::string current_labels = env_labels ? env_labels : config_.labels;
        if (current_labels != impl_->cached_labels) {
            encode_labels(current_labels);
        }
        if (impl_->cached_text_feats.empty()) return;
    }

    auto kept = impl_->run_inference(config_, image_bgr);

    for (const auto& det : kept) {
        DetectedObject obj{};
        obj.label_id   = det.class_id;
        obj.confidence = det.score * 100.0f;
        // position stays zero (no point cloud available)
        obj.bounding_box_2d[0] = det.x1; obj.bounding_box_2d[1] = det.y1;
        obj.bounding_box_2d[2] = det.x2; obj.bounding_box_2d[3] = det.y1;
        obj.bounding_box_2d[4] = det.x2; obj.bounding_box_2d[5] = det.y2;
        obj.bounding_box_2d[6] = det.x1; obj.bounding_box_2d[7] = det.y2;
        // bounding_box_3d stays zero
        detected_.push_back(obj);
    }
}

void OpenVocab::set_labels(const std::string& labels) {
    config_.labels = labels;
    encode_labels(labels);
}

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

std::string OpenVocab::label_names() const {
    return impl_->cached_labels.empty() ? config_.labels : impl_->cached_labels;
}

} // namespace engines