#include "Frame.h"
#include "backend.h"
#include "sc_map.h"
#include <cstddef>
#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 "lidar_utils/utils.h"
#include "map_builder/commons.h"
#include "file_system/file_path.h"
#include "map_builder/map_builder.h"
#include "map/map_manager.h"
#include "Map.h"

#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(const std::string &config_file_path)
    {
        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);
    }

    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;
    }

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;
};

void publishCloud(
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr pub,
    CloudType::Ptr cloud, std::string frame_id, const double &time)
{
    if (pub->get_subscription_count() <= 0)
        return;
    sensor_msgs::msg::PointCloud2 cloud_msg;
    pcl::toROSMsg(*cloud, cloud_msg);
    cloud_msg.header.frame_id = frame_id;
    cloud_msg.header.stamp = Utils::getTime(time);
    pub->publish(cloud_msg);
}

void publishPath(rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr path_pub, std::string frame_id, const double &time, std::shared_ptr<IESKF> &m_kf, StateData &m_state_data)
{
    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);
    m_state_data.path.header.stamp = Utils::getTime(time);
    m_state_data.path.header.frame_id = "map";

    // std::cout << pose.pose.position.x << std::endl;

    path_pub->publish(m_state_data.path);
}

#ifdef USE_RERUN

void visualizeLoopPairs(
    rerun::RecordingStream &stream,
    const Dobot::Map &map,
    const std::vector<Dobot::LoopPair> &loop_pairs,
    const std::string &entity_path = "/world/loop_pairs")
{
    if (loop_pairs.empty())
        return;

    std::vector<rerun::Vec3D> origins;
    std::vector<rerun::Vec3D> vectors;
    std::vector<rerun::Rgba32> colors;

    for (const auto &loop : loop_pairs)
    {
        // 获取 source 和 target 帧
        auto source_frame = map.getFrame(loop.source_id);
        auto target_frame = map.getFrame(loop.target_id);
        if (!source_frame || !target_frame)
        {
            continue; // 跳过无效帧
        }

        Eigen::Matrix4d T_world_source = source_frame->getPose();
        Eigen::Matrix4d T_world_target = target_frame->getPose();

        Eigen::Vector3d source_pos = T_world_source.block<3, 1>(0, 3);
        Eigen::Vector3d target_pos = T_world_target.block<3, 1>(0, 3);

        origins.push_back({static_cast<float>(source_pos.x()),
                           static_cast<float>(source_pos.y()),
                           static_cast<float>(source_pos.z())});

        // 箭头方向 = target - source
        Eigen::Vector3d direction = target_pos - source_pos;
        vectors.push_back({static_cast<float>(direction.x()),
                           static_cast<float>(direction.y()),
                           static_cast<float>(direction.z())});

        // 颜色：根据 score 映射（假设 score 越大越好）
        // 这里简单用 green=score, red=1-score（需归一化）
        float normalized_score = std::clamp(static_cast<float>(loop.score), 0.0f, 1.0f);
        uint8_t r = static_cast<uint8_t>((1.0f - normalized_score) * 255);
        uint8_t g = static_cast<uint8_t>(normalized_score * 255);
        uint8_t b = 0;
        colors.push_back(rerun::Rgba32(r, g, b));
    }

    if (!origins.empty())
    {
        stream.log(entity_path,
                   rerun::Arrows3D::from_vectors(vectors)
                       .with_origins(origins)
                       .with_colors(colors) // 箭头粗细
        );
    }
}

#endif

inline bool isKeyFrame(const Dobot::Frame &last, const Dobot::Frame &cur,
                       double trans_thresh = 0.2,
                       double rot_thresh_deg = 90.0)
{
    // 平移距离
    Eigen::Vector3d t_last = last.getPose().block<3, 1>(0, 3);
    Eigen::Vector3d t_cur = cur.getPose().block<3, 1>(0, 3);
    double trans = (t_cur - t_last).norm();

    // 旋转角度
    Eigen::Matrix3d R_last = last.getPose().block<3, 3>(0, 0);
    Eigen::Matrix3d R_cur = cur.getPose().block<3, 3>(0, 0);
    double rot_rad = std::acos(std::min(
        1.0, std::max(-1.0, ((R_last.transpose() * R_cur).trace() - 1.0) / 2.0)));
    double rot_deg = rot_rad * 180.0 / M_PI;

    return trans > trans_thresh || rot_deg > rot_thresh_deg;
}

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

    if (argc != 3)
    {
        cerr << "Usage: " << argv[0] << " <path_to_bag> <config>" << 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");

    const std::string bag_path = argv[1];
    const std::string config_file_path(argv[2]);

    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;
    }

    LIONodeInference LIO_node_inference(config_file_path);

    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr m_body_cloud_pub;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr m_world_cloud_pub;
    rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr m_path_pub;
    rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr m_odom_pub;

    m_world_cloud_pub = node->create_publisher<sensor_msgs::msg::PointCloud2>("world_cloud", 10000);
    m_path_pub = node->create_publisher<nav_msgs::msg::Path>("lio_path", 10000);

    // LIO_node_inference.m_builder = std::make_shared<MapBuilder>(
    //     LIO_node_inference.m_builder_config, LIO_node_inference.m_kf);

    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;
    }

    // CloudType::Ptr map(new CloudType);

    Map2::MapManager map_manager;

    map_manager.CreateOrGetLayer<Dobot::Map>();
    map_manager.CreateOrGetLayer<Map2::SCMap>();
    auto map = map_manager.GetLayer<Dobot::Map>();
    auto sc_map = map_manager.GetLayer<Map2::SCMap>();

    auto back_end = std::make_shared<Dobot::BackEnd>(std::make_shared<Map2::MapManager>(map_manager));

    Dobot::Frame::Ptr last_frame = nullptr;

