#include "core/basic_obj_types.hpp"
#include "robot/fracture_robot/fracture_robot_defs.h"
#include "utility/config_utility.hpp"
#include "utility/coro_worker.hpp"
#include "utility/coro_worker_helper_func.hpp"
#include "utility/name_translator.hpp"
#include "utility/singleton_helper.hpp"
#include "utility/variable_helper.hpp"
#include "utility/versatile_buffer2.hpp"

#include <boost/asio/serial_port.hpp>
#include <boost/asio/use_awaitable.hpp>
#include <boost/smart_ptr/scoped_ptr.hpp>

#include <spdlog/spdlog.h>

#include "utility/assert_utility.h"

DEFAULT_TRISTATE_OBJ_DEF(fracture_robot_interface)

namespace sophiar {

    using boost::asio::async_read;
    using boost::asio::async_write;
    using boost::asio::awaitable;
    using boost::asio::serial_port;
    using boost::asio::use_awaitable;
    using boost::scoped_ptr;

    struct fracture_robot_interface::impl {

        using control_value_type = std::array<int32_t, 6>;
        using control_mode_type = fracture_robot_control_mode_type;

        struct control_mode_translator_type :
                public singleton_helper<control_mode_translator_type>,
                public name_translator<control_mode_type> {
            void initialize() override {
                register_item("direct", control_mode_type::DIRECT);
                register_item("cartesian", control_mode_type::CARTESIAN);
            }
        };

        inline static auto control_mode_translator
                = control_mode_translator_type::get_instance();

        static constexpr auto robot_endian = boost::endian::order::little;
        static constexpr auto default_max_velocity = 5.0; // mm/s
        static constexpr auto default_controller_gain = 5.0;
        static constexpr size_t control_packet_length = 6 * sizeof(int32_t);
        static constexpr uint16_t control_packet_tail = 0x7fff;

        struct screw_info {
            Eigen::Vector3d origin; // Point on the screw
            Eigen::Vector3d direction; // Direction of the screw
        };

        fracture_robot_interface *q_this = nullptr;

        screw_info top_screw, side_screw;
        double max_velocity = default_max_velocity;
        double controller_gain = default_controller_gain;
        control_mode_type current_control_mode;

        variable_index_type current_control_mode_index = -1;
        variable_index_type is_moving_index = -1;
        // direct control
        variable_index_type target_motor_speed_index = -1;
        // cartesian control
        variable_index_type current_tcp_pose_index = -1, target_tcp_pose_index = -1;

        scoped_ptr<serial_port> conn;
        coro_worker::pointer control_worker;

        void fake_inverse_kinematics(const transform_obj::value_type &trans,
                                     array6_obj::value_type &pos) const {
            static constexpr auto side_z_offset = -60.0; // mm
            static constexpr auto top1_y_offset = 440.0;
            static constexpr auto top2_y_offset = 590.0;

            { // side cross
                auto origin = trans * side_screw.origin;
                auto direction = trans.linear() * side_screw.direction;
                auto t = (side_z_offset - origin.z()) / direction.z();
                auto cross = origin + t * direction;
                assert(std::abs(cross.z() - side_z_offset) < 1e-3);
                pos[0] = cross.x();
                pos[3] = cross.y();
            }
            { // top cross
                auto origin = trans * top_screw.origin;
                auto direction = trans.linear() * top_screw.direction;
                auto t1 = (top1_y_offset - origin.y()) / direction.y();
                auto t2 = (top2_y_offset - origin.y()) / direction.y();
                auto cross1 = origin + t1 * direction;
                auto cross2 = origin + t2 * direction;
                assert(std::abs(cross1.y() - top1_y_offset) < 1e-3);
                assert(std::abs(cross2.y() - top2_y_offset) < 1e-3);
                pos[4] = -cross1.x();
                pos[5] = -cross1.z();
                pos[2] = -cross2.x();
                pos[1] = -cross2.z();
            }
        }

        static void encode_arr(const array6_obj::value_type &vec_val,
                               control_value_type &arr_val) {
            static constexpr int base_freq = 1000000; // Hz
            static constexpr auto dis_per_clock = 0.00625; // mm
            static constexpr auto minimum_velocity = 0.5; // mm/s, lower than this will cause stop

            for (int i = 0; i < 6; ++i) {
                if (std::abs(vec_val[i]) < minimum_velocity) {
                    arr_val[i] = 0; // stop
                } else {
                    auto tar_freq = vec_val[i] / dis_per_clock;
                    arr_val[i] = static_cast<int32_t>(base_freq / tar_freq);
                    assert(std::abs(arr_val[i]) < 0x7fff);
                    if ((arr_val[i] & 0xff) == 0x7f) [[unlikely]] { // make the LSB not 0x7f
                        ++arr_val[i];
                    }
                }
            }
        }

