#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <cstring>
#include <mutex>
#include <regex>
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "image_transport/publisher_plugin.hpp"
#include "pluginlib/class_loader.hpp"
#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/msg/image.hpp"
#include "sensor_msgs/msg/point_cloud2.hpp"
#include "sensor_msgs/msg/point_field.hpp"

using Image = sensor_msgs::msg::Image;
using PointCloud2 = sensor_msgs::msg::PointCloud2;

class Adapter final : public rclcpp::Node {
public:
  Adapter() : Node("ros_viz_adapter") {
    output_ns_ = declare_parameter<std::string>("output_namespace", "/viz");
    input_fps_ = std::max(0.0, declare_parameter<double>("input_fps_limit", 30.0));
    output_fps_ = std::max(0.0, declare_parameter<double>("output_fps", 10.0));
    voxel_ = std::max(0.0, declare_parameter<double>("voxel_size_m", 0.05));
    max_points_ = std::max<int>(1, declare_parameter<int>("max_points", 100000));
    qos_ = rclcpp::SensorDataQoS().keep_last(1);
    discover();
    timer_ = create_wall_timer(std::chrono::seconds(2), [this] { discover(); });
  }

  ~Adapter() override {
    stop_ = true; wake_.notify_all();
    if (rgb_thread_.joinable()) rgb_thread_.join();
    if (depth_thread_.joinable()) depth_thread_.join();
    if (cloud_thread_.joinable()) cloud_thread_.join();
  }

private:
  enum class ImageKind { RGB, DEPTH, UNKNOWN };

  static std::string safe_name(const std::string & topic) {
    auto name = std::regex_replace(topic, std::regex("[^A-Za-z0-9_]+"), "_");
    while (!name.empty() && name.front() == '_') name.erase(name.begin());
    while (!name.empty() && name.back() == '_') name.pop_back();
    return name.empty() ? "stream" : name;
  }

  static ImageKind classify(const std::string & encoding) {
    std::string e = encoding;
    std::transform(e.begin(), e.end(), e.begin(), [](unsigned char c) { return std::tolower(c); });
    if (e == "16uc1" || e == "32fc1" || e == "mono16" || e == "32sc1") return ImageKind::DEPTH;
    if (e == "rgb8" || e == "bgr8" || e == "rgba8" || e == "bgra8" ||
        e == "mono8" || e == "8uc3" || e == "8uc4" || e == "r8g8b8" ||
        e == "b8g8r8" || e == "rgb_int8" || e == "rgba_int8" ||
        e == "bgr_int8" || e == "bgra_int8") return ImageKind::RGB;
    return ImageKind::UNKNOWN;
  }

  bool due(const std::string & key, double fps, bool output) {
    if (fps <= 0.0) return true;
    std::lock_guard<std::mutex> lock(rate_mutex_);
    auto & table = output ? output_times_ : input_times_;
    const auto now = std::chrono::steady_clock::now();
    auto it = table.find(key);
    if (it != table.end() && std::chrono::duration<double>(now - it->second).count() < 1.0 / fps) return false;
    table[key] = now; return true;
  }

  void discover() {
    for (const auto & entry : get_topic_names_and_types()) {
      const auto & topic = entry.first;
      if (topic == output_ns_ || topic.rfind(output_ns_ + "/", 0) == 0) continue;
      if (subscriptions_.count(topic) || cloud_subscriptions_.count(topic)) continue;
      const auto & types = entry.second;
      if (std::find(types.begin(), types.end(), "sensor_msgs/msg/Image") != types.end()) {
        RCLCPP_INFO(get_logger(), "Discovered image topic: %s", topic.c_str());
        subscriptions_[topic] = create_subscription<Image>(topic, qos_, [this, topic](Image::ConstSharedPtr msg) { on_image(topic, msg); });
      } else if (std::find(types.begin(), types.end(), "sensor_msgs/msg/PointCloud2") != types.end()) {
        cloud_subscriptions_[topic] = create_subscription<PointCloud2>(topic, qos_, [this, topic](PointCloud2::ConstSharedPtr msg) { on_cloud(topic, msg); });
      }
    }
  }