#ifdef USE_RERUN
    auto rec = rerun::RecordingStream("lio_mapping");
    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;
#endif

    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::cout << "IMU: " << imu->header.stamp.sec << "." << imu->header.stamp.nanosec << std::endl;

            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)
            {
                // RCLCPP_WARN(LIO_node_inference.get_logger(), "Lidar Message is out of order");
                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();

            // LOG(INFO) << LIO_node_inference.m_builder->lidar_processor()->r_wl();
            // LOG(INFO) << LIO_node_inference.m_builder->lidar_processor()->t_wl();

            Dobot::Frame::Ptr cur_frame = std::make_shared<Dobot::Frame>();
            cur_frame->setCloud(LIO_node_inference.m_package.cloud);
            cur_frame->setPose(lidar_r, lidar_t);

            // LOG(INFO) << "Frame id: " << cur_frame->getId();

            if (last_frame == nullptr || isKeyFrame(*last_frame, *cur_frame))
            {

                // LOG(INFO) << "Add frame to map";
                map->addFrame(cur_frame);
                sc_map->addFrame(*cur_frame->getCloud(), cur_frame->getPose(),
                                 cur_frame->getId());

                SCRelocInfo sc_reloc_info =
                    sc_map->detectLoopClosureID(*cur_frame->getCloud());

                if (sc_reloc_info.frame_id != -1)
                    LOG(INFO) << sc_reloc_info.toString();

                if (sc_reloc_info.frame_id != -1)
                {
                    Dobot::LoopPair loop_info =
                        Dobot::GetLoopInfo(*cur_frame,
                                           *map->getFrame(sc_reloc_info.frame_id));
                    if (loop_info.source_id >= 0 && loop_info.target_id >= 0 &&
                        loop_info.score > 0.001)
                    {
                        loop_info.Print();
                        back_end->AddLoopInfo(loop_info);
                    }
                }

                last_frame = cur_frame;
            }

            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());

            // map += world_cloud;
            // map->points.insert(map->points.end(), world_cloud->points.begin(),
            //                    world_cloud->points.end());

#ifdef USE_RERUN
            // rec.log("world/path",
            //     rerun::Transform3D(
            //         rerun::to_rerun(lidar_t),
            //         rerun::to_rerun(lidar_r)));

            std::vector<rerun::Position3D> rec_points = rerun::to_rerun(*world_cloud);

            rec_path.push_back(rerun::to_rerun(lidar_t));

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

            publishCloud(m_world_cloud_pub, world_cloud,
                         "map", timestamp);
            publishPath(m_path_pub, "map", timestamp, LIO_node_inference.m_kf, LIO_node_inference.m_state_data);
#endif
        }

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

    auto kf_traj_3dof_raw = map->GetKFTrajectory3dof();

#ifdef USE_RERUN

    auto kf_traj_3dof_raw_rerrun = rerun::to_rerun(kf_traj_3dof_raw);

    rec.log("world/path_kf", rerun::LineStrips3D(rerun::LineStrip3D(kf_traj_3dof_raw_rerrun)));

    visualizeLoopPairs(rec, *map, back_end->GetLoopInfo());

#endif

    back_end->Optimize();

    auto kf_traj_3dof_opt = map->GetKFTrajectory3dof();

#ifdef USE_RERUN

    auto kf_traj_3dof_opt_rerrun = rerun::to_rerun(kf_traj_3dof_opt);

    rec.log("world/path_kf_opt", rerun::LineStrips3D(rerun::LineStrip3D(kf_traj_3dof_opt_rerrun)));

#endif

    const std::string map_folder_path = FileSystem2::JoinPath(bag_path, "map");
    std::filesystem::create_directories(map_folder_path);

    auto pointcloud_map = map->toPointCloudMap();

    pointcloud_map->saveFromFolder(map_folder_path);

    return 0;
}