import time
import threading
import numpy as np
import cv2
import netifaces
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.go2.sport.sport_client import SportClient
from unitree_sdk2py.go2.video.video_client import VideoClient
from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_


error_code = {
    100: "灵动",
    1001: "阻尼",
    1002: "站立锁定",
    1004: "蹲下",
    2006: "蹲下",
    1006: "打招呼/伸懒腰/舞蹈/拜年/比心/开心",
    1007: "坐下",
    1008: "前跳",
    1009: "扑人",
    1013: "平衡站立",
    1015: "常规行走",
    1016: "常规跑步",
    1017: "常规续航",
    1091: "摆姿势",
    2004: "翻身",   # ???
    2007: "闪避",
    2008: "并腿跑",
    2009: "跳跃跑",
    2010: "经典",
    2011: "倒立",
    2012: "前空翻",
    2013: "后空翻",
    2014: "左空翻",
    2016: "交叉步",
    2017: "直立",
    2019: "牵引",
}


class Go2:
    """
    宇树Go2 运动控制封装类
    适用于：Unitree SDK2 Python接口
    """
    def __init__(self, interface=None, timeout=20.0):
        """
        初始化控制器
        :param interface: 网卡接口名称
        :param timeout: 超时时间（秒）
        """
        self.interface = interface
        self.timeout = timeout
        self.sport_client = None
        self.video_client = None
        self.cap = None
        self.error_code = 0  # 状态码

        if self.interface is None:
            self.get_interface()

    def get_interface(self):
        """获取当前接口"""
        
        interfaces = netifaces.interfaces()
        
        # 优先寻找已知的Go2接口
        go2_interfaces = ["enx00e0986113a6"]  # pipstream.py中使用的接口
        
        for go2_iface in go2_interfaces:
            if go2_iface in interfaces:
                self.interface = go2_iface
                print(f"找到Go2接口: {self.interface}")
                return
        
        # 如果没找到，寻找其他en开头的接口
        for iface in interfaces:
            if iface.startswith("en"):
                self.interface = iface
                print(f"使用网络接口: {self.interface}")
                break
        
        if self.interface is None:
            print("警告: 未找到合适的网络接口")
            print("可用接口:", interfaces)
        



    def init(self):
        """初始化与Go2的连接"""
        
        ChannelFactoryInitialize(0, self.interface)

        # 启动状态订阅
        self.sub_state(self.callback)

        # 初始化运动控制客户端
        self.sport_client = SportClient()
        self.sport_client.SetTimeout(self.timeout)
        self.sport_client.Init()
        # 初始化视频流客户端
        # self.video_client = VideoClient()
        # self.video_client.SetTimeout(self.timeout)
        # self.video_client.Init()
        self.cap = self.open_video()
        

    def open_video(self, width: int = 480, height: int = 320):
        """打开视频流"""
        gstreamer_str = (
            f"udpsrc address=230.1.1.1 port=1720 multicast-iface={self.interface} "
            "! application/x-rtp, media=video, encoding-name=H264 "
            "! rtph264depay ! h264parse "
            "! avdec_h264 "  # 解码H.264
            "! videoscale "  # 添加缩放元素，用于调整分辨率
            f"! video/x-raw,width={width},height={height} "  # 目标分辨率
            "! videoconvert ! video/x-raw, format=BGR "  # 转换为OpenCV支持的BGR格式
            "! appsink drop=1"
        )
        cap = cv2.VideoCapture(gstreamer_str, cv2.CAP_GSTREAMER)
        if not cap.isOpened():
            print("视频流打开失败")
            return
        return cap

    def read_image(self):
        """从视频流获取一帧图像"""
        ret, frame = self.cap.read()
        if not ret:
            print("读取图像失败")
            return None
        return frame

    def open_video_pipe(self, width: int = 1280, height: int = 720):
        """使用命名管道方式打开视频流（适用于OpenCV GStreamer支持不完整的情况）"""
        import os
        import time
        
        pipe_path = "/tmp/go2_video_pipe"
        
        try:
            # 清理旧管道
            if os.path.exists(pipe_path):
                os.remove(pipe_path)
            
            # 创建命名管道
            os.mkfifo(pipe_path)
            print(f"创建视频管道: {pipe_path}")
            
            # 构建GStreamer命令（参考pipstream.py的成功实现）
            gst_cmd = (
                f"exec gst-launch-1.0 udpsrc address=230.1.1.1 port=1720 multicast-iface={self.interface} "
                "! application/x-rtp, media=video, encoding-name=H264 ! rtph264depay ! h264parse "
                "! avdec_h264 ! videoconvert ! video/x-raw,format=BGR,width={width},height={height} "
                f"! filesink location={pipe_path} > /dev/null 2>&1 &"
            )
            
            print("启动GStreamer视频流...")
            
            # 使用os.system后台运行（与pipstream.py一致）
            os.system(gst_cmd)
            
            # 等待GStreamer启动
            time.sleep(3)
            
            print("GStreamer视频流启动成功")
            return pipe_path, (width, height)
                
        except Exception as e:
            print(f"创建视频管道失败: {e}")
            # 清理
            if os.path.exists(pipe_path):
                os.remove(pipe_path)
            return None, None
    
    def read_image_pipe(self, pipe_path=None, size=None):
        """从命名管道读取图像（参考pipstream.py实现）"""
        import os
        import numpy as np
        
        if pipe_path is None or size is None:
            pipe_path, size = self.open_video_pipe()
            if pipe_path is None:
                return None
        
        width, height = size
        frame_size = width * height * 3
        
        try:
            with open(pipe_path, 'rb') as pipe:
                frame_data = pipe.read(frame_size)
                
                if len(frame_data) == frame_size:
                    # 转换为图像（与pipstream.py一致）
                    frame = np.frombuffer(frame_data, dtype=np.uint8).reshape((height, width, 3))
                    return frame
                else:
                    # 数据不完整时不打印错误，可能是正常的等待状态
                    return None
                    
        except Exception as e:
            # 管道读取异常时不打印错误，可能是正常的等待状态
            return None
    
    def close_video_pipe(self, pipe_path=None):
        """关闭视频管道（参考pipstream.py实现）"""
        import os
        
        # 终止GStreamer进程（与pipstream.py一致）
        os.system("pkill -f gst-launch")
        
        # 清理管道文件
        if pipe_path and os.path.exists(pipe_path):
            os.remove(pipe_path)
        
        print("视频管道已关闭")
    
    def get_video_reader(self, method='auto'):
        """
        获取视频读取器，自动选择最佳方法
        
        Args:
            method: 'auto', 'opencv', 'pipe'
            
        Returns:
            读取器函数和清理函数
        """
        if method == 'auto':
            # 尝试OpenCV方法
            try:
                cap = self.open_video()
                if cap and cap.isOpened():
                    print("使用OpenCV方法读取视频")
                    return (lambda: self.read_image()), (lambda: None)
            except:
                pass
            
            # 回退到管道方法
            print("OpenCV方法失败，使用管道方法")
            pipe_path, size = self.open_video_pipe()
            if pipe_path:
                return (lambda: self.read_image_pipe(pipe_path, size)), (lambda: self.close_video_pipe(pipe_path))
        
        elif method == 'opencv':
            cap = self.open_video()
            if cap and cap.isOpened():
                return (lambda: self.read_image()), (lambda: None)
        
        elif method == 'pipe':
            pipe_path, size = self.open_video_pipe()
            if pipe_path:
                return (lambda: self.read_image_pipe(pipe_path, size)), (lambda: self.close_video_pipe(pipe_path))
        
        print("无法创建视频读取器")
        return None, None

    
    def sub_state(self, callback, queue_size: int = 5):
        subscriber = ChannelSubscriber("rt/sportmodestate", SportModeState_)
        subscriber.Init(callback, queue_size)

    def callback(self, msg):
        self.error_code = msg.error_code


    # def read_image(self):
    #     """从视频流获取一帧图像"""
    #     code, data = self.video_client.GetImageSample()
    #     if code != 0 or data is None:
    #         print("获取图像样本失败，错误码:", code)
    #         return None
    #     image_data = np.frombuffer(bytes(data), dtype=np.uint8)
    #     image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    #     return image

    def Damp(self):
        """进入阻尼状态。"""
        pass

    def BalanceStand(self):
        """解除锁定。"""
        pass

    def StopMove(self):
        """停下当前动作，将绝大多数指令恢复成默认值。"""
        pass

    def _call(self, func, *args, **kwargs):
        """在线程中调用运动控制函数"""
        def fun_thread():
            ret = func(*args, **kwargs)
            print(f"{func.__name__} 执行结果:", ret)
        
        t = threading.Thread(target=fun_thread)
        t.start()



    def StandUp(self):
        """
        关节锁定，站高。
        执行后状态: 1002:站立锁定
        """
        # 执行前判断状态
        if self.error_code in [1002]:  # 1002 :站立锁定
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1001, 1007, 1013]:  # 100 :灵动, 1001 :阻尼, 1013 :平衡站立
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        self._call(self.sport_client.StandUp)



    def StandDown(self):
        """
        关节锁定，站低。
        执行后状态: 1001:阻尼
        """
        # 执行前判断状态
        if self.error_code in [1001, 1004, 2006]:  # 1001:阻尼 1004 :蹲下 2006 :蹲下
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 100 :灵动, 1001 :阻尼, 1013 :平衡站立
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.StandDown)

    def RecoveryStand(self):
        """	恢复站立。"""
        pass

    def Euler(self, roll, pitch, yaw):
        """站立和行走时的姿态。"""
        pass

    def Move(self, vx, vy, vyaw):
        """移动。"""
        pass

    def Sit(self):
        """
        坐下。
        执行后状态: 1007: 坐下
        """
        # 执行前判断状态
        if self.error_code in [1007]:  # 1007 : 坐下
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 100 :灵动, 1001 :阻尼, 1013 :平衡站立
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Sit)

    def RiseSit(self):
        """
        站起（相对于坐下）。
        执行后状态: 
        """
        # 执行前判断状态
        if self.error_code in [1007]:  # 1002 :站立锁定
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1007, 1013]:  # 100 :灵动, 1007 : 坐下, 1013 :平衡站立
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.RiseSit)



    def SpeedLevel(self, level: int):
        """设置速度档位。"""
        pass

    def Hello(self):
        """
        打招呼
        执行后状态: 1013:平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return

        # 执行指令
        self._call(self.sport_client.Hello)

    def Stretch(self):
        """
        伸懒腰。
        执行后状态: 1013:平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 站立且空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Stretch)



    def Content(self):
        """
        开心。
        执行后状态: 1013:平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Content)


    def Heart(self):
        """
        比心。
        执行后状态: 1013:平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Heart)


    def Pose(self, flag):
        """摆姿势。"""
        pass

    def Scrape(self):
        """ 
        拜年作揖。
        执行后状态: 1013:平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Scrape)



    def FrontJump(self):
        """ 
        前跳。
        执行后状态: 1013: 平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1008]:  # 1008 : 前跳
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.FrontJump)



    def FrontPounce(self):
        """ 
        向前扑人。
        执行后状态: 1013: 平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1009]:  # 1009 : 扑人
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.FrontPounce)


    def Dance1(self):
        """
        舞蹈段落1。
        执行后状态: 1013: 平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Dance1)
    
    def Dance2(self):
        """
        舞蹈段落2。
        执行后状态: 1013: 平衡站立
        """
        # 执行前判断状态
        if self.error_code in [1006]:  # 1006 :打招呼/伸懒腰/舞蹈/拜年/比心/开心
            print("当前状态:", self.error_code, error_code[self.error_code], "无需执行")
            return
        if self.error_code not in [100, 1002, 1013]:  # 空闲状态
            print("繁忙中,当前状态:", self.error_code, error_code[self.error_code])
            return
        
        # 执行指令
        self._call(self.sport_client.Dance2)
    
    def HandStand(self, flag: int):
        """倒立行走。"""
        pass

    def LeftFlip(self):
        """左空翻。"""
        pass

    def BackFlip(self):
        """后空翻。"""
        pass

    def FreeWalk(self, flag: int):
        """ 灵动模式（默认步态）。"""
        pass

    def FreeBound(self, flag: int):
        """ 并腿跑模式。"""
        pass

    def FreeJump(self, flag: int):
        """ 跳跃模式。"""
        pass

    def FreeAvoid(self, flag: int):
        """ 闪避模式。"""
        pass

    def WalkUpright(self, flag: int):
        """ 后腿直立模式。"""
        pass

    def CrossStep(self, flag: int):
        """ 交叉步模式。"""
        pass

    def AutoRecoverSet(self, flag: int):
        """ 设置自动翻身是否生效。"""
        pass

    def AutoRecoverGet(self):
        """ 查询自动翻身是否生效。"""
        pass

    def ClassicWalk(self, flag: int):
        """ 经典步态。"""
        pass

    def TrotRun(self):
        """ 进入常规跑步模式 """
        pass

    def StaticWalk(self):
        """ 进入常规行走模式"""
        pass

    def EconomicGait(self):
        """ 进入常规续航模式 """
        pass

    def SwitchAvoidMode(self):
        """ 闪避模式下，关闭摇杆未推时前方障碍物的闪避以及后方的障碍物躲避"""
        pass




# -------------------- 测试 --------------------
if __name__ == "__main__":
    interface = "enx00e0986113a6"  # 替换为你的Go2网卡接口名称
    
    go2 = Go2(interface=interface)

    go2.stand_down()
    go2.stand_up()
