#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <nav_msgs/msg/path.hpp>
#include <nav_msgs/msg/occupancy_grid.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <sensor_msgs/msg/laser_scan.hpp>
#include <std_msgs/msg/string.hpp>
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/msg/transform_stamped.hpp>

#include <memory>
#include <mutex>
#include <cmath>
#include <vector>

#include "car_nav_lite/types.hpp"
#include "car_nav_lite/grid_map.hpp"
#include "car_nav_lite/localizer.hpp"
#include "car_nav_lite/bt_executor.hpp"
#include "car_nav_lite/fence_model.hpp"
#include "car_nav_lite/session_mapper.hpp"

using namespace std::chrono_literals;
using namespace car_nav_lite;
using LaserScanMsg = sensor_msgs::msg::LaserScan;
using OdometryMsg  = nav_msgs::msg::Odometry;
using ImuMsg       = sensor_msgs::msg::Imu;
using PoseStMsg    = geometry_msgs::msg::PoseStamped;

class NavLiteNode : public rclcpp::Node {
public:
    NavLiteNode() : Node("nav_lite_node") {
        // ── ROS2 dynamic parameters ──
        declare_parameter("map_file",   std::string(""));
        declare_parameter("test_mode",  false);
        declare_parameter("lookahead",  0.5);
        declare_parameter("max_speed",  1.0);
        declare_parameter("min_turn_radius", 0.20);
        declare_parameter("steer_limit", 0.6);
        declare_parameter("speed_curv_factor", 6.0);
        declare_parameter("calib_scan_count", 30);
        declare_parameter("fence_expected_size", 5.0);
        declare_parameter("fence_size_tolerance", 0.35);
        declare_parameter("wall_ransac_thresh", 0.04);
        declare_parameter("cone_mark_radius", 0.18);
        declare_parameter("dynamic_obstacle_enable", false);
        _read_params();

        param_cb_ = add_on_set_parameters_callback(
            [this](const std::vector<rclcpp::Parameter>& params) {
                for (auto& p : params) {
                    std::string n = p.get_name();
                    if (n=="lookahead") la_=p.as_double();
                    else if (n=="max_speed") ms_=p.as_double();
                    else if (n=="min_turn_radius") mr_=p.as_double();
                    else if (n=="steer_limit") sl_=p.as_double();
                    else if (n=="speed_curv_factor") cf_=p.as_double();
                    else if (n=="test_mode") { test_=p.as_bool(); tps_=false; }
                    else if (n=="calib_scan_count") calib_scan_count_=p.as_int();
                    else if (n=="fence_expected_size") fence_expected_size_=p.as_double();
                    else if (n=="fence_size_tolerance") fence_size_tolerance_=p.as_double();
                    else if (n=="wall_ransac_thresh") wall_ransac_thresh_=p.as_double();
                    else if (n=="cone_mark_radius") cone_mark_radius_=p.as_double();
                    else if (n=="dynamic_obstacle_enable") dynamic_obstacle_enable_=p.as_bool();
                }
                RCLCPP_INFO(get_logger(),"params: la=%.2f ms=%.1f mr=%.2f sl=%.2f cf=%.1f test=%d",
                            la_,ms_,mr_,sl_,cf_,test_);
                return rcl_interfaces::msg::SetParametersResult().set__successful(true);
            });

        // ── Subscribers ──
        sub_scan_ = create_subscription<LaserScanMsg>("/scan", 10,
            std::bind(&NavLiteNode::scan_cb, this, std::placeholders::_1));
        sub_odom_ = create_subscription<OdometryMsg>("/odom", 10,
            std::bind(&NavLiteNode::odom_cb, this, std::placeholders::_1));
        sub_imu_  = create_subscription<ImuMsg>("/imu/data_raw", 10,
            std::bind(&NavLiteNode::imu_cb, this, std::placeholders::_1));
        sub_goal_ = create_subscription<PoseStMsg>("/goal_pose", 10,
            std::bind(&NavLiteNode::goal_cb, this, std::placeholders::_1));

        // ── Publishers ──
        pub_cmd_  = create_publisher<geometry_msgs::msg::Twist>("/cmd_vel", 10);
        pub_path_ = create_publisher<nav_msgs::msg::Path>("/global_path", 10);
        pub_map_  = create_publisher<nav_msgs::msg::OccupancyGrid>("/map", rclcpp::QoS(1).transient_local());
        pub_status_ = create_publisher<std_msgs::msg::String>("/nav_status", 10);
        pub_metrics_ = create_publisher<std_msgs::msg::String>("/nav_metrics", 10);
        bt_.setFailCallback([this](){
            std_msgs::msg::String msg; msg.data = "fail";
            pub_status_->publish(msg);
        });
        tf_broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>(*this);

        if (!mf_.empty()) {
            if (!map_.loadPGM(mf_)) RCLCPP_WARN(get_logger(),"Map load fail");
            else RCLCPP_INFO(get_logger(),"Map OK %dx%d",map_.width(),map_.height());
        }

        timer_ = create_wall_timer(50ms, std::bind(&NavLiteNode::tick, this));
        map_timer_ = create_wall_timer(1s, std::bind(&NavLiteNode::publish_map, this));
        metrics_timer_ = create_wall_timer(200ms, std::bind(&NavLiteNode::publish_metrics, this));
        morph_timer_ = create_wall_timer(100ms, std::bind(&NavLiteNode::morph_close, this));

        RCLCPP_INFO(get_logger(),"car_nav_lite ready. test_mode:=true for auto path. "
                    "ros2 param set /nav_lite_node <name> <value>");
        RCLCPP_INFO(get_logger(),"Params: la=%.2f ms=%.1f mr=%.2f sl=%.2f cf=%.1f",
                    la_,ms_,mr_,sl_,cf_);
    }

private:
    // ── Dynamic params ──
    std::string mf_;
    bool test_=false, tps_=false;
    double la_=0.5, ms_=1.0, mr_=0.20, sl_=0.6, cf_=6.0;
    int calib_scan_count_=30;
    double fence_expected_size_=5.0, fence_size_tolerance_=0.35;
    double wall_ransac_thresh_=0.04, cone_mark_radius_=0.18;
    bool dynamic_obstacle_enable_=false;
    rclcpp::node_interfaces::OnSetParametersCallbackHandle::SharedPtr param_cb_;
    void _read_params(){
        mf_=get_parameter("map_file").as_string();
        test_=get_parameter("test_mode").as_bool();
        la_=get_parameter("lookahead").as_double();
        ms_=get_parameter("max_speed").as_double();
        mr_=get_parameter("min_turn_radius").as_double();
        sl_=get_parameter("steer_limit").as_double();
        cf_=get_parameter("speed_curv_factor").as_double();
        calib_scan_count_=get_parameter("calib_scan_count").as_int();
        fence_expected_size_=get_parameter("fence_expected_size").as_double();
        fence_size_tolerance_=get_parameter("fence_size_tolerance").as_double();
        wall_ransac_thresh_=get_parameter("wall_ransac_thresh").as_double();
        cone_mark_radius_=get_parameter("cone_mark_radius").as_double();
        dynamic_obstacle_enable_=get_parameter("dynamic_obstacle_enable").as_bool();
    }

