#include "pointcloud_map.h"
#include "sc_map.h"
#include <mutex>
#include <rcl/event.h>
#include <rclcpp/node.hpp>
#include <thread>
#include <vector>
#include <queue>
#include <memory>
#include <iostream>
#include <filesystem>
#include <chrono>
// #include <filesystem>
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <livox_ros_driver2/msg/custom_msg.hpp>

#ifdef USE_RERUN
#include <rerun.hpp>
#include "rerun/collection_adapter.h"
#endif

#include "lidar_utils/utils.h"
#include "map_builder/commons.h"
#include "file_system/file_path.h"
#include "map_builder/map_builder.h"
#include "relocalizer.h"
#include "icp_localizer.h"
#include "map/map_manager.h"

#include "sophus/se3.hpp"

#include <pcl_conversions/pcl_conversions.h>
#include "tf2_ros/transform_broadcaster.h"
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <nav_msgs/msg/path.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <yaml-cpp/yaml.h>

#include <iostream>
#include <memory>
#include <string>

#include <glog/logging.h>

#include "rclcpp/rclcpp.hpp"
#include "rosbag2_cpp/reader.hpp"
#include "sensor_msgs/msg/image.hpp"

using namespace std;

struct NodeConfig
{
    std::string imu_topic = "/livox/imu";
    std::string lidar_topic = "/livox/lidar";
    std::string body_frame = "body";
    std::string world_frame = "lidar";
    bool print_time_cost = false;
};
struct StateData
{
    bool lidar_pushed = false;
    std::mutex imu_mutex;
    std::mutex lidar_mutex;
    double last_lidar_time = -1.0;
    double last_imu_time = -1.0;
    std::deque<IMUData> imu_buffer;
    std::deque<std::pair<double, pcl::PointCloud<pcl::PointXYZINormal>::Ptr>> lidar_buffer;
    nav_msgs::msg::Path path;
};

class LIONodeInference
{

public:
    LIONodeInference()
    {
        std::string config_file_path(
            "/home/xiarui/catkin_ws/src/FASTLIO2_ROS2/fastlio2/config/lio_YJ.yaml");
        loadParameters(config_file_path);

        m_state_data.path.poses.clear();
        m_state_data.path.header.frame_id = m_node_config.world_frame;

        m_kf = std::make_shared<IESKF>();
        m_builder = std::make_shared<MapBuilder>(m_builder_config, m_kf);
        // m_timer = this->create_wall_timer(20ms, std::bind(&LIONode::timerCB, this));
    }

    void loadParameters(const std::string &config_path)
    {

        YAML::Node config = YAML::LoadFile(config_path);

        m_node_config.imu_topic = config["imu_topic"].as<std::string>();
        m_node_config.lidar_topic = config["lidar_topic"].as<std::string>();
        m_node_config.body_frame = config["body_frame"].as<std::string>();
        m_node_config.world_frame = config["world_frame"].as<std::string>();
        m_node_config.print_time_cost = config["print_time_cost"].as<bool>();

        m_builder_config.lidar_filter_num = config["lidar_filter_num"].as<int>();
        m_builder_config.lidar_min_range = config["lidar_min_range"].as<double>();
        m_builder_config.lidar_max_range = config["lidar_max_range"].as<double>();
        m_builder_config.scan_resolution = config["scan_resolution"].as<double>();
        m_builder_config.map_resolution = config["map_resolution"].as<double>();
        m_builder_config.cube_len = config["cube_len"].as<double>();
        m_builder_config.det_range = config["det_range"].as<double>();
        m_builder_config.move_thresh = config["move_thresh"].as<double>();
        m_builder_config.na = config["na"].as<double>();
        m_builder_config.ng = config["ng"].as<double>();
        m_builder_config.nba = config["nba"].as<double>();
        m_builder_config.nbg = config["nbg"].as<double>();

        m_builder_config.imu_init_num = config["imu_init_num"].as<int>();
        m_builder_config.near_search_num = config["near_search_num"].as<int>();
        m_builder_config.ieskf_max_iter = config["ieskf_max_iter"].as<int>();
        m_builder_config.gravity_align = config["gravity_align"].as<bool>();
        m_builder_config.esti_il = config["esti_il"].as<bool>();
        std::vector<double> t_il_vec = config["t_il"].as<std::vector<double>>();
        std::vector<double> r_il_vec = config["r_il"].as<std::vector<double>>();
        m_builder_config.t_il << t_il_vec[0], t_il_vec[1], t_il_vec[2];
        m_builder_config.r_il << r_il_vec[0], r_il_vec[1], r_il_vec[2], r_il_vec[3], r_il_vec[4], r_il_vec[5], r_il_vec[6], r_il_vec[7], r_il_vec[8];
        m_builder_config.lidar_cov_inv = config["lidar_cov_inv"].as<double>();
    }