        void calc_target_velocity(const array6_obj::value_type &cur_pos,
                                  const array6_obj::value_type &tar_pos,
                                  array6_obj::value_type &vec_val) const {
            static constexpr auto dis_tolerance = 0.5; // 0.1mm, lower than this will cause stop

            double cur_max_vec = 0;
            for (int i = 0; i < 6; ++i) {
                auto dis_offset = tar_pos[i] - cur_pos[i];
                if (std::abs(dis_offset) < dis_tolerance) {
                    vec_val[i] = 0; // stop
                } else {
                    vec_val[i] = controller_gain * dis_offset;
                    cur_max_vec = std::max(std::abs(vec_val[i]), cur_max_vec);
                }
            }

            // clamp velocity
            if (cur_max_vec > max_velocity) {
                auto shrink_factor = max_velocity / cur_max_vec;
                for (int i = 0; i < 6; ++i) {
                    vec_val[i] *= shrink_factor;
                }
            }
        }

        awaitable<void> send_control_packet(const array6_obj::value_type &vec_val) {
            // encode packet
            auto arr_val = control_value_type{};
            encode_arr(vec_val, arr_val);

            // write packet
            bool is_moving = false;
            auto mem = static_memory<control_packet_length>{};
            auto writer = versatile_writer<robot_endian>(mem);
            for (auto item: arr_val) {
                writer << item;
                if (item != 0) {
                    is_moving = true;
                }
            }

            // send packet
            assert(conn->is_open());
            assert(writer.remaining_bytes() == 0);
            co_await async_write_memory_to(*conn, mem);

            // update moving status
            try_update_variable_value<bool_obj>(is_moving_index, is_moving);
            co_return;
        }

        auto create_direct_control_func() {
            assert(current_control_mode == control_mode_type::DIRECT);
            assert(is_valid_id(target_motor_speed_index));
            auto worker_func = [=, this,
                    motor_speed = VARIABLE_MANUAL_DELEGATE(array6_obj, target_motor_speed_index)]() mutable
                    -> awaitable<bool> {
                co_await motor_speed.coro_wait_update();
                if (motor_speed.empty()) co_return true;
                co_await send_control_packet(motor_speed->value);
                co_return true;
            };
            return std::move(worker_func);
        };

        auto create_cartesian_control_func() {
            assert(current_control_mode == control_mode_type::CARTESIAN);
            assert(is_valid_id(current_tcp_pose_index));
            assert(is_valid_id(target_tcp_pose_index));
            auto worker_func = [=, this,
                    current_tcp = VARIABLE_MANUAL_DELEGATE(transform_obj, current_tcp_pose_index),
                    target_tcp = VARIABLE_MANUAL_DELEGATE(transform_obj, target_tcp_pose_index),
                    target_pos = array6_obj::value_type{}]() mutable
                    -> awaitable<bool> {
                co_await current_tcp.coro_wait_update();
                if (current_tcp.empty()) co_return true;
                if (target_tcp.manual_sync()) { // target pose has been updated
                    fake_inverse_kinematics(target_tcp->value, target_pos);
                }
                auto current_pos = array6_obj::value_type{};
                fake_inverse_kinematics(current_tcp->value, current_pos);
                auto vec_val = array6_obj::value_type{};
                calc_target_velocity(current_pos, target_pos, vec_val);
                co_await send_control_packet(vec_val);
                co_return true;
            };
            return std::move(worker_func);
        };

        void create_control_worker() {
            assert(control_worker == nullptr);

            auto exception_handle_func = [](std::exception &) {
                // TODO show log
            };
            auto exit_func = reset_on_exit_func(q_this);

            switch (current_control_mode) {
                case control_mode_type::DIRECT: {
                    auto noexcept_wrapper = make_noexcept_func(
                            create_direct_control_func(), std::move(exception_handle_func));
                    control_worker = make_infinite_coro_worker(std::move(noexcept_wrapper),
                                                               std::move(exit_func));
                    break;
                }
                case control_mode_type::CARTESIAN: {
                    auto noexcept_wrapper = make_noexcept_func(
                            create_cartesian_control_func(), std::move(exception_handle_func));
                    control_worker = make_infinite_coro_worker(std::move(noexcept_wrapper),
                                                               std::move(exit_func));
                    break;
                }
                default: {
                    assert(false);
                    return;
                }
            }

            control_worker->run();

            // set initial state
            try_update_variable_value<bool_obj>(is_moving_index, false);
        }

