#include "custom_robot.hpp"
#include <thread>
#include <chrono>
#include <functional>
#include <ctime>
#include <random>
#include <sstream>
#include <iomanip>
#include <nlohmann/json.hpp>
#include <unitree/robot/channel/channel_factory.hpp>
#include <unitree/robot/go2/sport/sport_client.hpp>
#include "recharge.hpp"
#include <cmath>

namespace custom {

CustomRobot::CustomRobot()
    : controller_(nullptr)
    , navigation_(nullptr)
    , rsc_(nullptr)
    , mqttClient_(nullptr)
    , running_(false)
    , initialized_(false)
    , continuous_move_enabled_(false)
    , continuous_vx_(0.0f)
    , continuous_vy_(0.0f)
    , continuous_vyaw_(0.0f) {
    config_.loadDefaults();
    
    try {
        unitree::robot::ChannelFactory::Instance()->Init(0, config_.network_interface);
        LOG_INFO("ChannelFactory initialized with interface: " + config_.network_interface);
    } catch (const std::exception& e) {
        LOG_ERROR("Failed to initialize ChannelFactory: " + std::string(e.what()));
    }
}

CustomRobot::~CustomRobot() {
    LOG_INFO("Shutting down CustomRobot");

    if (mqttClient_) {
        mqttClient_->stopMessageProcessor();
        mqttClient_->disconnect();
        mqttClient_.reset();
        LOG_INFO("MQTT client disconnected and cleaned up");
    }
    
    if (controller_) {
        controller_->shutdown();
        controller_.reset();
    }

    // low_controller_ = std::make_unique<LowController>();
    // low_controller_->initialize();
    
    if (rsc_) {
        rsc_.reset();
    }

    if (navigation_) {
        navigation_.reset();
    }
    
    if (recharge_) {
        recharge_.reset();
    }
    
    if (monitor_) {
        monitor_.reset();
    }
    
    try {
        unitree::robot::ChannelFactory::Instance()->Release();
        LOG_INFO("ChannelFactory released");
    } catch (const std::exception& e) {
        LOG_ERROR("Failed to release ChannelFactory: " + std::string(e.what()));
    }
    
    initialized_ = false;
}

bool CustomRobot::initialize() {
    try {
        controller_ = std::make_unique<Controller>();
        controller_->initialize();

        // low_controller_ = std::make_unique<LowController>();
        // low_controller_->initialize();

        navigation_ = std::make_unique<Navigation>();
        navigation_->Init();

        rsc_ = std::make_unique<unitree::robot::go2::RobotStateClient>();
        rsc_->SetTimeout(3.0f);
        rsc_->Init();
        
        recharge_ = std::make_unique<Recharge>();
        recharge_->Init();

        monitor_ = std::make_unique<Monitor>();
        monitor_->initialize();

        if (!initializeMqtt()) {
            LOG_ERROR("Failed to initialize MQTT client");
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Failed to initialize : " + std::string(e.what()));
        return false;
    }
    
    initialized_ = true;
    return true;
}


bool CustomRobot::start() {
    if (!initialized_) {
        LOG_ERROR("Cannot start: not initialized");
        return false;
    }
    
    if (running_) {
        LOG_WARN("CustomRobot already running");
        return true;
    }
    
    std::vector<unitree::robot::go2::ServiceState> serviceList;
    if (!GetServiceList(serviceList)) {
        LOG_ERROR("Failed to get service list");
        return false;
    }
    
    if (!SwitchService("sport_mode", false)) {
        LOG_ERROR("Failed to stop sport_mode service");
        return false;
    }
    
    std::this_thread::sleep_for(std::chrono::seconds(1));
    
    if (!controller_->start()) {
        LOG_ERROR("Failed to start robot controller");
        return false;
    }

    // if (!low_controller_->start()) {
    //     LOG_ERROR("Failed to start low-level controller");
    //     return false;
    // }

    running_ = true;
    
    // Start continuous movement thread
    traj_th_ = unitree::common::CreateRecurrentThreadEx("continuous_move", UT_CPU_ID_NONE, 50, 
        &CustomRobot::continuousMoveLoop, this);
    
    LOG_INFO("CustomRobot started successfully");
    return true;
}


bool CustomRobot::GetServiceList(std::vector<unitree::robot::go2::ServiceState>& serviceList) {
    if (!initialized_ || !rsc_) {
        LOG_ERROR("Robot not initialized or stateClient is null");
        return false;
    }
    
    try {
        int32_t ret = rsc_->ServiceList(serviceList);
        if (ret != 0) {
            LOG_ERROR("Failed to get service list, error code: " + std::to_string(ret));
            return false;
        }
        
        // LOG_INFO("Successfully retrieved service list with " + std::to_string(serviceList.size()) + " services");
        // for (const auto& service : serviceList) {
        //     std::string statusStr = (service.status == 1) ? "ACTIVE" : "INACTIVE";
        //     LOG_INFO("Service: " + service.name + " | Status: " + statusStr + " | Protect: " + 
        //             (service.protect ? "YES" : "NO"));
        // }
        
        return true;
    } catch (const std::exception& e) {
        LOG_ERROR("Exception in getServiceList: " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::SwitchService(const std::string& serviceName, bool enable) {
    if (!initialized_ || !rsc_) {
        LOG_ERROR("Robot not initialized or stateClient is null");
        return false;
    }
    
    try {
        int32_t status;
        int32_t ret = rsc_->ServiceSwitch(serviceName, enable ? 1 : 0, status);
        if (ret != 0) {
            LOG_ERROR("Failed to switch service " + serviceName + ", error code: " + std::to_string(ret));
            return false;
        }
        LOG_INFO("Successfully switched service " + serviceName + " to " + (enable ? "enabled" : "disabled"));
        return true;
    } catch (const std::exception& e) {
        LOG_ERROR("Exception in switchService: " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::SetReportFreq(int32_t interval, int32_t duration) {
    if (!initialized_ || !rsc_) {
        LOG_ERROR("Robot not initialized or stateClient is null");
        return false;
    }
    
    try {
        int32_t ret = rsc_->SetReportFreq(interval, duration);
        if (ret != 0) {
            LOG_ERROR("Failed to set report frequency, error code: " + std::to_string(ret));
            return false;
        }
        LOG_INFO("Successfully set report frequency: interval=" + std::to_string(interval) + ", duration=" + std::to_string(duration));
        return true;
    } catch (const std::exception& e) {
        LOG_ERROR("Exception in setReportFreq: " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::initializeMqtt() {
    try {
        std::string broker = config_.mqtt_broker;
        int port = config_.mqtt_port;
        std::string clientId = generateRandomClientId();  
        
        mqttClient_ = std::make_unique<MqttClient>(broker, port, clientId);
        using namespace std::placeholders;
        mqttClient_->setMessageCb(std::bind(&CustomRobot::onMqttMessage, this, _1, _2));
        mqttClient_->setConnectionCb(std::bind(&CustomRobot::onMqttConnection, this, _1));
        
        mqttClient_->startMessageProcessor();
        
        if (!mqttClient_->connect(config_.mqtt_username, config_.mqtt_password)) {
            LOG_ERROR("Failed to connect to MQTT broker");
            return false;
        }
        
        std::string controlTopic = config_.topic_prefix + "/" + config_.topic_cmd;
        
        mqttClient_->subscribe(controlTopic, 1);
        
        return true;
        
    } catch (const std::exception& e) {
        LOG_ERROR("Exception during MQTT initialization: " + std::string(e.what()));
        return false;
    }
}

void CustomRobot::onMqttMessage(const std::string& topic, const std::string& payload) {
    try {

        nlohmann::json message = nlohmann::json::parse(payload);
        
        std::string controlTopic = config_.topic_prefix + "/" + config_.topic_cmd;
        
        if (topic == controlTopic) {
            processCmd(message);
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error processing MQTT message: " + std::string(e.what()));
    }
}

void CustomRobot::onMqttConnection(bool connected) {
    if (connected) {
        LOG_INFO("MQTT client connected to broker");
    } else {
        LOG_WARN("MQTT client disconnected from broker");
    }
}


void CustomRobot::processCmd(const nlohmann::json& message) {
    try {
        if (!message.contains("type") || !message.contains("cmd")) {
            LOG_ERROR("Invalid command format: missing 'type' or 'cmd' field");
            return;
        }
        
        std::string type = message["type"];
        std::string cmd = message["cmd"];
        
        LOG_INFO("Processing cmd - type: " + type + ", cmd: " + cmd);
        
        bool success = false;
        
        if (type == "oac") {
            success = processOacCmd(cmd, message);
        } else if (type == "sport") {
            success = processSportCmd(cmd, message);
        } else if (type == "rsc") {
            success = processRscCmd(cmd, message);
        } else if (type == "nav") {
            success = processNavCmd(cmd, message);
        } else if (type == "msc") {
            success = processMscCmd(cmd, message);
        } else if (type == "low") {
            success = processLowCmd(cmd, message);
        } else if (type == "recharge") {
            success = processRechargeCmd(cmd, message);
        } else {
            LOG_ERROR("Unknown command type: " + type);
            return;
        }
        LOG_INFO(cmd + ", success: " + std::string(success ? "true" : "false"));
        
    } catch (const std::exception& e) {
        LOG_ERROR("Error processing cmd: " + std::string(e.what()));
    }
}

bool CustomRobot::processLowCmd(const std::string& cmd, const nlohmann::json& message)
{
    // if (!low_controller_)
    // {
    //     LOG_ERROR("LowController not initialized");
    //     return false;
    // }
    //
    // try
    // {
    //     if (cmd == "requestAutoCharge")
    //     {
    //         if (!message.contains("param") || !message["param"].contains("enable"))
    //         {
    //             LOG_ERROR("requestAutoCharge cmd missing 'enable' parameter");
    //             return false;
    //         }
    //         bool enable = message["param"]["enable"];
    //         low_controller_->requestAutoCharge(enable);
    //         return true;
    //     }
    //     else if (cmd == "powerOff")
    //     {
    //         low_controller_->powerOff();
    //         return true;
    //     }
    //     else
    //     {
    //         LOG_ERROR("Unknown Low command: %s", cmd.c_str());
    //         return false;
    //     }
    // }
    // catch (const std::exception& e)
    // {
    //     LOG_ERROR("Error executing Low command %s: %s", cmd.c_str(), e.what());
    //     return false;
    // }
    
    LOG_WARN("LowController functionality is currently disabled");
    return false;
}

bool CustomRobot::processRechargeCmd(const std::string& cmd, const nlohmann::json& message) {
    if (!recharge_) {
        LOG_ERROR("Recharge module not initialized");
        return false;
    }

    try {
        if (cmd == "start") {
            LOG_INFO("Starting recharge process...");
            std::string result = recharge_->StartRecharge();
            LOG_INFO("Recharge process finished with result: " + result);
            return result == "aruco_success";
        } else if (cmd == "stop") {
            LOG_INFO("Stopping recharge process...");
            return recharge_->StopRecharge();
        } else {
            LOG_ERROR("Unknown recharge command: " + cmd);
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error executing recharge command " + cmd + ": " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::processOacCmd(const std::string& cmd, const nlohmann::json& message) {
    if (!controller_) {
        LOG_ERROR("Controller not initialized");
        return false;
    }
    
    try {
        if (cmd == "SwitchSet") {
            if (message.find("param") == message.end() || message["param"].find("enable") == message["param"].end()) {
                LOG_ERROR("Missing 'enable' parameter for Switch command");
                return false;
            }
            bool enable = message["param"]["enable"];
            return controller_->SwitchSet(enable);
        } else if (cmd == "Move") {
            if (!message.contains("param")) {
                LOG_ERROR("Move cmd missing 'param'");
                return false;
            }
            auto param = message["param"];
            return controller_->ObstacleMove(param["x"], param["y"], param["yaw"]);
            
        } else if (cmd == "UseRemoteCommandFromApi") {
            if (!message.contains("param") || !message["param"].contains("enable")) {
                LOG_ERROR("UseRemoteCommandFromApi cmd missing 'enable' parameter");
                return false;
            }
            bool enable = message["param"]["enable"];
            return controller_->UseRemoteCommandFromApi(enable);
            
        } else if (cmd == "MoveToAbsolutePosition") {
            if (!message.contains("param")) {
                LOG_ERROR("MoveToAbsolutePosition cmd missing 'param'");
                return false;
            }
            auto param = message["param"];
            if (!param.contains("x") || !param.contains("y") || !param.contains("yaw")) {
                LOG_ERROR("MoveToAbsolutePosition missing x, y, or yaw parameters");
                return false;
            }
            float x = param["x"];
            float y = param["y"];
            float yaw = param["yaw"];
            return controller_->MoveToAbsolutePosition(x, y, yaw);
            
        } else if (cmd == "MoveToIncrementPosition") {
            if (!message.contains("param")) {
                LOG_ERROR("MoveToIncrementPosition cmd missing 'param'");
                return false;
            }
            auto param = message["param"];
            if (!param.contains("x") || !param.contains("y") || !param.contains("yaw")) {
                LOG_ERROR("MoveToIncrementPosition missing x, y, or yaw parameters");
                return false;
            }
            float x = param["x"];
            float y = param["y"];
            float yaw = param["yaw"];
            return controller_->MoveToIncrementPosition(x, y, yaw);
            
        } else {
            LOG_ERROR("Unknown OAC command: " + cmd);
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error executing OAC command " + cmd + ": " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::processNavCmd(const std::string& cmd, const nlohmann::json& message) {
    if (!navigation_) {
        LOG_ERROR("Navigation not initialized");
        return false;
    }

    try {
        if (cmd == "start_mapping") {
            return navigation_->startMapping();
        } else if (cmd == "end_mapping") {
            if (!message.contains("param") || !message["param"].contains("address")) {
                LOG_ERROR("end_mapping cmd missing 'address' parameter");
                return false;
            }
            std::string address = message["param"]["address"];
            return navigation_->endMapping(address);
        } else if (cmd == "initialize_pose") {
            if (!message.contains("param")) {
                LOG_ERROR("initialize_pose cmd missing 'param'");
                return false;
            }
            auto& param = message["param"];
            const std::vector<std::string> keys = {"x", "y", "z", "q_x", "q_y", "q_z", "q_w", "address"};
            for (const auto& key : keys) {
                if (!param.contains(key)) {
                    LOG_ERROR("initialize_pose missing parameter: " + key);
                    return false;
                }
            }
            return navigation_->initializePose(param["x"], param["y"], param["z"], param["q_x"], param["q_y"], param["q_z"], param["q_w"], param["address"]);
        } else if (cmd == "pose_navigation") {
            if (!message.contains("param")) {
                LOG_ERROR("pose_navigation cmd missing 'param'");
                return false;
            }
            auto& param = message["param"];
            if (!param.contains("targetPose") || !param.contains("mode") || !param.contains("speed")) {
                LOG_ERROR("pose_navigation missing 'targetPose', 'mode', or 'speed'");
                return false;
            }
            auto& targetPose = param["targetPose"];
            const std::vector<std::string> pose_keys = {"x", "y", "z", "q_x", "q_y", "q_z", "q_w"};
             for (const auto& key : pose_keys) {
                if (!targetPose.contains(key)) {
                    LOG_ERROR("pose_navigation.targetPose missing parameter: " + key);
                    return false;
                }
            }
            return navigation_->poseNavigation(targetPose["x"], targetPose["y"], targetPose["z"], targetPose["q_x"], targetPose["q_y"], targetPose["q_z"], targetPose["q_w"], param["mode"], param["speed"]);
        } else if (cmd == "pause_navigation") {
            return navigation_->pauseNavigation();
        } else if (cmd == "resume_navigation") {
            return navigation_->resumeNavigation();
        } else if (cmd == "close_slam") {
            return navigation_->closeSlam();
        } else {
            LOG_ERROR("Unknown Nav command: " + cmd);
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error executing Nav command " + cmd + ": " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::processRscCmd(const std::string& cmd, const nlohmann::json& message) {
    try {
        if (cmd == "list") {
            LOG_INFO("Received service list query command");
            
            std::vector<unitree::robot::go2::ServiceState> serviceList;
            if (!GetServiceList(serviceList)) {
                LOG_ERROR("Failed to get service list");
                return false;
            }
            
            LOG_INFO("Successfully retrieved service list with " + std::to_string(serviceList.size()) + " services");
            
            int filterStatus = -1; 
            if (message.contains("param") && message["param"].contains("status")) {
                filterStatus = message["param"]["status"];
                LOG_INFO("Filtering services by status: " + std::to_string(filterStatus));
            }
            
            printServiceList(serviceList, filterStatus);
            return true;
            
        } else if (cmd == "switch") {
            if (!message.contains("param") || 
                !message["param"].contains("name") || 
                !message["param"].contains("enable")) {
                LOG_ERROR("Service switch command missing required parameters (name, enable)");
                return false;
            }
            
            std::string serviceName = message["param"]["name"];
            bool enable = message["param"]["enable"];
            
            bool result = SwitchService(serviceName, enable);
            if (result) {
                LOG_INFO("Service switch completed successfully");
            }
            return result;
            
        } else {
            LOG_ERROR("Unknown service command: " + cmd);
            return false;
        }
        
    } catch (const std::exception& e) {
        LOG_ERROR("Error processing service cmd: " + std::string(e.what()));
        return false;
    }
}

bool CustomRobot::processMscCmd(const std::string& cmd, const nlohmann::json& message) {
    if (!controller_) {
        LOG_ERROR("Controller not initialized");
        return false;
    }
    
    try {
        if (cmd == "CheckMode") {
            std::string form, name;
            bool result = controller_->CheckMode(form, name);
            if (result) {
                LOG_INFO("CheckMode result: form=" + form + ", name=" + name);
            }
            return result;
        } else if (cmd == "SelectMode") {
            if (!message.contains("param") || !message["param"].contains("name")) {
                LOG_ERROR("SelectMode cmd missing 'name' parameter");
                return false;
            }
            std::string name = message["param"]["name"];
            return controller_->SelectMode(name);
        } else if (cmd == "ReleaseMode") {
            return controller_->ReleaseMode();
        } else {
            LOG_ERROR("Unknown MSC command: " + cmd);
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error executing MSC command " + cmd + ": " + std::string(e.what()));
        return false;
    }
}

void CustomRobot::printServiceList(const std::vector<unitree::robot::go2::ServiceState>& serviceList, int filterStatus) {
    std::string filterDesc = "All Services";
    if (filterStatus == 0) {
        filterDesc = "INACTIVE Services Only";
    } else if (filterStatus == 1) {
        filterDesc = "ACTIVE Services Only";
    }
    
    LOG_INFO("=== Service Status Report (" + filterDesc + ") ===");
    
    int matchCount = 0;
    for (const auto& service : serviceList) {
        if (filterStatus != -1 && service.status != filterStatus) {
            continue;
        }
        std::string statusStr = (service.status == 1) ? "ACTIVE" : "INACTIVE";
        
        LOG_INFO("Service: " + service.name + " | Status: " + statusStr + " | Protect: " + 
                (service.protect ? "YES" : "NO"));
        matchCount++;
    }
    
    if (filterStatus != -1) {
        LOG_INFO("Found " + std::to_string(matchCount) + " services matching filter criteria");
    }
    
    LOG_INFO("=== End Service Report ===");
}


bool CustomRobot::processSportCmd(const std::string& cmd, const nlohmann::json& message) {
    if (!controller_) {
        LOG_ERROR("Controller not initialized");
        return false;
    }
    
    try {
        if (cmd == "Damp") {
            return controller_->Damp();
        } else if (cmd == "BalanceStand") {
            return controller_->BalanceStand();
        } else if (cmd == "StopMove") {
            return controller_->StopMove();
        } else if (cmd == "StandUp") {
            return controller_->StandUp();
        } else if (cmd == "StandDown") {
            return controller_->StandDown();
        } else if (cmd == "RecoveryStand") {
            return controller_->RecoveryStand();
        } else if (cmd == "Euler") {
            return controller_->Euler(message["param"]["roll"], message["param"]["pitch"], message["param"]["yaw"]);
        } else if (cmd == "Move") {
            return controller_->Move(message["param"]["vx"], message["param"]["vy"], message["param"]["vyaw"]);
        } else if (cmd == "Sit") {
            return controller_->Sit();
        } else if (cmd == "RiseSit") {
            return controller_->RiseSit();
        } else if (cmd == "SpeedLevel") {
            return controller_->SpeedLevel(message["param"]["level"]);
        } else if (cmd == "Hello") {
            return controller_->Hello();
        } else if (cmd == "Stretch") {
            return controller_->Stretch();
        } else if (cmd == "SwitchJoystick") {
            return controller_->SwitchJoystick(message["param"]["flag"]);
        } else if (cmd == "Content") {
            return controller_->Content();
        } else if (cmd == "Heart") {
            return controller_->Heart();
        } else if (cmd == "Pose") {
            return controller_->Pose(message["param"]["flag"]);
        } else if (cmd == "Scrape") {
            return controller_->Scrape();
        } else if (cmd == "FrontFlip") {
            return controller_->FrontFlip();
        } else if (cmd == "FrontJump") {
            return controller_->FrontJump();
        } else if (cmd == "FrontPounce") {
            return controller_->FrontPounce();
        } else if (cmd == "Dance1") {
            return controller_->Dance1();
        } else if (cmd == "Dance2") {
            return controller_->Dance2();
        } else if (cmd == "LeftFlip") {
            return controller_->LeftFlip();
        } else if (cmd == "BackFlip") {
            return controller_->BackFlip();
        } else if (cmd == "FreeWalk") {
            return controller_->FreeWalk();
        } else if (cmd == "FreeBound") {
            return controller_->FreeBound(message["param"]["flag"]);
        } else if (cmd == "FreeJump") {
            return controller_->FreeJump(message["param"]["flag"]);
        } else if (cmd == "FreeAvoid") {
            return controller_->FreeAvoid(message["param"]["flag"]);
        } else if (cmd == "WalkUpright") {
            return controller_->WalkUpright(message["param"]["flag"]);
        } else if (cmd == "CrossStep") {
            return controller_->CrossStep(message["param"]["flag"]);
        } else if (cmd == "StartContinuousMove") {
            if (!message.contains("param")) {
                LOG_ERROR("StartContinuousMove cmd missing 'param'");
                return false;
            }
            auto param = message["param"];
            if (!param.contains("vx") || !param.contains("vy") || !param.contains("vyaw")) {
                LOG_ERROR("StartContinuousMove missing vx, vy, or vyaw parameters");
                return false;
            }
            float vx = param["vx"];
            float vy = param["vy"];
            float vyaw = param["vyaw"];
            return startContinuousMove(vx, vy, vyaw);
        } else if (cmd == "StopContinuousMove") {
            return stopContinuousMove();
        } 
        else {
            LOG_ERROR("Unknown Sport command: " + cmd);
            return false;
        }
    } catch (const std::exception& e) {
        LOG_ERROR("Error executing Sport command " + cmd + ": " + std::string(e.what()));
        return false;
    }
}

// Continuous movement control implementation
bool CustomRobot::startContinuousMove(float vx, float vy, float vyaw) {
    if (!controller_) {
        LOG_ERROR("Controller not initialized");
        return false;
    }
    
    continuous_vx_ = vx;
    continuous_vy_ = vy;
    continuous_vyaw_ = vyaw;
    continuous_move_enabled_ = true;
    
    LOG_INFO("Started continuous move with vx=" + std::to_string(vx) + 
             ", vy=" + std::to_string(vy) + ", vyaw=" + std::to_string(vyaw));
    return true;
}

bool CustomRobot::stopContinuousMove() {
    continuous_move_enabled_ = false;
    
    if (controller_) {
        controller_->StopMove();
    }
    
    LOG_INFO("Stopped continuous move");
    return true;
}

void CustomRobot::continuousMoveLoop() {
    if (!continuous_move_enabled_ || !controller_) {
        return;
    }
    
    float vx = continuous_vx_.load();
    float vy = continuous_vy_.load();
    float vyaw = continuous_vyaw_.load();
    
    controller_->Move(vx, vy, vyaw);
}

} // namespace custom