    bool syncPackage()
    {
        if (m_state_data.imu_buffer.empty() || m_state_data.lidar_buffer.empty())
            return false;
        if (!m_state_data.lidar_pushed)
        {
            m_package.cloud = m_state_data.lidar_buffer.front().second;
            std::sort(m_package.cloud->points.begin(), m_package.cloud->points.end(), [](PointType &p1, PointType &p2)
                      { return p1.curvature < p2.curvature; });
            m_package.cloud_start_time = m_state_data.lidar_buffer.front().first;
            m_package.cloud_end_time = m_package.cloud_start_time + m_package.cloud->points.back().curvature / 1000.0;
            m_state_data.lidar_pushed = true;
        }
        if (m_state_data.last_imu_time < m_package.cloud_end_time)
            return false;

        Vec<IMUData>().swap(m_package.imus);
        while (!m_state_data.imu_buffer.empty() && m_state_data.imu_buffer.front().time < m_package.cloud_end_time)
        {
            m_package.imus.emplace_back(m_state_data.imu_buffer.front());
            m_state_data.imu_buffer.pop_front();
        }
        m_state_data.lidar_buffer.pop_front();
        m_state_data.lidar_pushed = false;
        return true;
    }

    // void publishPath(rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr path_pub, std::string frame_id, const double &time)
    // {
    //     if (path_pub->get_subscription_count() <= 0)
    //         return;
    //     geometry_msgs::msg::PoseStamped pose;
    //     pose.header.frame_id = frame_id;
    //     pose.header.stamp = Utils::getTime(time);
    //     pose.pose.position.x = m_kf->x().t_wi.x();
    //     pose.pose.position.y = m_kf->x().t_wi.y();
    //     pose.pose.position.z = m_kf->x().t_wi.z();
    //     Eigen::Quaterniond q(m_kf->x().r_wi);
    //     pose.pose.orientation.x = q.x();
    //     pose.pose.orientation.y = q.y();
    //     pose.pose.orientation.z = q.z();
    //     pose.pose.orientation.w = q.w();
    //     m_state_data.path.poses.push_back(pose);
    //     path_pub->publish(m_state_data.path);
    // }

public:
    StateData m_state_data;
    SyncPackage m_package;
    NodeConfig m_node_config;
    Config m_builder_config;
    std::shared_ptr<IESKF> m_kf;
    std::shared_ptr<MapBuilder> m_builder;
};

