#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
biohead.py — 基于gRPC的仿生头控制脚本
功能：
  1. 控制单个表情参数
  2. 控制多个表情参数
  3. 退出
  4. 自动增减模式：从0到1再到0循环变化
  5. 设置自动模式参数数量
  6. 眨眼模式：执行预设的眨眼动作
  7. 打哈欠模式：模拟打哈欠动作
  8. 说话模式：模拟嘴巴开合动作
"""
import sys
import time
import random
import os
import grpc
from typing import Generator

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 导入gRPC相关模块
from generated.cmvr.api import biohead_command_pb2
from cmvr import CMVRGrpcClient, CMVRErrorCode, FacialExpressionState

def angle_to_normalized(angle):
    """将角度值(0-180)转换为归一化值(0-1)"""
    return max(0.0, min(1.0, angle / 180.0))

def normalized_to_angle(normalized):
    """将归一化值(0-1)转换为角度值(0-180)"""
    return int(normalized * 180)

def init_biohead_client(server_addr="192.168.1.222:50051"):
    """初始化仿生头gRPC客户端"""
    client = CMVRGrpcClient(server_addr)
    if not client.is_connected():
        print("连接服务器失败")
        return None
    
    biohead = client.get_biohead("bio_head")
    if not biohead:
        print("获取仿生头客户端失败")
        client.close()
        return None
        
    return client, biohead

def set_single_expression(biohead):
    """设置单个表情参数"""
    print("\n--- 设置单个表情参数 ---")
    # 表情参数映射表
    param_map = {
        1: "left_eyebrow_outside_y",
        2: "left_eyebrow_inside_y",
        3: "right_eyebrow_outside_y",
        4: "right_eyebrow_inside_y",
        5: "left_eye_upper_lid_y",
        6: "left_eye_lower_lid_y",
        7: "right_eye_upper_lid_y",
        8: "right_eye_lower_lid_y",
        9: "left_eye_ball_x",
        10: "left_eye_ball_y",
        11: "right_eye_ball_x",
        12: "right_eye_ball_y",
        13: "upper_lip_y",
        14: "lower_lip_y",
        15: "jaw_x",
        16: "jaw_y"
    }
    
    print("可用表情参数:")
    for idx, name in param_map.items():
        print(f"{idx}: {name}")
    
    try:
        param_id = int(input("选择参数ID: "))
        if param_id not in param_map:
            print("无效的参数ID")
            return
            
        angle = int(input("输入角度(0-180): "))
        if angle < 0 or angle > 180:
            print("角度必须在0-180之间")
            return
            
        # 转换为归一化值
        value = angle_to_normalized(angle)
        
        # 创建表情对象
        expression = FacialExpressionState()
        # 设置对应的参数
        setattr(expression, param_map[param_id], value)
        
        # 发送表情设置
        result = biohead.set_expression(expression)
        if result == CMVRErrorCode.CMVR_SUCCESS:
            print(f"设置 {param_map[param_id]} 成功，角度: {angle}°，归一化值: {value:.2f}")
        else:
            print(f"设置失败，错误码: {result}")
            
    except ValueError:
        print("输入无效，请输入数字")

def set_multiple_expressions(biohead):
    """设置多个表情参数"""
    print("\n--- 设置多个表情参数 ---")
    # 表情参数映射表
    param_map = {
        1: "left_eyebrow_outside_y",
        2: "left_eyebrow_inside_y",
        3: "right_eyebrow_outside_y",
        4: "right_eyebrow_inside_y",
        5: "left_eye_upper_lid_y",
        6: "left_eye_lower_lid_y",
        7: "right_eye_upper_lid_y",
        8: "right_eye_lower_lid_y",
        9: "left_eye_ball_x",
        10: "left_eye_ball_y",
        11: "right_eye_ball_x",
        12: "right_eye_ball_y",
        13: "upper_lip_y",
        14: "lower_lip_y",
        15: "jaw_x",
        16: "jaw_y"
    }
    
    print("可用表情参数:")
    for idx, name in param_map.items():
        print(f"{idx}: {name}")
    
    try:
        count = int(input("输入要设置的参数数量: "))
        if count <= 0:
            print("数量必须大于0")
            return
            
        # 创建表情对象
        expression = FacialExpressionState()
        
        for i in range(count):
            print(f"\n第{i+1}个参数:")
            param_id = int(input("选择参数ID: "))
            if param_id not in param_map:
                print("无效的参数ID，跳过该参数")
                continue
                
            angle = int(input("输入角度(0-180): "))
            if angle < 0 or angle > 180:
                print("角度必须在0-180之间，跳过该参数")
                continue
                
            # 转换为归一化值
            value = angle_to_normalized(angle)
            # 设置对应的参数
            setattr(expression, param_map[param_id], value)
            print(f"已添加: {param_map[param_id]} = {value:.2f} (对应角度: {angle}°)")
        
        # 发送表情设置
        result = biohead.set_expression(expression)
        if result == CMVRErrorCode.CMVR_SUCCESS:
            print("设置多个表情参数成功")
        else:
            print(f"设置失败，错误码: {result}")
            
    except ValueError:
        print("输入无效，请输入数字")

def auto_mode(biohead, params):
    """自动增减模式：0→1→0，循环变化"""
    print("\n--- 自动增减模式 ---")
    try:
        # 表情参数映射表
        param_map = {
            1: "left_eyebrow_outside_y",
            2: "left_eyebrow_inside_y",
            3: "right_eyebrow_outside_y",
            4: "right_eyebrow_inside_y",
            5: "left_eye_upper_lid_y",
            6: "left_eye_lower_lid_y",
            7: "right_eye_upper_lid_y",
            8: "right_eye_lower_lid_y",
            9: "left_eye_ball_x",
            10: "left_eye_ball_y",
            11: "right_eye_ball_x",
            12: "right_eye_ball_y",
            13: "upper_lip_y",
            14: "lower_lip_y",
            15: "jaw_x",
            16: "jaw_y"
        }
        
        # 过滤有效的参数ID
        valid_params = [p for p in params if p in param_map]
        if not valid_params:
            print("没有有效的参数，返回菜单")
            return
            
        print(f"将自动控制以下参数: {[param_map[p] for p in valid_params]}")
        
        frequency = 50.0  # 50Hz
        interval = 1.0 / frequency
        value = 0.0
        direction = 0.02  # 步长
        
        # 启动流式控制
        result = biohead.start_stream()
        if result != CMVRErrorCode.CMVR_SUCCESS:
            print(f"开始流式会话失败，错误码: {result}")
            return
            
        print("启动自动模式，按Ctrl+C停止")
        while True:
            # 创建protobuf消息
            proto_expr = biohead_command_pb2.FacialExpression()
            
            # 设置所有参数
            for param_id in valid_params:
                param_name = param_map[param_id]
                
                # 根据参数类型设置到不同的字段
                if param_name in ["left_eyebrow_outside_y", "left_eyebrow_inside_y", 
                                 "right_eyebrow_outside_y", "right_eyebrow_inside_y"]:
                    setattr(proto_expr.eyebrow, param_name, value)
                elif param_name in ["left_eye_upper_lid_y", "left_eye_lower_lid_y",
                                   "right_eye_upper_lid_y", "right_eye_lower_lid_y"]:
                    setattr(proto_expr.eyelid, param_name, value)
                elif param_name in ["left_eye_ball_x", "left_eye_ball_y",
                                   "right_eye_ball_x", "right_eye_ball_y"]:
                    setattr(proto_expr.eye_ball, param_name, value)
                elif param_name in ["upper_lip_y", "lower_lip_y"]:
                    setattr(proto_expr.mouth, param_name, value)
                elif param_name in ["jaw_x", "jaw_y"]:
                    setattr(proto_expr.jaw, param_name, value)
            
            # 发送流式表情
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"流式控制失败，错误码: {result}")
                break
            
            # 更新值
            value += direction
            if value >= 1.0:
                value = 1.0
                direction = -direction
            elif value <= 0.0:
                value = 0.0
                direction = -direction
                
            time.sleep(interval)
            
    except KeyboardInterrupt:
        print("\n自动模式已停止")
    finally:
        # 结束流式会话
        biohead.end_stream()

def set_auto_params():
    """设置自动模式的参数列表"""
    print("\n--- 设置自动模式参数 ---")
    # 表情参数映射表
    param_map = {
        1: "left_eyebrow_outside_y",
        2: "left_eyebrow_inside_y",
        3: "right_eyebrow_outside_y",
        4: "right_eyebrow_inside_y",
        5: "left_eye_upper_lid_y",
        6: "left_eye_lower_lid_y",
        7: "right_eye_upper_lid_y",
        8: "right_eye_lower_lid_y",
        9: "left_eye_ball_x",
        10: "left_eye_ball_y",
        11: "right_eye_ball_x",
        12: "right_eye_ball_y",
        13: "upper_lip_y",
        14: "lower_lip_y",
        15: "jaw_x",
        16: "jaw_y"
    }
    
    print("可用表情参数:")
    for idx, name in param_map.items():
        print(f"{idx}: {name}")
    
    try:
        count = int(input("输入要自动控制的参数数量: "))
        if count <= 0:
            print("数量必须大于0")
            return None
            
        params = []
        for i in range(count):
            param_id = int(input(f"第{i+1}个参数ID: "))
            if param_id not in param_map:
                print("无效的参数ID，跳过")
                continue
            params.append(param_id)
            
        if not params:
            print("未设置有效的参数")
            return None
            
        print(f"已设置自动模式参数: {[param_map[p] for p in params]}")
        return params
        
    except ValueError:
        print("输入无效，请输入数字")
        return None

def blink_mode(biohead):
    """眨眼模式：执行预设的眨眼动作"""
    print("\n--- 眨眼模式 ---")
    try:
        blink_count = int(input("输入眨眼次数(默认10次): ") or "10")
        interval = float(input("输入间隔时间(秒，默认0.1): ") or "0.1")
        
        # 启动流式控制
        result = biohead.start_stream()
        if result != CMVRErrorCode.CMVR_SUCCESS:
            print(f"开始流式会话失败，错误码: {result}")
            return
            
        print(f"启动眨眼模式，共{blink_count}次，间隔{interval}秒")
        
        for i in range(blink_count):
            print(f"第{i+1}次眨眼...")
            
            # 闭眼动作
            proto_expr = biohead_command_pb2.FacialExpression()
            proto_expr.eyelid.left_upper_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.left_lower_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.right_upper_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.right_lower_lid_y = angle_to_normalized(90)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"闭眼动作失败，错误码: {result}")
                break
                
            time.sleep(interval)
            
            # 睁眼动作
            proto_expr = biohead_command_pb2.FacialExpression()
            proto_expr.eyelid.left_upper_lid_y = angle_to_normalized(50)
            proto_expr.eyelid.left_lower_lid_y = angle_to_normalized(130)
            proto_expr.eyelid.right_upper_lid_y = angle_to_normalized(130)
            proto_expr.eyelid.right_lower_lid_y = angle_to_normalized(60)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"睁眼动作失败，错误码: {result}")
                break
                
            time.sleep(interval)
            
        print("眨眼模式结束")
        
    except KeyboardInterrupt:
        print("\n眨眼模式已停止")
    except ValueError:
        print("输入无效，请输入数字")
    finally:
        # 结束流式会话
        biohead.end_stream()

def yawn_mode(biohead):
    """打哈欠模式：模拟打哈欠动作"""
    print("\n--- 打哈欠模式 ---")
    try:
        yawn_count = int(input("输入打哈欠次数(默认5次): ") or "5")
        interval = float(input("输入间隔时间(秒，默认1): ") or "1")
        
        # 启动流式控制
        result = biohead.start_stream()
        if result != CMVRErrorCode.CMVR_SUCCESS:
            print(f"开始流式会话失败，错误码: {result}")
            return
            
        print(f"启动打哈欠模式，共{yawn_count}次，间隔{interval}秒")
        
        for i in range(yawn_count):
            print(f"第{i+1}次打哈欠...")
            
            # 打哈欠动作
            proto_expr = biohead_command_pb2.FacialExpression()
            # 眉毛
            proto_expr.eyebrow.left_eyebrow_outside_y = angle_to_normalized(70)
            proto_expr.eyebrow.right_eyebrow_outside_y = angle_to_normalized(110)
            # 眼睑
            proto_expr.eyelid.left_upper_lid_y = angle_to_normalized(60)
            proto_expr.eyelid.left_lower_lid_y = angle_to_normalized(130)
            proto_expr.eyelid.right_upper_lid_y = angle_to_normalized(120)
            proto_expr.eyelid.right_lower_lid_y = angle_to_normalized(80)
            # 嘴巴和下巴
            proto_expr.mouth.upper_lip_y = angle_to_normalized(120)
            proto_expr.mouth.lower_lip_y = angle_to_normalized(45)
            proto_expr.jaw.x = angle_to_normalized(60)
            proto_expr.jaw.y = angle_to_normalized(120)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"打哈欠动作失败，错误码: {result}")
                break
                
            time.sleep(interval)
            
            # 恢复动作
            proto_expr = biohead_command_pb2.FacialExpression()
            # 恢复到正常状态
            proto_expr.eyebrow.left_eyebrow_outside_y = angle_to_normalized(90)
            proto_expr.eyebrow.right_eyebrow_outside_y = angle_to_normalized(90)
            proto_expr.eyelid.left_upper_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.left_lower_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.right_upper_lid_y = angle_to_normalized(90)
            proto_expr.eyelid.right_lower_lid_y = angle_to_normalized(90)
            proto_expr.mouth.upper_lip_y = angle_to_normalized(90)
            proto_expr.mouth.lower_lip_y = angle_to_normalized(90)
            proto_expr.jaw.x = angle_to_normalized(90)
            proto_expr.jaw.y = angle_to_normalized(90)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"恢复动作失败，错误码: {result}")
                break
                
            time.sleep(interval)
            
        print("打哈欠模式结束")
        
    except KeyboardInterrupt:
        print("\n打哈欠模式已停止")
    except ValueError:
        print("输入无效，请输入数字")
    finally:
        # 结束流式会话
        biohead.end_stream()

def talk_mode(biohead):
    """说话模式：模拟嘴巴开合动作"""
    print("\n--- 说话模式 ---")
    try:
        talk_count = int(input("输入说话循环次数(默认50次): ") or "50")
        
        # 启动流式控制
        result = biohead.start_stream()
        if result != CMVRErrorCode.CMVR_SUCCESS:
            print(f"开始流式会话失败，错误码: {result}")
            return
            
        print(f"启动说话模式，共{talk_count}次")
        
        for i in range(talk_count):
            print(f"第{i+1}次说话...")
            
            # 开口动作
            proto_expr = biohead_command_pb2.FacialExpression()
            proto_expr.mouth.upper_lip_y = angle_to_normalized(120)
            proto_expr.mouth.lower_lip_y = angle_to_normalized(45)
            proto_expr.jaw.x = angle_to_normalized(80)
            proto_expr.jaw.y = angle_to_normalized(110)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"开口动作失败，错误码: {result}")
                break
                
            # 随机间隔
            interval = random.uniform(0, 0.3)
            time.sleep(interval)
            
            # 闭口动作
            proto_expr = biohead_command_pb2.FacialExpression()
            proto_expr.mouth.upper_lip_y = angle_to_normalized(120)
            proto_expr.mouth.lower_lip_y = angle_to_normalized(45)
            proto_expr.jaw.x = angle_to_normalized(90)
            proto_expr.jaw.y = angle_to_normalized(90)
            
            result = biohead.stream_expression(proto_expr)
            if result != CMVRErrorCode.CMVR_SUCCESS:
                print(f"闭口动作失败，错误码: {result}")
                break
                
            # 随机间隔
            interval = random.uniform(0, 0.3)
            time.sleep(interval)
            
        print("说话模式结束")
        
    except KeyboardInterrupt:
        print("\n说话模式已停止")
    except ValueError:
        print("输入无效，请输入数字")
    finally:
        # 结束流式会话
        biohead.end_stream()

def main():
    # 初始化gRPC客户端
    server_addr = "192.168.1.222:50051"
    if len(sys.argv) > 1:
        server_addr = sys.argv[1]
        
    print(f"连接到仿生头服务器: {server_addr}")
    client, biohead = init_biohead_client(server_addr)
    if not client or not biohead:
        sys.exit(1)
        
    # 默认自动模式参数
    auto_params = [5, 6, 7, 8]  # 默认控制眼睑参数
    
    try:
        while True:
            print("\n=== 仿生头控制菜单 ===")
            print("1. 控制单个表情参数")
            print("2. 控制多个表情参数")
            print("3. 退出")
            print("4. 自动增减模式(50Hz)")
            print("5. 设置自动模式参数")
            print("6. 眨眼模式")
            print("7. 打哈欠模式")
            print("8. 说话模式")
            opt = input("选择操作: ").strip()
            
            if opt == '1':
                set_single_expression(biohead)
            elif opt == '2':
                set_multiple_expressions(biohead)
            elif opt == '3':
                print("退出程序")
                break
            elif opt == '4':
                auto_mode(biohead, auto_params)
            elif opt == '5':
                new_params = set_auto_params()
                if new_params:
                    auto_params = new_params
            elif opt == '6':
                blink_mode(biohead)
            elif opt == '7':
                yawn_mode(biohead)
            elif opt == '8':
                talk_mode(biohead)
            else:
                print("无效选项，请重试")
                
    except KeyboardInterrupt:
        print("\n用户中断，退出程序")
    finally:
        client.close()
        print("已关闭连接")

if __name__ == '__main__':
    main()