    // ── State ──
    std::mutex mtx_;
    LaserScan latest_scan_; bool has_scan_=false;
    Pose2D scan_pose_snapshot_;
    double latest_vx_=0, latest_wz_=0, latest_gyro_z_=0;
    rclcpp::Time last_odom_t_{0,0,RCL_ROS_TIME};
    rclcpp::Time last_imu_t_{0,0,RCL_ROS_TIME};
    GridMap map_; Localizer localizer_; BTExecutor bt_;
    SessionMapper session_mapper_;
    FenceModel fence_model_;
    enum SessionState { CALIBRATING, READY };
    SessionState session_state_ = CALIBRATING;
    int calib_cnt_ = 0;
    rclcpp::Subscription<LaserScanMsg>::SharedPtr sub_scan_;
    rclcpp::Subscription<OdometryMsg>::SharedPtr  sub_odom_;
    rclcpp::Subscription<ImuMsg>::SharedPtr       sub_imu_;
    rclcpp::Subscription<PoseStMsg>::SharedPtr    sub_goal_;
    rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr pub_cmd_;
    rclcpp::Publisher<nav_msgs::msg::Path>::SharedPtr pub_path_;
    rclcpp::Publisher<nav_msgs::msg::OccupancyGrid>::SharedPtr pub_map_;
    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr pub_status_;
    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr pub_metrics_;
    std::shared_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
    rclcpp::TimerBase::SharedPtr timer_, map_timer_, metrics_timer_, morph_timer_;

