#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UR5 集成测试脚本
用于验证 UR5 机器人在 Unity-MoveIt2 系统中的完整集成
"""

import sys
import time
import yaml
from pathlib import Path

# 设置控制台编码为 UTF-8
if sys.platform == 'win32':
    import os
    os.system('chcp 65001 >nul 2>&1')

# ANSI 颜色代码
class Colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_header(text):
    """打印标题"""
    print(f"\n{Colors.HEADER}{Colors.BOLD}{'='*60}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{text:^60}{Colors.ENDC}")
    print(f"{Colors.HEADER}{Colors.BOLD}{'='*60}{Colors.ENDC}\n")

def print_success(text):
    """打印成功信息"""
    print(f"{Colors.OKGREEN}[OK] {text}{Colors.ENDC}")

def print_error(text):
    """打印错误信息"""
    print(f"{Colors.FAIL}[ERROR] {text}{Colors.ENDC}")

def print_warning(text):
    """打印警告信息"""
    print(f"{Colors.WARNING}[WARN] {text}{Colors.ENDC}")

def print_info(text):
    """打印信息"""
    print(f"{Colors.OKCYAN}[INFO] {text}{Colors.ENDC}")

class UR5IntegrationTest:
    """UR5 集成测试类"""

    def __init__(self):
        self.project_root = Path(__file__).parent.resolve()
        self.ur5_dir = self.project_root / "UR5"
        self.config_dir = self.project_root / "configs" / "robots"
        self.ros2_ws = self.project_root / "ros2-workspace"
        self.test_results = {}

    def test_file_structure(self):
        """测试文件结构完整性"""
        print_header("测试 1: 文件结构完整性检查")

        required_files = {
            "URDF 文件": self.ur5_dir / "urdf" / "ur5.urdf",
            "SRDF 文件": self.ur5_dir / "urdf" / "ur5.srdf",
            "YAML 配置": self.config_dir / "ur5.yaml",
            "ROS2 配置": self.ros2_ws / "src" / "unity_moveit2_bridge" / "config" / "ur5_robot_config.yaml",
        }

        all_passed = True
        for name, path in required_files.items():
            if path.exists():
                print_success(f"{name} 存在: {path.name}")
            else:
                print_error(f"{name} 缺失: {path}")
                all_passed = False

        # 检查 mesh 文件
        mesh_dirs = {
            "可视化网格": self.ur5_dir / "meshes" / "ur5" / "visual",
            "碰撞网格": self.ur5_dir / "meshes" / "ur5" / "collision",
        }

        for name, path in mesh_dirs.items():
            if path.exists():
                mesh_count = len(list(path.glob("*")))
                print_success(f"{name} 目录存在，包含 {mesh_count} 个文件")
            else:
                print_error(f"{name} 目录缺失")
                all_passed = False

        self.test_results['file_structure'] = all_passed
        return all_passed

    def test_urdf_validity(self):
        """测试 URDF 文件有效性"""
        print_header("测试 2: URDF 文件有效性检查")

        urdf_path = self.ur5_dir / "urdf" / "ur5.urdf"

        try:
            with open(urdf_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # 基本 XML 结构检查
            if '<?xml version="1.0"' in content:
                print_success("XML 声明正确")
            else:
                print_warning("缺少 XML 声明")

            # 检查关键元素
            required_elements = [
                ('<robot name="ur5"', "机器人名称"),
                ('<link name="base_link"', "base_link"),
                ('<link name="tool0"', "tool0 末端执行器"),
                ('<joint name="shoulder_pan_joint"', "shoulder_pan_joint"),
                ('<joint name="wrist_3_joint"', "wrist_3_joint"),
            ]

            all_passed = True
            for element, desc in required_elements:
                if element in content:
                    print_success(f"包含 {desc}")
                else:
                    print_error(f"缺少 {desc}")
                    all_passed = False

            # 统计关节数量
            joint_count = content.count('<joint name=')
            link_count = content.count('<link name=')
            print_info(f"检测到 {link_count} 个连杆和 {joint_count} 个关节")

            if joint_count == 7 and link_count == 8:  # 6 个运动关节 + 1 个固定关节 (tool0)
                print_success("关节和连杆数量正确 (6-DOF + tool0)")
            else:
                print_warning(f"期望 7 个关节和 8 个连杆，实际为 {joint_count} 个关节和 {link_count} 个连杆")

            self.test_results['urdf_validity'] = all_passed
            return all_passed

        except Exception as e:
            print_error(f"读取 URDF 文件失败: {e}")
            self.test_results['urdf_validity'] = False
            return False

    def test_yaml_config(self):
        """测试 YAML 配置文件"""
        print_header("测试 3: YAML 配置文件检查")

        yaml_path = self.config_dir / "ur5.yaml"

        try:
            with open(yaml_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 检查关键配置项
            required_sections = {
                'robot_info': '机器人信息',
                'file_paths': '文件路径',
                'joints': '关节配置',
                'kinematics': '运动学配置',
                'planning': '路径规划配置',
                'controllers': '控制器配置',
            }

            all_passed = True
            for section, desc in required_sections.items():
                if section in config:
                    print_success(f"包含 {desc} ({section})")
                else:
                    print_error(f"缺少 {desc} ({section})")
                    all_passed = False

            # 检查关节数量
            if 'joints' in config:
                joint_count = len(config['joints'])
                print_info(f"配置了 {joint_count} 个关节")
                if joint_count == 6:
                    print_success("关节数量正确 (6-DOF)")
                else:
                    print_warning(f"期望 6 个关节，实际为 {joint_count}")

            # 检查运动学配置
            if 'kinematics' in config:
                kin = config['kinematics']
                if kin.get('tip_link') == 'tool0':
                    print_success("末端执行器设置为 tool0")
                else:
                    print_error(f"末端执行器应为 tool0，实际为 {kin.get('tip_link')}")
                    all_passed = False

            self.test_results['yaml_config'] = all_passed
            return all_passed

        except Exception as e:
            print_error(f"读取 YAML 配置失败: {e}")
            self.test_results['yaml_config'] = False
            return False

    def test_ros2_config(self):
        """测试 ROS2 配置文件"""
        print_header("测试 4: ROS2 配置文件检查")

        ros2_config_path = self.ros2_ws / "src" / "unity_moveit2_bridge" / "config" / "ur5_robot_config.yaml"

        try:
            with open(ros2_config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 检查机器人信息
            if config.get('robot', {}).get('name') == 'ur5':
                print_success("机器人名称正确: ur5")
            else:
                print_error(f"机器人名称错误: {config.get('robot', {}).get('name')}")

            # 检查规划组
            if 'planning_groups' in config:
                groups = config['planning_groups']
                if 'manipulator' in groups:
                    print_success("包含 manipulator 规划组")
                    manip = groups['manipulator']

                    # 检查关节列表
                    expected_joints = [
                        "shoulder_pan_joint", "shoulder_lift_joint", "elbow_joint",
                        "wrist_1_joint", "wrist_2_joint", "wrist_3_joint"
                    ]
                    actual_joints = manip.get('joints', [])

                    if actual_joints == expected_joints:
                        print_success(f"关节配置正确 ({len(actual_joints)} 个关节)")
                    else:
                        print_error("关节配置不正确")
                        print_info(f"期望: {expected_joints}")
                        print_info(f"实际: {actual_joints}")
                else:
                    print_error("缺少 manipulator 规划组")

            # 检查工作空间
            if 'workspace' in config:
                ws = config['workspace']
                print_success("包含工作空间配置")
                print_info(f"  X: [{ws.get('min_x')}, {ws.get('max_x')}]")
                print_info(f"  Y: [{ws.get('min_y')}, {ws.get('max_y')}]")
                print_info(f"  Z: [{ws.get('min_z')}, {ws.get('max_z')}]")

            self.test_results['ros2_config'] = True
            return True

        except Exception as e:
            print_error(f"读取 ROS2 配置失败: {e}")
            self.test_results['ros2_config'] = False
            return False

    def test_docker_environment(self):
        """测试 Docker 环境"""
        print_header("测试 5: Docker 环境检查")

        import subprocess

        try:
            # 检查 Docker 容器状态
            result = subprocess.run(
                ['docker', 'ps', '--filter', 'name=unity-ros2-system', '--format', '{{.Status}}'],
                capture_output=True,
                text=True,
                timeout=5
            )

            if result.returncode == 0:
                status = result.stdout.strip()
                if 'Up' in status:
                    print_success(f"Docker 容器运行中: {status}")
                    self.test_results['docker_environment'] = True
                    return True
                else:
                    print_error(f"Docker 容器未运行: {status}")
                    print_info("请运行: cd docker && docker-compose up -d")
                    self.test_results['docker_environment'] = False
                    return False
            else:
                print_error("无法检查 Docker 状态")
                self.test_results['docker_environment'] = False
                return False

        except subprocess.TimeoutExpired:
            print_error("Docker 命令超时")
            self.test_results['docker_environment'] = False
            return False
        except FileNotFoundError:
            print_error("Docker 未安装或不在 PATH 中")
            self.test_results['docker_environment'] = False
            return False
        except Exception as e:
            print_error(f"检查 Docker 环境失败: {e}")
            self.test_results['docker_environment'] = False
            return False

    def print_summary(self):
        """打印测试摘要"""
        print_header("测试摘要")

        total_tests = len(self.test_results)
        passed_tests = sum(1 for result in self.test_results.values() if result)

        print(f"\n总测试数: {total_tests}")
        print(f"通过: {Colors.OKGREEN}{passed_tests}{Colors.ENDC}")
        print(f"失败: {Colors.FAIL}{total_tests - passed_tests}{Colors.ENDC}")
        print(f"通过率: {passed_tests / total_tests * 100:.1f}%\n")

        print("详细结果:")
        for test_name, result in self.test_results.items():
            status = f"{Colors.OKGREEN}[PASS]{Colors.ENDC}" if result else f"{Colors.FAIL}[FAIL]{Colors.ENDC}"
            print(f"  {test_name}: {status}")

        print()

        if passed_tests == total_tests:
            print_success("所有测试通过！UR5 已准备好进行集成测试。")
            return True
        else:
            print_warning("部分测试失败，请检查上述错误信息。")
            return False

    def run_all_tests(self):
        """运行所有测试"""
        print_header("UR5 集成测试开始")
        print_info(f"项目路径: {self.project_root}")
        print_info(f"测试时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")

        # 运行所有测试
        self.test_file_structure()
        self.test_urdf_validity()
        self.test_yaml_config()
        self.test_ros2_config()
        self.test_docker_environment()

        # 打印摘要
        return self.print_summary()

def main():
    """主函数"""
    test = UR5IntegrationTest()
    success = test.run_all_tests()

    # 返回退出码
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()