// Copyright 2025 Enactic, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <atomic>
#include <array>
#include <algorithm>
#include <cmath>
#include <chrono>
#include <controller/control.hpp>
#include <controller/dynamics.hpp>
#include <csignal>
#include <filesystem>
#include <iostream>
#include <openarm/can/socket/openarm.hpp>
#include <openarm/damiao_motor/dm_motor_constants.hpp>
#include <linux/can.h>
#include <linux/can/raw.h>
#include <net/if.h>
#include <openarm_port/openarm_init.hpp>
#include <periodic_timer_thread.hpp>
#include <poll.h>
#include <robot_state.hpp>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <thread>
#include <unistd.h>
#include <yamlloader.hpp>

std::atomic<bool> keep_running(true);
constexpr std::size_t SENSOR_JOINT_NUM = 7;
std::array<std::atomic<double>, SENSOR_JOINT_NUM> sensor_angle_rad;
std::array<std::atomic<double>, SENSOR_JOINT_NUM> sensor_velocity_rad_s;
std::array<std::atomic<bool>, SENSOR_JOINT_NUM> sensor_angle_ready;

struct openarm_joint_sensor_status
{
    uint8_t id;
    double angle_rad;
    double velocity_rad_s;
    double torque_Nm;
};


constexpr std::array<canid_t, SENSOR_JOINT_NUM> SENSOR_CAN_ID = {0x01, 0x02, 0x03, 0x04,
                                                                  0x05, 0x06, 0x07};
// Per-joint sign: use +1.0 or -1.0 to match each joint's positive rotation direction.
constexpr std::array<double, SENSOR_JOINT_NUM> SENSOR_DIRECTION_SIGN = {
    -1.0, 1.0, -1.0, 1.0, 1.0, 1.0, 1.0};
// Per-joint zero offset in radians: joint_angle = sign * raw_angle + offset.
constexpr std::array<double, SENSOR_JOINT_NUM> SENSOR_ZERO_OFFSET_RAD = {
    2.75, -2.49, 0.08, 0.0, 0.0, 0.0, 0.0};
// constexpr double SENSOR_SCALE = 1e-3;     // TODO: replace with your sensor scaling
// constexpr double SENSOR_OFFSET = 0.0;
constexpr double SENSOR_FILTER_ALPHA = 0.05;  // 0~1, larger = less smoothing
constexpr double SENSOR_VEL_FILTER_ALPHA = 0.2;  // 0~1, larger = less smoothing
constexpr double pi = 3.14159265358979323846;
constexpr double TWO_PI = 2.0 * pi;

double wrap_to_pi(double angle_rad) {
    return std::atan2(std::sin(angle_rad), std::cos(angle_rad));
}

double convert_0_to_2pi_to_signed(double angle_rad_0_to_2pi) {
    double angle = std::fmod(angle_rad_0_to_2pi, TWO_PI);
    if (angle < 0.0) {
        angle += TWO_PI;
    }
    if (angle > pi) {
        angle -= TWO_PI;
    }
    return angle;
}

int sensor_id_to_joint_index(canid_t id) {
    for (std::size_t i = 0; i < SENSOR_CAN_ID.size(); ++i) {
        if (SENSOR_CAN_ID[i] == id) {
            return static_cast<int>(i);
        }
    }
    return -1;
}

int open_can_socket(const std::string &can_interface) {
    int sock = socket(PF_CAN, SOCK_RAW, CAN_RAW);
    if (sock < 0) {
        throw std::runtime_error("Failed to open CAN socket");
    }

    struct ifreq ifr;
    std::snprintf(ifr.ifr_name, IFNAMSIZ, "%s", can_interface.c_str());
    if (ioctl(sock, SIOCGIFINDEX, &ifr) < 0) {
        close(sock);
        throw std::runtime_error("Failed to get CAN interface index");
    }

    struct sockaddr_can addr;
    std::memset(&addr, 0, sizeof(addr));
    addr.can_family = AF_CAN;
    addr.can_ifindex = ifr.ifr_ifindex;

    if (bind(sock, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) < 0) {
        close(sock);
        throw std::runtime_error("Failed to bind CAN socket");
    }

    struct can_filter filter;
    filter.can_id = 0x01;
    filter.can_mask = CAN_SFF_MASK & 0xF0;
    // filter.can_mask = CAN_SFF_MASK;
    if (setsockopt(sock, SOL_CAN_RAW, CAN_RAW_FILTER, &filter, sizeof(filter)) < 0) {
        close(sock);
        throw std::runtime_error("Failed to set CAN filter");
    }

    return sock;
}