    void scan_cb(LaserScanMsg::SharedPtr m){std::lock_guard lk(mtx_);
        latest_scan_.angle_min=m->angle_min;latest_scan_.angle_max=m->angle_max;
        latest_scan_.angle_increment=m->angle_increment;
        latest_scan_.range_min=m->range_min;latest_scan_.range_max=m->range_max;
        latest_scan_.ranges=m->ranges;has_scan_=true;
        scan_pose_snapshot_ = localizer_.pose();

        if (session_state_ == CALIBRATING && calib_cnt_ < calib_scan_count_) {
            session_mapper_.addScan(latest_scan_, localizer_.pose());
            if (++calib_cnt_ >= calib_scan_count_) {
                std::vector<Eigen::Vector2d> cones;
                if (session_mapper_.build(fence_model_, cones, fence_expected_size_, fence_size_tolerance_, wall_ransac_thresh_)) {
                    localizer_.setFenceModel(fence_model_);
                    for (auto& c : cones) map_.markCone(c.x(), c.y(), cone_mark_radius_, 255);
                    map_.markFenceModel(fence_model_);
                    map_.morphologyClose(3);
                    map_.freezeSessionMap();
                    session_state_ = READY;
                    fprintf(stderr, "[Session] READY — fence %.2fx%.2f yaw=%.1f°, %zu cones\n",
                            fence_model_.width, fence_model_.height, fence_model_.yaw_field*57.3, cones.size());
                } else {
                    calib_cnt_ = 0; session_mapper_.reset();
                    fprintf(stderr, "[Session] calibration failed, retrying...\n");
                }
            }
        } else {
            localizer_.correctWithScanCV(latest_scan_, map_, &scan_pose_snapshot_);
        }}
    void odom_cb(OdometryMsg::SharedPtr m){std::lock_guard lk(mtx_);
        // EXACT-MPPI pattern: trust odom pose directly (sim = ground truth,
        // real robot = wheel-encoder / EKF output). No manual integration.
        auto& p = m->pose.pose;
        double qx=p.orientation.x, qy=p.orientation.y, qz=p.orientation.z, qw=p.orientation.w;
        double yaw = std::atan2(2.0*(qw*qz + qx*qy), 1.0 - 2.0*(qy*qy + qz*qz));
        localizer_.setOdomPose({p.position.x, p.position.y, yaw});
        last_odom_t_ = rclcpp::Time(m->header.stamp);
        latest_vx_ = m->twist.twist.linear.x;
        latest_wz_ = m->twist.twist.angular.z;}
    void imu_cb(ImuMsg::SharedPtr m){std::lock_guard lk(mtx_);
        auto now=rclcpp::Time(m->header.stamp);
        double dt=last_imu_t_.nanoseconds()>0?(now-last_imu_t_).seconds():0.005;
        last_imu_t_=now;latest_gyro_z_=m->angular_velocity.z;
        if(dt>0&&dt<1.0)localizer_.updateIMU(latest_gyro_z_,dt);}
    void goal_cb(PoseStMsg::SharedPtr m){std::lock_guard lk(mtx_);
        if(m->header.frame_id!="map"){RCLCPP_WARN(get_logger(),"goal frame not map");return;}
        Pose2D g(m->pose.position.x,m->pose.position.y,0.0);
        double qw=m->pose.orientation.w,qz=m->pose.orientation.z;
        if(qw!=0||qz!=0)g.yaw=std::atan2(2.0*qw*qz,qw*qw-qz*qz);
        bt_.setGoal(g);tps_=false;}