        void setup_connection(const nlohmann::json &config) {
            assert(conn == nullptr);
            conn.reset(new serial_port(*global_context));

            // open serial port
            auto com_port_name = LOAD_STRING_ITEM("com_port");
            conn->open(com_port_name);
            assert(conn->is_open());

            // config
            conn->set_option(serial_port::baud_rate(9600));
            conn->set_option(serial_port::stop_bits(serial_port::stop_bits::one));
            conn->set_option(serial_port::parity(serial_port::parity::none));
            conn->set_option(serial_port::character_size(8));
            conn->set_option(serial_port::flow_control(serial_port::flow_control::none));
        }

        void load_init_config(const nlohmann::json &config) {
            assert(config.contains("screw"));
            auto &screw_config = config["screw"];
            assert(screw_config.contains("top"));
            { // top screw
                auto &top_screw_config = screw_config["top"];
                assert(top_screw_config.contains("origin"));
                top_screw.origin = scalarxyz_obj::read_value_from_json_array(top_screw_config["origin"]);
                assert(top_screw_config.contains("direction"));
                top_screw.direction = scalarxyz_obj::read_value_from_json_array(top_screw_config["direction"]);
            }
            { // side screw
                auto &side_screw_config = screw_config["side"];
                assert(side_screw_config.contains("origin"));
                side_screw.origin = scalarxyz_obj::read_value_from_json_array(side_screw_config["origin"]);
                assert(side_screw_config.contains("direction"));
                side_screw.direction = scalarxyz_obj::read_value_from_json_array(side_screw_config["direction"]);
            }
        }

        void load_start_config(const nlohmann::json &config) {
            assert(config.contains("input_config"));
            { // output config
                auto &output_config = config["output_config"];
                current_control_mode_index = LOAD_VARIABLE_INDEX2(output_config, u64int_obj, "current_control_mode");
                is_moving_index = TRY_LOAD_VARIABLE_INDEX2(output_config, bool_obj, "is_moving");
            }
            { // input config
                // control mode
                auto &input_config = config["input_config"];
                auto control_mode_str = LOAD_STRING_ITEM2(input_config, "control_mode");
                current_control_mode = control_mode_translator->translate(control_mode_str);
                UPDATE_VARIABLE_VAL(u64int_obj, current_control_mode_index, (uint8_t) current_control_mode);
                // controller parameters
                max_velocity = TRY_LOAD_FLOAT_ITEM2(input_config, "max_velocity", default_max_velocity);
                controller_gain = TRY_LOAD_FLOAT_ITEM2(input_config, "controller_gain", default_controller_gain);
                // control input
                switch (current_control_mode) {
                    case control_mode_type::DIRECT: {
                        target_motor_speed_index
                                = LOAD_VARIABLE_INDEX2(input_config, array6_obj, "target_motor_velocity");
                        break;
                    }
                    case control_mode_type::CARTESIAN: {
                        current_tcp_pose_index = LOAD_VARIABLE_INDEX2(input_config, transform_obj, "current_tcp");
                        target_tcp_pose_index = LOAD_VARIABLE_INDEX2(input_config, transform_obj, "target_tcp");
                        break;
                    }
                    default: {
                        assert(false);
                        return;
                    }
                }
            }
        }
    };

    awaitable<bool> fracture_robot_interface::on_init(const nlohmann::json &config) noexcept {
        pimpl->load_init_config(config);
        try {
            pimpl->setup_connection(config);
        } catch (std::exception &e) {
            SPDLOG_ERROR("Failed to init fracture robot interface: {}", e.what());
            co_return false;
        }
        co_return true;
    }

    awaitable<bool> fracture_robot_interface::on_start(const nlohmann::json &config) noexcept {
        pimpl->load_start_config(config);
        pimpl->create_control_worker();
        co_return true;
    }

    awaitable<void> fracture_robot_interface::on_stop() noexcept {
        SAFE_RESET_WORKER(pimpl->control_worker)
        // update status
        try_update_variable_value<bool_obj>(pimpl->is_moving_index, false);
        co_return;
    }

    awaitable<void> fracture_robot_interface::on_reset() noexcept {
        pimpl->conn.reset(nullptr);
        co_return;
    }

    fracture_robot_interface::fracture_robot_interface()
            : pimpl(std::make_unique<impl>()) {
        pimpl->q_this = this;
    }

    fracture_robot_interface::~fracture_robot_interface() = default;

}