bool read_sensor(int sock, openarm_joint_sensor_status &status_out) {
    struct pollfd pfd;
    pfd.fd = sock;
    pfd.events = POLLIN;
    int ret = poll(&pfd, 1, 10);
    if (ret <= 0 || !(pfd.revents & POLLIN)) {
        return false;
    }

    struct can_frame frame;
    int nbytes = read(sock, &frame, sizeof(frame));
    if (nbytes != static_cast<int>(sizeof(frame))) {
        return false;
    }

    const canid_t frame_id = frame.can_id & CAN_SFF_MASK;
    bool valid_id = false;
    for (const auto &id : SENSOR_CAN_ID) {
        if (frame_id == id) {
            valid_id = true;
            break;
        }
        // std::cout << "Received CAN ID: " << frame_id << ", expected: " << id << std::endl;
    }
    if (!valid_id || frame.can_dlc < 5) {
        return false;
    }

    int32_t rawPos = 0;
    rawPos |= static_cast<int32_t>(frame.data[1]) << 8;
    rawPos |= static_cast<int32_t>(frame.data[2]);
    int16_t rawVel = 0;
    rawVel |= static_cast<int32_t>(frame.data[3]) << 8;
    rawVel |= static_cast<int32_t>(frame.data[4]);

    status_out.id = frame_id;
    status_out.velocity_rad_s = (static_cast<double>(rawVel) / 100.0 / 4096.0) * 360.0 * (pi / 180.0);

    // Sensor reports angle in [0, 2pi). Keep it as-is here and normalize later.
    status_out.angle_rad = (static_cast<double>(rawPos) / 4096.0) * TWO_PI;
    std::cout << "CAN ID: " << frame_id << " angle: " << status_out.angle_rad << " velocity: " << status_out.velocity_rad_s << std::endl;
    return true;
}

void signal_handler(int signal) {
    if (signal == SIGINT) {
        std::cout << "\nCtrl+C detected. Exiting loop..." << std::endl;
        keep_running = false;
    }
}

class FollowerArmThread : public PeriodicTimerThread {
public:
    FollowerArmThread(std::shared_ptr<RobotSystemState> robot_state, Control *control_f,
                                            std::size_t joint_num, double hz = 500.0)
                : PeriodicTimerThread(hz), robot_state_(robot_state), control_f_(control_f),
                    joint_num_(joint_num) {}

protected:
    void before_start() override { std::cout << "follower start thread " << std::endl; }

    void after_stop() override { std::cout << "follower stop thread " << std::endl; }

    void on_timer() override {
        static auto prev_time = std::chrono::steady_clock::now();
        for (std::size_t i = 0; i < joint_num_; ++i) {
            if (!sensor_angle_ready[i].load()) {
                continue;
            }
            control_f_->SetArmJointReference(static_cast<int>(i), sensor_angle_rad[i].load(),
                                             sensor_velocity_rad_s[i].load());
            // std::cout << "Joint " << i << " angle: " << sensor_angle_rad[i].load()
            //           << " velocity: " << sensor_velocity_rad_s[i].load() << std::endl;
        }

        control_f_->unilateral_step();

        auto now = std::chrono::steady_clock::now();

        auto elapsed_us =
            std::chrono::duration_cast<std::chrono::microseconds>(now - prev_time).count();
        prev_time = now;
        // std::cout << "Leader joint 0 position: " << robot_state_->arm_state().get_reference(0).position
        //           << std::endl;
        // std::cout << "Leader joint 0 velocity: " << robot_state_->arm_state().get_reference(0).velocity
        //           << std::endl;
        // std::cout << "[Follower] Period: " << elapsed_us << " us" << std::endl;
    }

private:
    std::shared_ptr<RobotSystemState> robot_state_;
    Control *control_f_;
    std::size_t joint_num_;
};