    // ── Test path: figure-8 with 30cm circles ──
    std::vector<Pose2D> _genTestPath(){
        std::vector<Pose2D> wp;
        double cx=0,cy=-0.5,r=0.15;  // 15cm radius = 30cm diameter
        int N=30;
        // Circle 1: counter-clockwise
        for(int i=0;i<=N;++i){
            double a=2.0*M_PI*i/N;
            wp.push_back({cx-r*std::cos(a),cy-r+r*std::sin(a),0});
        }
        // Straight connecting segment
        wp.push_back({cx,cy-2*r,0});
        // Circle 2: clockwise (other direction), centered right
        for(int i=0;i<=N;++i){
            double a=-2.0*M_PI*i/N;
            wp.push_back({cx+r*std::cos(a),cy-3*r+r*std::sin(a),0});
        }
        // Straight back
        wp.push_back({cx,cy,0});
        // Rectangle 2x1.5m
        for(auto& p:std::vector<Pose2D>{{cx-1,cy-0.75,0},{cx+1,cy-0.75,0},{cx+1,cy+0.75,0},{cx-1,cy+0.75,0},{cx-1,cy-0.75,0}})
            wp.push_back(p);
        // Interpolate
        std::vector<Pose2D> dense;
        for(size_t i=0;i+1<wp.size();++i){
            double d=wp[i].distTo(wp[i+1]);
            int steps=std::max(1,(int)(d/0.05));
            for(int s=0;s<steps;++s){
                double t=(double)s/steps;
                dense.push_back({wp[i].x+t*(wp[i+1].x-wp[i].x),wp[i].y+t*(wp[i+1].y-wp[i].y),0});
            }
        }
        for(size_t i=0;i+1<dense.size();++i)dense[i].yaw=dense[i].angleTo(dense[i+1]);
        if(!dense.empty())dense.back().yaw=dense[dense.size()-2].yaw;
        return dense;
    }

    void tick(){
        std::lock_guard lk(mtx_);
        // CALIBRATING: force zero cmd, wait for fence model
        if (session_state_ == CALIBRATING) {
            auto zero = geometry_msgs::msg::Twist();
            pub_cmd_->publish(zero);
            return;
        }
        auto pose=localizer_.pose();publish_tf(pose);
        // READY: only update scan if dynamic obstacles enabled
        if(has_scan_ && dynamic_obstacle_enable_)
            map_.updateScan(latest_scan_,pose);

        // Test mode: auto-activate figure-8 after 3s
        if(test_&&!tps_){
            static int tc=0;
            if(++tc>60){
                auto tp=_genTestPath();
                bt_.setTestPath(tp);tps_=true;
                RCLCPP_INFO(get_logger(),"Test path: fig8+rect, %zu pts",tp.size());
            }
        }

        // Update BT with current params
        bt_.setParams(la_,ms_,mr_,sl_,cf_);
        Twist cmd=bt_.tick(pose,map_);

        // Slow down if localization confidence is low
        if (localizer_.loc_conf_ < 0.3) { cmd.vx *= 0.4; cmd.wz *= 0.4; }

        auto tw=geometry_msgs::msg::Twist();tw.linear.x=cmd.vx;tw.angular.z=cmd.wz;
        pub_cmd_->publish(tw);
        if(!bt_.currentPath().empty()){
            auto pm=nav_msgs::msg::Path();pm.header.stamp=now();pm.header.frame_id="map";
            for(auto&p:bt_.currentPath()){
                geometry_msgs::msg::PoseStamped ps;ps.header.frame_id="map";
                ps.pose.position.x=p.x;ps.pose.position.y=p.y;
                ps.pose.orientation.z=std::sin(p.yaw/2.0);ps.pose.orientation.w=std::cos(p.yaw/2.0);
                pm.poses.push_back(ps);
            }pub_path_->publish(pm);
        }
        static int ct=0;
        if(++ct%60==0)RCLCPP_INFO(get_logger(),"pose=(%.2f,%.2f,%.0f) cmd=(%.2f,%.2f)",
                                  pose.x,pose.y,pose.yaw*57.3,cmd.vx,cmd.wz);
    }