  void on_image(const std::string & topic, Image::ConstSharedPtr msg) {
    if (!due(topic, input_fps_, false)) return;
    const auto kind = classify(msg->encoding);
    if (kind == ImageKind::UNKNOWN) {
      RCLCPP_WARN_THROTTLE(get_logger(), *get_clock(), 5000, "Ignoring image topic %s with encoding '%s'", topic.c_str(), msg->encoding.c_str());
      return;
    }
    RCLCPP_INFO_ONCE(get_logger(), "Classified image topic %s as %s (encoding=%s)", topic.c_str(), kind == ImageKind::RGB ? "RGB" : "DEPTH", msg->encoding.c_str());
    std::lock_guard<std::mutex> lock(data_mutex_);
    auto & slot = kind == ImageKind::RGB ? rgb_latest_ : depth_latest_;
    slot[topic] = msg;
    if (kind == ImageKind::RGB && !rgb_thread_.joinable()) rgb_thread_ = std::thread([this] { image_loop(rgb_latest_, false); });
    if (kind == ImageKind::DEPTH && !depth_thread_.joinable()) depth_thread_ = std::thread([this] { image_loop(depth_latest_, true); });
    wake_.notify_all();
  }

  void on_cloud(const std::string & topic, PointCloud2::ConstSharedPtr msg) {
    if (!due(topic, input_fps_, false)) return;
    std::lock_guard<std::mutex> lock(data_mutex_);
    cloud_latest_[topic] = msg;
    if (!cloud_thread_.joinable()) cloud_thread_ = std::thread([this] { cloud_loop(); });
    wake_.notify_all();
  }

  std::string image_topic(const std::string & topic, bool depth) const {
    // Publisher plugins append their own transport suffix.
    return (output_ns_.empty() ? "/viz" : output_ns_) + "/" + safe_name(topic) + (depth ? "/depth" : "/rgb");
  }

  void image_loop(std::unordered_map<std::string, Image::ConstSharedPtr> & slots, bool depth) {
    while (!stop_) {
      std::unordered_map<std::string, Image::ConstSharedPtr> pending;
      {
        std::unique_lock<std::mutex> lock(data_mutex_);
        wake_.wait_for(lock, std::chrono::milliseconds(100), [this, &slots] { return stop_ || !slots.empty(); });
        pending.swap(slots);
      }
      for (const auto & item : pending) {
        const auto key = item.first + (depth ? ":depth" : ":rgb");
        if (!due(key, output_fps_, true)) continue;
        std::lock_guard<std::mutex> lock(pub_mutex_);
        auto it = image_publishers_.find(key);
        if (it == image_publishers_.end()) {
          const auto lookup = depth ? "image_transport/compressedDepth_pub" : "image_transport/compressed_pub";
          try {
            auto plugin = image_loader_->createSharedInstance(lookup);
            plugin->advertise(this, image_topic(item.first, depth), qos_.get_rmw_qos_profile());
            RCLCPP_INFO(get_logger(), "Loaded %s -> %s", lookup, plugin->getTopic().c_str());
            it = image_publishers_.emplace(key, std::move(plugin)).first;
          } catch (const pluginlib::PluginlibException & ex) {
            RCLCPP_ERROR(get_logger(), "Cannot load image_transport plugin %s for %s: %s", lookup, item.first.c_str(), ex.what());
            continue;
          }
        }
        it->second->publishPtr(item.second);
      }
    }
  }

  void cloud_loop() {
    while (!stop_) {
      std::unordered_map<std::string, PointCloud2::ConstSharedPtr> pending;
      {
        std::unique_lock<std::mutex> lock(data_mutex_);
        wake_.wait_for(lock, std::chrono::milliseconds(100), [this] { return stop_ || !cloud_latest_.empty(); });
        pending.swap(cloud_latest_);
      }
      for (const auto & item : pending) {
        if (!due(item.first + ":points", output_fps_, true)) continue;
        auto output = downsample(*item.second);
        std::lock_guard<std::mutex> lock(pub_mutex_);
        auto it = cloud_publishers_.find(item.first);
        if (it == cloud_publishers_.end()) {
          it = cloud_publishers_.emplace(item.first, create_publisher<PointCloud2>(output_ns_ + "/" + safe_name(item.first) + "/points", qos_)).first;
        }
        it->second->publish(output);
      }
    }
  }