int main(int argc, char **argv) {
    try {
        std::signal(SIGINT, signal_handler);
        for (std::size_t i = 0; i < SENSOR_JOINT_NUM; ++i) {
            sensor_angle_rad[i].store(0.0);
            sensor_velocity_rad_s[i].store(0.0);
            sensor_angle_ready[i].store(false);
        }

        // default configration
        std::string arm_side = "right_arm";
        std::string leader_urdf_path;
        std::string follower_urdf_path;
        std::string follower_can_interface = "can2";
        std::string sensor_can_interface = "can0";

        if (argc < 3) {
            std::cerr
                << "Usage: " << argv[0]
                     << " <leader_urdf_path> <follower_urdf_path> [arm_side] [leader_can] [follower_can]"
                << std::endl;
            return 1;
        }

        // Required: URDF paths
        leader_urdf_path = argv[1];
        follower_urdf_path = argv[2];

        // Optional: arm_side
        if (argc >= 4) {
            arm_side = argv[3];
            if (arm_side != "left_arm" && arm_side != "right_arm") {
                std::cerr << "[ERROR] Invalid arm_side: " << arm_side
                          << ". Must be 'left_arm' or 'right_arm'." << std::endl;
                return 1;
            }
        }

        // Optional: CAN interfaces
        if (argc >= 6) {
            follower_can_interface = argv[5];
        }

        // Optional: sensor CAN interface
        if (argc >= 7) {
            sensor_can_interface = argv[6];
        }

        // URDF file existence check
        if (!std::filesystem::exists(leader_urdf_path)) {
            std::cerr << "[ERROR] Leader URDF not found: " << leader_urdf_path << std::endl;
            return 1;
        }
        if (!std::filesystem::exists(follower_urdf_path)) {
            std::cerr << "[ERROR] Follower URDF not found: " << follower_urdf_path << std::endl;
            return 1;
        }

        // Setup dynamics
        std::string root_link = "openarm_body_link0";
        std::string leaf_link =
            (arm_side == "left_arm") ? "openarm_left_hand" : "openarm_right_hand";

        // Output confirmation
        std::cout << "=== OpenArm Unilateral Control ===" << std::endl;
        std::cout << "Arm side         : " << arm_side << std::endl;
        std::cout << "Follower CAN     : " << follower_can_interface << std::endl;
        std::cout << "Leader URDF path : " << leader_urdf_path << std::endl;
        std::cout << "Follower URDF path: " << follower_urdf_path << std::endl;
        std::cout << "Root link         : " << root_link << std::endl;
        std::cout << "Leaf link         : " << leaf_link << std::endl;
        std::cout << "Sensor CAN        : " << sensor_can_interface << std::endl;

        int sensor_can_socket = open_can_socket(sensor_can_interface);
        std::thread can_read_thread([&]() {
            std::array<double, SENSOR_JOINT_NUM> filtered_angle{};
            std::array<double, SENSOR_JOINT_NUM> filtered_velocity{};
            std::array<bool, SENSOR_JOINT_NUM> filter_initialized{};
            openarm_joint_sensor_status status;
            while (keep_running) {
                if (read_sensor(sensor_can_socket, status)) {
                    const int joint_idx = sensor_id_to_joint_index(status.id);
                    if (joint_idx < 0) {
                        continue;
                    }

                    const std::size_t i = static_cast<std::size_t>(joint_idx);
                    const double sensor_angle_signed = convert_0_to_2pi_to_signed(status.angle_rad);
                    const double angle = wrap_to_pi(
                        SENSOR_DIRECTION_SIGN[i] * sensor_angle_signed + SENSOR_ZERO_OFFSET_RAD[i]);
                    const double velocity = SENSOR_DIRECTION_SIGN[i] * status.velocity_rad_s;

                    if (!filter_initialized[i]) {
                        filtered_angle[i] = angle;
                        filtered_velocity[i] = velocity;
                        filter_initialized[i] = true;
                    } else {
                        filtered_angle[i] = SENSOR_FILTER_ALPHA * angle +
                                            (1.0 - SENSOR_FILTER_ALPHA) * filtered_angle[i];
                        filtered_velocity[i] = SENSOR_VEL_FILTER_ALPHA * velocity +
                                               (1.0 - SENSOR_VEL_FILTER_ALPHA) *
                                                   filtered_velocity[i];
                    }

                    sensor_angle_rad[i].store(filtered_angle[i]);
                    sensor_velocity_rad_s[i].store(filtered_velocity[i]);
                    sensor_angle_ready[i].store(true);
                    std::cout << "Joint " << i << " filtered_angle: " << filtered_angle[i]
                              << ", filtered_velocity: " << filtered_velocity[i] << std::endl;
                }
            }
        });

        YamlLoader leader_loader("/home/shen/openarm_ros2_ws/src/openarm_teleop/config/leader.yaml");
        YamlLoader follower_loader("/home/shen/openarm_ros2_ws/src/openarm_teleop/config/follower.yaml");

        // Follower parameters
        std::vector<double> follower_kp = follower_loader.get_vector("FollowerArmParam", "Kp");
        std::vector<double> follower_kd = follower_loader.get_vector("FollowerArmParam", "Kd");
        std::vector<double> follower_Fc = follower_loader.get_vector("FollowerArmParam", "Fc");
        std::vector<double> follower_k = follower_loader.get_vector("FollowerArmParam", "k");
        std::vector<double> follower_Fv = follower_loader.get_vector("FollowerArmParam", "Fv");
        std::vector<double> follower_Fo = follower_loader.get_vector("FollowerArmParam", "Fo");

        Dynamics *leader_arm_dynamics = new Dynamics(leader_urdf_path, root_link, leaf_link);
        leader_arm_dynamics->Init();

        Dynamics *follower_arm_dynamics = new Dynamics(follower_urdf_path, root_link, leaf_link);
        follower_arm_dynamics->Init();

        std::cout << "=== Initializing Follower OpenArm ===" << std::endl;
        openarm::can::socket::OpenArm *follower_openarm =
            openarm_init::OpenArmInitializer::initialize_openarm(follower_can_interface, true);

        size_t follower_arm_motor_num = follower_openarm->get_arm().get_motors().size();
        size_t follower_hand_motor_num = follower_openarm->get_gripper().get_motors().size();

        std::cout << "follower arm motor num : " << follower_arm_motor_num << std::endl;
        std::cout << "follower hand motor num : " << follower_hand_motor_num << std::endl;

        // Declare robot_state (Joint and motor counts are assumed to be equal)

        std::shared_ptr<RobotSystemState> follower_state =
            std::make_shared<RobotSystemState>(follower_arm_motor_num, follower_hand_motor_num);

        Control *control_follower =
            new Control(follower_openarm, leader_arm_dynamics, follower_arm_dynamics,
                        follower_state, 1.0 / FREQUENCY, ROLE_FOLLOWER, arm_side,
                        follower_arm_motor_num, follower_hand_motor_num);

        control_follower->SetParameter(follower_kp, follower_kd, follower_Fc, follower_k,
                                       follower_Fv, follower_Fo);

        // set home postion
        // std::thread thread_l(&Control::AdjustPosition, control_leader);
        // std::thread thread_f(&Control::AdjustPosition, control_follower);
        // thread_l.join();
        // thread_f.join();

        // Start control process
        const std::size_t controlled_joint_num =
            std::min<std::size_t>(follower_arm_motor_num, SENSOR_JOINT_NUM);
        FollowerArmThread follower_thread(follower_state, control_follower, controlled_joint_num,
                                          FREQUENCY);

        follower_thread.start_thread();

        while (keep_running) {
            std::this_thread::sleep_for(std::chrono::milliseconds(1000));
        }

        follower_thread.stop_thread();

        follower_openarm->disable_all();

        if (can_read_thread.joinable()) {
            can_read_thread.join();
        }
        close(sensor_can_socket);

    } catch (const std::exception &e) {
        std::cerr << e.what() << '\n';
    }

    return 0;
}