    void publish_map(){std::lock_guard lk(mtx_);
        auto m=nav_msgs::msg::OccupancyGrid();m.header.stamp=now();m.header.frame_id="map";
        m.info.resolution=GRID_RES;m.info.width=m.info.height=GRID_SIZE;
        m.info.origin.position.x=GRID_ORIGIN_X;m.info.origin.position.y=GRID_ORIGIN_Y;
        m.data.resize(GRID_SIZE*GRID_SIZE);
        for(size_t i=0;i<map_.data().size();++i){uint8_t v=map_.data()[i];m.data[i]=(v==0)?0:((v>128)?100:-1);}
        pub_map_->publish(m);}

    void publish_metrics(){
        std::lock_guard lk(mtx_);
        NavMetrics m;
        auto pose = localizer_.pose();
        m.pose_x = pose.x; m.pose_y = pose.y; m.pose_yaw = pose.yaw;
        m.loc_lines = localizer_.loc_lines_;
        m.loc_dx = localizer_.loc_dx_; m.loc_dy = localizer_.loc_dy_; m.loc_dyaw = localizer_.loc_dyaw_;
        bt_.fillMetrics(m);
        static const char* states[] = {"IDLE","PLANNING","TRACKING","ARRIVED","FAILED"};
        auto s = bt_.state(); int si=(int)s; m.state = (si>=0&&si<5)?states[si]:"?";

        char buf[512];
        snprintf(buf, sizeof(buf),
            "{\"state\":\"%s\",\"x\":%.2f,\"y\":%.2f,\"yaw\":%.1f,"
            "\"loc_conf\":%.2f,\"loc_lines\":%d,\"loc_dx\":%.3f,\"loc_dy\":%.3f,\"loc_dyaw\":%.2f,"
            "\"mppi_alive\":%d,\"mppi_total\":%d,\"mppi_minC\":%.1f}",
            m.state, m.pose_x, m.pose_y, m.pose_yaw*57.3,
            localizer_.loc_conf_, m.loc_lines, m.loc_dx, m.loc_dy, m.loc_dyaw*57.3,
            m.mppi_alive, m.mppi_total, m.mppi_min_cost);
        std_msgs::msg::String msg; msg.data = buf;
        pub_metrics_->publish(msg);
    }

    void morph_close(){std::lock_guard lk(mtx_);
        if (session_state_ == READY) return;  // frozen map, no morph needed
        map_.morphologyClose(3);
    }

    void publish_tf(const Pose2D& pose){
        geometry_msgs::msg::TransformStamped tf;
        tf.header.stamp=now();tf.header.frame_id="map";tf.child_frame_id="base_link";
        tf.transform.translation.x=pose.x;tf.transform.translation.y=pose.y;
        double hy=pose.yaw/2.0;tf.transform.rotation.z=std::sin(hy);tf.transform.rotation.w=std::cos(hy);
        tf_broadcaster_->sendTransform(tf);}
};

int main(int argc,char*argv[]){rclcpp::init(argc,argv);rclcpp::spin(std::make_shared<NavLiteNode>());rclcpp::shutdown();return 0;}