int main(int argc, char **argv)
{

    if (argc != 2)
    {
        cerr << "Usage: " << argv[0] << " <path_to_bag>" << endl;
        return 1;
    }
    std::string log_dir = "./log";
    google::SetLogDestination(google::GLOG_INFO, (log_dir + "/").c_str());
    google::SetLogDestination(google::GLOG_WARNING, (log_dir + "/").c_str());
    google::SetLogDestination(google::GLOG_ERROR, (log_dir + "/").c_str());
    google::SetLogDestination(google::GLOG_FATAL, (log_dir + "/").c_str());

    // 3. 初始化 Glog（必须在 SetLogDestination 之后）
    google::InitGoogleLogging("loc");

    // 可选：关闭日志文件的 stderr 输出（只写文件）
    FLAGS_logtostderr = 1; // 不打印到终端

    // 可选：设置日志文件名前缀（可选，便于区分）
    FLAGS_log_prefix = true;

    LOG(INFO) << "LIO Node Inference Started";

    rclcpp::init(argc, argv);
    auto node = rclcpp::Node::make_shared("lidar_mapping");

    rosbag2_cpp::Reader reader;
    rosbag2_storage::StorageOptions storage_options;
    storage_options.uri = argv[1];
    storage_options.storage_id = "sqlite3"; // 或 "mcap"

    rosbag2_cpp::ConverterOptions converter_options;
    converter_options.input_serialization_format = "cdr";
    converter_options.output_serialization_format = "cdr";

    try
    {
        reader.open(storage_options, converter_options);
    }
    catch (const exception &e)
    {
        cerr << "Open failed: " << e.what() << endl;
        return 1;
    }

    LOG(INFO) << "Bag size: " << reader.get_metadata().bag_size;
    auto metadata = reader.get_metadata();
    LOG(INFO) << "Number of messages: " << metadata.message_count << std::endl;
    LOG(INFO) << "Topics:" << std::endl;
    for (const auto &topic : metadata.topics_with_message_count)
    {
        LOG(INFO) << " - " << topic.topic_metadata.name << " (" << topic.topic_metadata.type << ") : " << topic.message_count << " msgs" << std::endl;
    }

    LIONodeInference LIO_node_inference;

    std::shared_ptr<Map2::MapManager> map_manager = std::make_shared<Map2::MapManager>();

    map_manager->CreateOrGetLayer<Map2::PointCloudMap>();
    map_manager->CreateOrGetLayer<Map2::SCMap>();

    const std::string map_dir = "/home/xiarui/Dataset/YJ/2025-09-18/bag_20250919_050818/map/";
    map_manager->LoadAll(map_dir);

    std::shared_ptr<Dobot::Relocalizer> relocalizer =
        std::make_shared<Dobot::Relocalizer>(map_manager);
    auto pointcloud_map = map_manager->GetLayer<Map2::PointCloudMap>();
#ifdef USE_RERUN
    auto rec = rerun::RecordingStream("lio_localization");
    rec.spawn().exit_on_failure();
    rec.log_static("world", rerun::ViewCoordinates::RIGHT_HAND_Z_UP); // Set an up-axis
    std::vector<rerun::datatypes::Vec3D> rec_path;

    std::vector<rerun::Position3D> rec_map_points = rerun::to_rerun(*pointcloud_map->getCloud());
    rec.log("world/map", rerun::Points3D(rec_map_points));

    rec.log(
        "world/robot",
        rerun::Arrows3D::from_vectors({{0.3, 0.0, 0.0}, {0.0, 0.3, 0.0}, {0.0, 0.0, 0.3}}).with_colors({{255, 0, 0}, {0, 255, 0}, {0, 0, 255}}));

    rec.log(
        "world/origin",
        rerun::Arrows3D::from_vectors({{0.5, 0.0, 0.0}, {0.0, 0.5, 0.0}, {0.0, 0.0, 0.5}}).with_colors({{255, 0, 0}, {0, 255, 0}, {0, 0, 255}}));

#endif

    bool map_initialized = false;

    auto sc_map = map_manager->GetLayer<Map2::SCMap>();

    Dobot::ICPConfig icp_config;
    auto icp_localizer = std::make_shared<Dobot::ICPLocalizer>(icp_config);

    icp_localizer->setMap(pointcloud_map->getCloud());

    Sophus::SE3<double> cur_pose;

    while (reader.has_next() && rclcpp::ok())
    {
        auto msg = reader.read_next();

        if (msg->topic_name.find("imu") != std::string::npos)
        {
            rclcpp::SerializedMessage serialized(*msg->serialized_data);
            auto imu = make_shared<sensor_msgs::msg::Imu>();
            rclcpp::Serialization<sensor_msgs::msg::Imu> serialization;
            serialization.deserialize_message(&serialized, imu.get());

            std::lock_guard<std::mutex> lock(LIO_node_inference.m_state_data.imu_mutex);
            double timestamp = Utils::getSec(imu->header);
            if (timestamp < LIO_node_inference.m_state_data.last_imu_time)
            {
                std::deque<IMUData>().swap(LIO_node_inference.m_state_data.imu_buffer);
            }
            LIO_node_inference.m_state_data.imu_buffer.emplace_back(V3D(imu->linear_acceleration.x, imu->linear_acceleration.y, imu->linear_acceleration.z) * 10.0,
                                                                    V3D(imu->angular_velocity.x, imu->angular_velocity.y, imu->angular_velocity.z),
                                                                    timestamp);
            LIO_node_inference.m_state_data.last_imu_time = timestamp;
        }
        else if (msg->topic_name.find("pointcloud") != std::string::npos)
        {

            rclcpp::SerializedMessage serialized(*msg->serialized_data);
            auto cloud = make_shared<sensor_msgs::msg::PointCloud2>();
            rclcpp::Serialization<sensor_msgs::msg::PointCloud2> serialization;
            serialization.deserialize_message(&serialized, cloud.get());
            // std::cout << "Cloud: " << cloud->header.stamp.sec << "." << cloud->header.stamp.nanosec << std::endl;
            auto cloud_normal = Utils::convertToXYZINormal(cloud, 100, 0.1, 100.0);

            std::lock_guard<std::mutex> lock(LIO_node_inference.m_state_data.lidar_mutex);
            double timestamp = Utils::getSec(cloud->header);
            if (timestamp < LIO_node_inference.m_state_data.last_lidar_time)
            {
                std::deque<std::pair<double, pcl::PointCloud<pcl::PointXYZINormal>::Ptr>>().swap(LIO_node_inference.m_state_data.lidar_buffer);
            }
            LIO_node_inference.m_state_data.lidar_buffer.emplace_back(timestamp, cloud_normal);
            LIO_node_inference.m_state_data.last_lidar_time = timestamp;

            bool sync_res = LIO_node_inference.syncPackage();

            if (!sync_res)
            {
                continue;
            }

            // LIO_node_inference.m_builder->process(LIO_node_inference.m_package);

            // Eigen::Vector3d lidar_t = LIO_node_inference.m_builder->lidar_processor()->t_wl();
            // Eigen::Matrix3d lidar_r = LIO_node_inference.m_builder->lidar_processor()->r_wl();

            // CloudType::Ptr world_cloud =
            //     LIO_node_inference.m_builder->lidar_processor()->transformCloud(
            //         LIO_node_inference.m_package.cloud,
            //         LIO_node_inference.m_builder->lidar_processor()->r_wl(),
            //         LIO_node_inference.m_builder->lidar_processor()->t_wl());

            if (map_initialized)
            {

                icp_localizer->setInput(LIO_node_inference.m_package.cloud);

                Eigen::Matrix4f pose_f = cur_pose.matrix().cast<float>();

                icp_localizer->align(pose_f);

                Eigen::Matrix4d pose_d = pose_f.cast<double>();

                auto t = pose_d.block<3, 1>(0, 3);
                Eigen::Quaterniond q = Eigen::Quaterniond(
                    pose_d.block<3, 3>(0, 0));

                cur_pose = Sophus::SE3d(q.normalized().toRotationMatrix(), t);

                LOG(INFO) << "cur pose: " << cur_pose.translation().transpose();
            }
            else
            {

                LOG(INFO) << "Querying SCMap";

                // sc_map->QueryMap(*cloud_normal);
                bool reloc_result = relocalizer->Reloc(cloud_normal);

                Dobot::Relocalizer::RelocResultInfo reloc_result_info =
                    relocalizer->GetRelocResult();

                if (reloc_result_info.valid)
                {
                    // cur_pose = reloc_result_info.pose;

                    LOG(INFO) << reloc_result_info.pose;
                    auto t = reloc_result_info.pose.block<3, 1>(0, 3);
                    Eigen::Quaterniond q = Eigen::Quaterniond(
                        reloc_result_info.pose.block<3, 3>(0, 0));

                    cur_pose = Sophus::SE3d(q.normalized().toRotationMatrix(), t);
                }

                map_initialized = true;
            }

#ifdef USE_RERUN

            if (map_initialized)
            {

                rec.log("world/robot",
                        rerun::Transform3D(
                            rerun::to_rerun(cur_pose.translation()),
                            rerun::to_rerun(cur_pose.rotationMatrix())));

                CloudType::Ptr world_cloud =
                    LIO_node_inference.m_builder->lidar_processor()
                        ->transformCloud(LIO_node_inference.m_package.cloud,
                                         cur_pose.rotationMatrix(),
                                         cur_pose.translation());
                std::vector<rerun::Position3D> rec_points =
                    rerun::to_rerun(*world_cloud);

                rec_path.push_back(rerun::to_rerun(cur_pose.translation()));

                rec.log("world/path",
                        rerun::LineStrips3D(rerun::LineStrip3D(rec_path)));
                rec.log("world/lidar_raw", rerun::Points3D(rec_points));
            }

#else

#endif
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }

    return 0;
}
