#pragma once

#include <memory>
#include <thread>
#include <atomic>
#include <mutex>
#include <condition_variable>
#include <chrono>
#include <nlohmann/json.hpp>

#include <unitree/robot/go2/sport/sport_client.hpp>
#include <unitree/robot/go2/obstacles_avoid/obstacles_avoid_client.hpp>
#include <unitree/robot/go2/video/video_client.hpp>
#include <unitree/robot/go2/vui/vui_client.hpp>
#include <unitree/robot/channel/channel_subscriber.hpp>
#include <unitree/idl/go2/SportModeState_.hpp>
#include <unitree/idl/go2/LowState_.hpp>

namespace custom {

enum class RobotMode {
    IDLE,
    SPORT,
    LOW_LEVEL,
    RECOVERY
};

enum class MotionState {
    STOPPED,
    STANDING,
    WALKING,
    RUNNING,
    SITTING,
    LYING,
    DANCING,
    SPECIAL_MOTION
};

struct RobotState {
    RobotMode mode = RobotMode::IDLE;
    MotionState motion_state = MotionState::STOPPED;
    
    // Position and orientation
    double position[3] = {0.0, 0.0, 0.0}; // x, y, z
    double orientation[3] = {0.0, 0.0, 0.0}; // roll, pitch, yaw
    double velocity[3] = {0.0, 0.0, 0.0}; // vx, vy, vyaw
    
    // IMU data
    double imu_acc[3] = {0.0, 0.0, 0.0};
    double imu_gyro[3] = {0.0, 0.0, 0.0};
    double imu_quat[4] = {1.0, 0.0, 0.0, 0.0};
    
    // Battery and system
    double battery_voltage = 0.0;
    double battery_current = 0.0;
    int battery_percentage = 0;
    double temperature = 0.0;
    
    // Motor states (20 motors)
    double motor_positions[20] = {0};
    double motor_velocities[20] = {0};
    double motor_torques[20] = {0};
    double motor_temperatures[20] = {0};
    
    // Status flags
    bool is_connected = false;
    bool emergency_stop = false;
    bool low_battery = false;
    bool overheated = false;
    
    std::chrono::steady_clock::time_point timestamp;
};

struct MotionCommand {
    enum Type {
        VELOCITY,
        POSITION,
        SPECIAL_ACTION,
        GAIT_CHANGE,
        BODY_POSE
    } type;
    
    // Velocity command
    double linear_velocity[3] = {0.0, 0.0, 0.0}; // vx, vy, vz
    double angular_velocity[3] = {0.0, 0.0, 0.0}; // wx, wy, wz
    
    // Position command
    double target_position[3] = {0.0, 0.0, 0.0};
    double target_orientation[3] = {0.0, 0.0, 0.0};
    
    // Body pose
    double body_height = 0.0;
    double body_roll = 0.0;
    double body_pitch = 0.0;
    double body_yaw = 0.0;
    
    // Special actions
    std::string action_name;
    nlohmann::json action_params;
    
    // Gait parameters
    int gait_type = 0;
    double step_height = 0.0;
    
    std::chrono::steady_clock::time_point timestamp;
    double duration = 0.0; // command duration in seconds
};

class Controller {
public:
    using StateCallback = std::function<void(const RobotState&)>;
    using ErrorCallback = std::function<void(const std::string&)>;

    explicit Controller(const std::string& networkInterface);
    ~Controller();

    bool initialize();
    void shutdown();

    // State management
    bool start();
    bool stop();
    bool isRunning() const { return running_; }
    
    RobotState getCurrentState() const;
    void setStateCallback(StateCallback callback);
    void setErrorCallback(ErrorCallback callback);

    // Motion control
    bool executeCommand(const MotionCommand& command);
    bool emergencyStop();
    bool recoveryStand();
    
    // Basic motions
    bool standUp();
    bool standDown();
    bool sit();
    bool lie();
    bool damp();
    
    // Movement
    bool move(double vx, double vy, double vyaw);
    bool moveToPosition(double x, double y, double yaw);
    bool stop();
    
    // Body control
    bool setBodyPose(double roll, double pitch, double yaw, double height);
    bool balanceStand();
    
    // Gait control
    bool switchGait(int gaitType);
    bool setSpeedLevel(int level);
    
    // Special actions
    bool performAction(const std::string& actionName, const nlohmann::json& params = {});
    
    // Dance and tricks
    bool dance1();
    bool dance2();
    bool frontFlip();
    bool backFlip();
    bool hello();
    bool stretch();
    
    // Obstacle avoidance
    bool enableObstacleAvoidance(bool enable);
    bool isObstacleAvoidanceEnabled();
    
    // Video control
    bool startVideo();
    bool stopVideo();
    bool isVideoActive();
    
    // Audio control
    bool playAudio(const std::string& text);
    bool stopAudio();

private:
    void controlLoop();
    void stateUpdateLoop();
    void processCommands();
    
    void updateSportModeState(const void* message);
    void updateLowLevelState(const void* message);
    
    bool validateCommand(const MotionCommand& command);
    void executeMotionCommand(const MotionCommand& command);
    
    void handleError(const std::string& error);
    void checkSafetyLimits();
    
    // SDK clients
    std::unique_ptr<unitree::robot::go2::SportClient> sportClient_;
    std::unique_ptr<unitree::robot::go2::ObstaclesAvoidClient> obstacleClient_;
    std::unique_ptr<unitree::robot::go2::VideoClient> videoClient_;
    std::unique_ptr<unitree::robot::go2::VuiClient> vuiClient_;
    
    // State subscribers
    std::unique_ptr<unitree::robot::ChannelSubscriber<unitree_go::msg::dds_::SportModeState_>> sportStateSubscriber_;
    std::unique_ptr<unitree::robot::ChannelSubscriber<unitree_go::msg::dds_::LowState_>> lowStateSubscriber_;
    
    // Current state
    mutable std::mutex stateMutex_;
    RobotState currentState_;
    
    // Command queue
    std::queue<MotionCommand> commandQueue_;
    std::mutex commandMutex_;
    std::condition_variable commandCondition_;
    
    // Control threads
    std::thread controlThread_;
    std::thread stateThread_;
    std::atomic<bool> running_;
    std::atomic<bool> initialized_;
    
    // Callbacks
    StateCallback stateCallback_;
    ErrorCallback errorCallback_;
    
    // Safety
    std::chrono::steady_clock::time_point lastCommandTime_;
    std::atomic<bool> emergencyStopActive_;
    
    // Network
    std::string networkInterface_;
    
    // Timing
    std::chrono::milliseconds controlPeriod_;
    std::chrono::milliseconds statePeriod_;
};

} // namespace custom