  PointCloud2 downsample(const PointCloud2 & input) const {
    int xoff = -1, yoff = -1, zoff = -1;
    for (const auto & field : input.fields) {
      if (field.name == "x") xoff = field.offset;
      if (field.name == "y") yoff = field.offset;
      if (field.name == "z") zoff = field.offset;
    }
    PointCloud2 output; output.header = input.header; output.height = 1; output.point_step = 12; output.is_dense = true;
    sensor_msgs::msg::PointField fx; fx.name = "x"; fx.offset = 0; fx.datatype = sensor_msgs::msg::PointField::FLOAT32; fx.count = 1;
    sensor_msgs::msg::PointField fy; fy.name = "y"; fy.offset = 4; fy.datatype = sensor_msgs::msg::PointField::FLOAT32; fy.count = 1;
    sensor_msgs::msg::PointField fz; fz.name = "z"; fz.offset = 8; fz.datatype = sensor_msgs::msg::PointField::FLOAT32; fz.count = 1;
    output.fields = {fx, fy, fz};
    if (xoff < 0 || yoff < 0 || zoff < 0 || input.point_step == 0) return output;
    std::unordered_set<std::string> voxels;
    for (size_t i = 0; i < input.width * input.height && output.width < static_cast<uint32_t>(max_points_); ++i) {
      const auto * raw = input.data.data() + i * input.point_step; float xyz[3];
      std::memcpy(&xyz[0], raw + xoff, 4); std::memcpy(&xyz[1], raw + yoff, 4); std::memcpy(&xyz[2], raw + zoff, 4);
      if (!std::isfinite(xyz[0]) || !std::isfinite(xyz[1]) || !std::isfinite(xyz[2])) continue;
      if (voxel_ > 0.0) { const auto key = std::to_string(static_cast<long long>(std::floor(xyz[0] / voxel_))) + ":" + std::to_string(static_cast<long long>(std::floor(xyz[1] / voxel_))) + ":" + std::to_string(static_cast<long long>(std::floor(xyz[2] / voxel_))); if (!voxels.insert(key).second) continue; }
      const auto old = output.data.size(); output.data.resize(old + 12); std::memcpy(output.data.data() + old, xyz, 12); ++output.width;
    }
    output.row_step = output.width * output.point_step; return output;
  }

  std::string output_ns_; double input_fps_{30.0}, output_fps_{10.0}, voxel_{0.05}; int max_points_{100000};
  std::shared_ptr<pluginlib::ClassLoader<image_transport::PublisherPlugin>> image_loader_{
    std::make_shared<pluginlib::ClassLoader<image_transport::PublisherPlugin>>("image_transport", "image_transport::PublisherPlugin")};
  rclcpp::QoS qos_{rclcpp::SensorDataQoS()}; rclcpp::TimerBase::SharedPtr timer_;
  std::atomic_bool stop_{false}; std::mutex data_mutex_, rate_mutex_, pub_mutex_; std::condition_variable wake_;
  std::thread rgb_thread_, depth_thread_, cloud_thread_;
  std::unordered_map<std::string, rclcpp::Subscription<Image>::SharedPtr> subscriptions_; std::unordered_map<std::string, rclcpp::Subscription<PointCloud2>::SharedPtr> cloud_subscriptions_;
  std::unordered_map<std::string, Image::ConstSharedPtr> rgb_latest_, depth_latest_; std::unordered_map<std::string, PointCloud2::ConstSharedPtr> cloud_latest_;
  std::unordered_map<std::string, std::chrono::steady_clock::time_point> input_times_, output_times_;
  std::unordered_map<std::string, std::shared_ptr<image_transport::PublisherPlugin>> image_publishers_; std::unordered_map<std::string, rclcpp::Publisher<PointCloud2>::SharedPtr> cloud_publishers_;
};

int main(int argc, char ** argv) { rclcpp::init(argc, argv); auto node = std::make_shared<Adapter>(); rclcpp::spin(node); rclcpp::shutdown(); return 0; }