#!/usr/bin/env python3
"""
工作手机Agent - 语音控制版
运行在手机上，支持语音命令和自主执行任务
"""

import asyncio
import json
import logging
import argparse
import base64
import time
from datetime import datetime
from typing import Optional, Dict, Any, List

try:
    import websockets
    import uiautomator2 as u2
except ImportError:
    print("请安装依赖: pip install websockets uiautomator2")
    exit(1)

# 尝试导入语音识别
try:
    import speech_recognition as sr
    VOICE_ENABLED = True
except ImportError:
    VOICE_ENABLED = False
    print("语音功能未启用，安装: pip install SpeechRecognition")

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
logger = logging.getLogger(__name__)


class VoiceAgent:
    """
    语音控制Agent
    
    功能：
    1. 手机本地运行，无需电脑
    2. 语音命令控制
    3. 自主执行任务队列
    4. 连接远程服务器接收命令
    """
    
    def __init__(
        self,
        device_id: str,
        server_url: str = None,
        voice_enabled: bool = True
    ):
        self.device_id = device_id
        self.server_url = server_url
        self.voice_enabled = voice_enabled and VOICE_ENABLED
        
        # 初始化uiautomator2
        self.d = u2.connect()
        self.d.implicitly_wait(10.0)
        
        # 任务队列
        self.task_queue: List[Dict] = []
        self.running = False
        
        # 语音识别器
        if self.voice_enabled:
            self.recognizer = sr.Recognizer()
        
        # 预设命令
        self.commands = {
            "打开微信": self._open_wechat,
            "打开豆包": self._open_doubao,
            "打开抖音": self._open_douyin,
            "打开设置": self._open_settings,
            "截图": self._screenshot,
            "返回": self._go_back,
            "回到桌面": self._go_home,
            "向上滑": self._swipe_up,
            "向下滑": self._swipe_down,
        }
        
        logger.info(f"VoiceAgent初始化完成: {device_id}")
        logger.info(f"语音控制: {'启用' if self.voice_enabled else '禁用'}")
    
    # ========== 基础操作 ==========
    
    def _open_wechat(self):
        """打开微信"""
        self.d.app_start("com.tencent.mm")
        return {"success": True, "message": "微信已打开"}
    
    def _open_doubao(self):
        """打开豆包"""
        self.d.app_start("com.larus.nova")
        return {"success": True, "message": "豆包已打开"}
    
    def _open_douyin(self):
        """打开抖音"""
        self.d.app_start("com.ss.android.ugc.aweme")
        return {"success": True, "message": "抖音已打开"}
    
    def _open_settings(self):
        """打开设置"""
        self.d.app_start("com.android.settings")
        return {"success": True, "message": "设置已打开"}
    
    def _screenshot(self):
        """截图"""
        img = self.d.screenshot(format='raw')
        # 保存到手机
        path = f"/sdcard/screenshot_{int(time.time())}.png"
        with open(path, 'wb') as f:
            f.write(img)
        return {"success": True, "message": f"截图已保存: {path}"}
    
    def _go_back(self):
        """返回"""
        self.d.press("back")
        return {"success": True, "message": "已返回"}
    
    def _go_home(self):
        """回到桌面"""
        self.d.press("home")
        return {"success": True, "message": "已回到桌面"}
    
    def _swipe_up(self):
        """向上滑动"""
        self.d.swipe_ext("up", scale=0.8)
        return {"success": True, "message": "已向上滑动"}
    
    def _swipe_down(self):
        """向下滑动"""
        self.d.swipe_ext("down", scale=0.8)
        return {"success": True, "message": "已向下滑动"}
    
    # ========== 语音识别 ==========
    
    def listen_voice(self, timeout: int = 5) -> Optional[str]:
        """监听语音命令"""
        if not self.voice_enabled:
            return None
        
        try:
            with sr.Microphone() as source:
                logger.info("请说话...")
                self.recognizer.adjust_for_ambient_noise(source, duration=0.5)
                audio = self.recognizer.listen(source, timeout=timeout)
            
            # 使用Google语音识别（中文）
            text = self.recognizer.recognize_google(audio, language="zh-CN")
            logger.info(f"识别到: {text}")
            return text
        except sr.WaitTimeoutError:
            return None
        except sr.UnknownValueError:
            logger.warning("无法识别语音")
            return None
        except Exception as e:
            logger.error(f"语音识别错误: {e}")
            return None
    
    def process_voice_command(self, text: str) -> Dict:
        """处理语音命令"""
        text = text.strip()
        
        # 精确匹配
        if text in self.commands:
            return self.commands[text]()
        
        # 模糊匹配
        for cmd, func in self.commands.items():
            if cmd in text or text in cmd:
                return func()
        
        # 发送消息命令: "给XXX发消息说XXX"
        if "发消息" in text or "发送" in text:
            return self._parse_send_message(text)
        
        # 打开APP命令: "打开XXX"
        if text.startswith("打开"):
            app_name = text[2:].strip()
            return self._open_app_by_name(app_name)
        
        return {"success": False, "message": f"未识别的命令: {text}"}
    
    def _parse_send_message(self, text: str) -> Dict:
        """解析发消息命令"""
        # 示例: "给张三发消息说你好"
        import re
        match = re.search(r'给(.+?)(发消息|发送)(.+)', text)
        if match:
            contact = match.group(1).strip()
            message = match.group(3).replace("说", "").strip()
            return self._send_wechat_message(contact, message)
        return {"success": False, "message": "无法解析发消息命令"}
    
    def _send_wechat_message(self, contact: str, message: str) -> Dict:
        """发送微信消息"""
        try:
            # 打开微信
            self.d.app_start("com.tencent.mm")
            time.sleep(2)
            
            # 搜索联系人
            if self.d(text="搜索").exists(timeout=3):
                self.d(text="搜索").click()
            time.sleep(0.5)
            
            self.d.send_keys(contact)
            time.sleep(1)
            
            if self.d(text=contact).exists(timeout=3):
                self.d(text=contact).click()
            time.sleep(1)
            
            # 发送消息
            self.d.send_keys(message)
            time.sleep(0.3)
            
            if self.d(text="发送").exists(timeout=3):
                self.d(text="发送").click()
            
            return {"success": True, "message": f"已发送给{contact}: {message}"}
        except Exception as e:
            return {"success": False, "message": str(e)}
    
    def _open_app_by_name(self, name: str) -> Dict:
        """根据名称打开APP"""
        app_map = {
            "微信": "com.tencent.mm",
            "豆包": "com.larus.nova",
            "抖音": "com.ss.android.ugc.aweme",
            "小红书": "com.xingin.xhs",
            "淘宝": "com.taobao.taobao",
            "支付宝": "com.eg.android.AlipayGphone",
            "设置": "com.android.settings",
            "相机": "com.android.camera",
            "浏览器": "com.android.chrome",
        }
        
        package = app_map.get(name)
        if package:
            self.d.app_start(package)
            return {"success": True, "message": f"{name}已打开"}
        
        return {"success": False, "message": f"未知APP: {name}"}
    
    # ========== 任务队列 ==========
    
    def add_task(self, task: Dict):
        """添加任务到队列"""
        task["id"] = f"task_{int(time.time() * 1000)}"
        task["status"] = "pending"
        task["created_at"] = datetime.now().isoformat()
        self.task_queue.append(task)
        logger.info(f"添加任务: {task}")
        return task["id"]
    
    async def process_task_queue(self):
        """处理任务队列"""
        while self.running:
            if self.task_queue:
                task = self.task_queue[0]
                if task["status"] == "pending":
                    task["status"] = "running"
                    logger.info(f"执行任务: {task['id']}")
                    
                    try:
                        result = await self._execute_task(task)
                        task["status"] = "completed"
                        task["result"] = result
                    except Exception as e:
                        task["status"] = "failed"
                        task["error"] = str(e)
                    
                    self.task_queue.pop(0)
            
            await asyncio.sleep(1)
    
    async def _execute_task(self, task: Dict) -> Dict:
        """执行单个任务"""
        task_type = task.get("type")
        params = task.get("params", {})
        
        if task_type == "voice_command":
            return self.process_voice_command(params.get("text", ""))
        
        elif task_type == "click":
            self.d.click(params["x"], params["y"])
            return {"success": True}
        
        elif task_type == "send_message":
            return self._send_wechat_message(
                params.get("contact"),
                params.get("message")
            )
        
        elif task_type == "open_app":
            return self._open_app_by_name(params.get("name"))
        
        return {"success": False, "message": f"未知任务类型: {task_type}"}
    
    # ========== 远程连接 ==========
    
    async def connect_server(self):
        """连接远程服务器"""
        if not self.server_url:
            return
        
        while self.running:
            try:
                logger.info(f"连接服务器: {self.server_url}")
                
                async with websockets.connect(self.server_url) as ws:
                    # 注册
                    await ws.send(json.dumps({
                        "type": "register",
                        "device_id": self.device_id,
                        "capabilities": ["voice", "ui_control"]
                    }))
                    
                    async for message in ws:
                        data = json.loads(message)
                        await self._handle_server_message(data, ws)
                        
            except Exception as e:
                logger.error(f"服务器连接错误: {e}")
                await asyncio.sleep(5)
    
    async def _handle_server_message(self, data: Dict, ws):
        """处理服务器消息"""
        msg_type = data.get("type")
        
        if msg_type == "execute":
            # 添加到任务队列
            task_id = self.add_task(data.get("task", {}))
            await ws.send(json.dumps({
                "type": "task_queued",
                "task_id": task_id
            }))
        
        elif msg_type == "voice_command":
            result = self.process_voice_command(data.get("text", ""))
            await ws.send(json.dumps({
                "type": "result",
                "command_id": data.get("command_id"),
                "result": result
            }))
    
    # ========== 主循环 ==========
    
    async def voice_loop(self):
        """语音监听循环"""
        if not self.voice_enabled:
            return
        
        logger.info("开始语音监听...")
        while self.running:
            text = self.listen_voice(timeout=3)
            if text:
                # 唤醒词检测
                if "小助手" in text or "你好" in text:
                    logger.info("唤醒成功，等待命令...")
                    command = self.listen_voice(timeout=5)
                    if command:
                        result = self.process_voice_command(command)
                        logger.info(f"执行结果: {result}")
            
            await asyncio.sleep(0.1)
    
    async def start(self):
        """启动Agent"""
        self.running = True
        logger.info("VoiceAgent启动")
        
        tasks = [
            self.process_task_queue(),
        ]
        
        if self.server_url:
            tasks.append(self.connect_server())
        
        if self.voice_enabled:
            tasks.append(self.voice_loop())
        
        await asyncio.gather(*tasks)
    
    def stop(self):
        """停止Agent"""
        self.running = False
        logger.info("VoiceAgent停止")


def main():
    parser = argparse.ArgumentParser(description='语音控制Agent')
    parser.add_argument('--device-id', default='voice-agent-001', help='设备ID')
    parser.add_argument('--server', help='远程服务器地址 (ws://xxx:8899/ws/device/xxx)')
    parser.add_argument('--no-voice', action='store_true', help='禁用语音控制')
    
    args = parser.parse_args()
    
    agent = VoiceAgent(
        device_id=args.device_id,
        server_url=args.server,
        voice_enabled=not args.no_voice
    )
    
    try:
        asyncio.run(agent.start())
    except KeyboardInterrupt:
        agent.stop()


if __name__ == "__main__":
    main()