#pragma once

#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <unistd.h>

#include <unitree/robot/channel/channel_factory.hpp>
#include <unitree/robot/channel/channel_publisher.hpp>
#include <unitree/robot/channel/channel_subscriber.hpp>
#include <unitree/robot/b2/motion_switcher/motion_switcher_client.hpp>
#include <unitree/idl/hg/LowCmd_.hpp>
#include <unitree/idl/hg/LowState_.hpp>

namespace r1_vision {

using unitree_hg::msg::dds_::LowCmd_;
using unitree_hg::msg::dds_::LowState_;
using namespace unitree::robot;

inline constexpr int kNumMotors = 12;

inline constexpr std::array<int, kNumMotors> kJointId = {
    15, 16, 17, 18, 19,
    22, 23, 24, 25, 26,
    29, 30
};

inline constexpr std::array<const char*, kNumMotors> kJointName = {
    "left_shoulder_pitch",
    "left_shoulder_roll",
    "left_shoulder_yaw",
    "left_elbow",
    "left_wrist_roll",
    "right_shoulder_pitch",
    "right_shoulder_roll",
    "right_shoulder_yaw",
    "right_elbow",
    "right_wrist_roll",
    "head_pitch",
    "head_yaw"
};

inline constexpr std::array<float, kNumMotors> kKp = {
    100.f, 100.f, 100.f, 100.f, 50.f,
    100.f, 100.f, 100.f, 100.f, 50.f,
    50.f, 10.f
};

inline constexpr std::array<float, kNumMotors> kKd = {
    2.f, 2.f, 2.f, 2.f, 2.f,
    2.f, 2.f, 2.f, 2.f, 2.f,
    2.f, 0.1f
};

inline uint32_t Crc32Core(uint32_t* ptr, uint32_t len) {
    uint32_t xbit = 0;
    uint32_t data = 0;
    uint32_t crc = 0xFFFFFFFFu;
    constexpr uint32_t polynomial = 0x04c11db7u;

    for (uint32_t i = 0; i < len; ++i) {
        xbit = 1u << 31;
        data = ptr[i];
        for (uint32_t bits = 0; bits < 32; ++bits) {
            if (crc & 0x80000000u) {
                crc <<= 1;
                crc ^= polynomial;
            } else {
                crc <<= 1;
            }
            if (data & xbit) crc ^= polynomial;
            xbit >>= 1;
        }
    }
    return crc;
}

class RobotController {
public:
    explicit RobotController(const std::string& network_interface) {
        ChannelFactory::Instance()->Init(0, network_interface);

        motion_switcher_ = std::make_shared<unitree::robot::b2::MotionSwitcherClient>();
        motion_switcher_->SetTimeout(5.0f);
        motion_switcher_->Init();

        std::string form, name;
        while (true) {
            motion_switcher_->CheckMode(form, name);
            if (name.empty()) break;
            std::cout << "[SYS] Releasing motion mode: " << name << std::endl;
            if (motion_switcher_->ReleaseMode()) {
                std::cerr << "[WARN] ReleaseMode failed" << std::endl;
            }
            ::sleep(2);
        }

        lowcmd_ = std::make_shared<ChannelPublisher<LowCmd_>>("rt/lowcmd");
        lowcmd_->InitChannel();

        lowstate_ = std::make_shared<ChannelSubscriber<LowState_>>("rt/lowstate");
        lowstate_->InitChannel(
            std::bind(&RobotController::StateHandler, this, std::placeholders::_1), 1);

        std::cout << "[SYS] Waiting for LowState..." << std::endl;
        while (!state_received_.load()) ::usleep(10000);

        {
            std::lock_guard<std::mutex> lock(mutex_);
            target_q_ = current_q_;
        }

        running_.store(true);
        writer_thread_ = std::thread(&RobotController::CommandLoop, this);
        std::cout << "[SYS] Low-level controller running at 500 Hz." << std::endl;
    }

    ~RobotController() {
        HoldCurrent();
        running_.store(false);
        if (writer_thread_.joinable()) writer_thread_.join();
    }

    void StateHandler(const void* message) {
        const auto& state = *static_cast<const LowState_*>(message);
        const uint32_t crc = Crc32Core(
            reinterpret_cast<uint32_t*>(const_cast<LowState_*>(&state)),
            (sizeof(LowState_) >> 2) - 1);
        if (state.crc() != crc) return;

        std::lock_guard<std::mutex> lock(mutex_);
        for (int i = 0; i < kNumMotors; ++i) {
            current_q_[i] = state.motor_state().at(kJointId[i]).q();
        }
        mode_machine_ = state.mode_machine();
        state_received_.store(true);
    }

    void CommandLoop() {
        auto next = std::chrono::steady_clock::now();
        while (running_.load()) {
            next += std::chrono::milliseconds(2);

            LowCmd_ cmd;
            {
                std::lock_guard<std::mutex> lock(mutex_);
                cmd.mode_pr() = 0;
                cmd.mode_machine() = mode_machine_;
                for (int i = 0; i < kNumMotors; ++i) {
                    auto& motor = cmd.motor_cmd().at(kJointId[i]);
                    motor.mode() = 1;
                    motor.q() = target_q_[i];
                    motor.dq() = 0.0f;
                    motor.kp() = kKp[i];
                    motor.kd() = kKd[i];
                    motor.tau() = 0.0f;
                }
            }

            cmd.crc() = Crc32Core(
                reinterpret_cast<uint32_t*>(&cmd),
                (sizeof(cmd) >> 2) - 1);
            lowcmd_->Write(cmd);
            std::this_thread::sleep_until(next);
        }
    }

    std::array<float, kNumMotors> CurrentQ() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return current_q_;
    }

    float CurrentJoint(int joint) const {
        std::lock_guard<std::mutex> lock(mutex_);
        return current_q_.at(joint);
    }

    void SetTargets(const std::array<float, kNumMotors>& q) {
        std::lock_guard<std::mutex> lock(mutex_);
        target_q_ = q;
    }

    void SetJointTarget(int joint, float q) {
        std::lock_guard<std::mutex> lock(mutex_);
        target_q_.at(joint) = q;
    }

    void MoveJointRelative(int joint, float delta_rad, double seconds = 0.6) {
        const float start = CurrentJoint(joint);
        const float target = start + delta_rad;
        const auto t0 = std::chrono::steady_clock::now();

        while (true) {
            const double elapsed = std::chrono::duration<double>(
                std::chrono::steady_clock::now() - t0).count();
            const double a = std::clamp(elapsed / seconds, 0.0, 1.0);
            const double s = a * a * (3.0 - 2.0 * a);
            SetJointTarget(joint, static_cast<float>(start + delta_rad * s));
            if (a >= 1.0) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
        }

        SetJointTarget(joint, target);
        std::this_thread::sleep_for(std::chrono::milliseconds(300));
    }

    void HoldCurrent() {
        SetTargets(CurrentQ());
        std::this_thread::sleep_for(std::chrono::milliseconds(250));
    }

private:
    std::shared_ptr<unitree::robot::b2::MotionSwitcherClient> motion_switcher_;
    ChannelPublisherPtr<LowCmd_> lowcmd_;
    ChannelSubscriberPtr<LowState_> lowstate_;
    std::thread writer_thread_;
    std::atomic<bool> running_{false};
    std::atomic<bool> state_received_{false};
    mutable std::mutex mutex_;
    std::array<float, kNumMotors> current_q_{};
    std::array<float, kNumMotors> target_q_{};
    uint8_t mode_machine_{0};
};

} // namespace r1_vision