"""
使用UBOX SDK的Agent能力示例

这个示例展示了如何使用UBOX SDK的Agent功能。
在创建UBox时传入模型配置，然后直接使用 device.explore("task") 执行任务。
"""
import base64
import os
import traceback
from datetime import datetime

from ubox_py_sdk import UBox
from ubox_py_sdk.models import OSType, RunMode
from examples.config import get_ubox_config, get_device_config
from ubox_py_sdk.phone_agent import ModelConfig, AgentConfig

# 从配置文件获取UBox配置
ubox_config = get_ubox_config()
device_config = get_device_config()

if __name__ == "__main__":
    try:
        # 1. 配置AI模型（使用本地或远程的LLM服务）
        model_config = ModelConfig(
            # base_url="https://open.bigmodel.cn/api/paas/v4",
            # api_key="1fb4297e4073481b99f40822c1f11e28.0PvKvroIJCa5UWst",
            # model_name="AutoGLM-Phone",
            base_url="https://ai.gitee.com/v1",
            api_key="MCYWIH6IBORZEV0UE7DWLTMQJEPBMAQBTHLEV51A",
            model_name="AutoGLM-Phone-9B-Multilingual",
            # base_url="http://v2.open.venus.oa.com/llmproxy",
            # api_key="MZOSTCj1owKwsMzepjRuJu67@2057",
            # model_name="server:272328",
        )

        # 2. 配置Agent（可选）
        agent_config = AgentConfig(
            max_steps=100,  # 最大执行步数
            verbose=False,  # 是否显示详细日志
        )
        ubox = UBox(
            mode=ubox_config.get("mode", RunMode.NORMAL),
            base_url=ubox_config.get("base_url", ''),
            secret_id=ubox_config.get('secret_id'),
            secret_key=ubox_config.get('secret_key'),
            log_level="debug",
            log_to_file=True,
            agent_config=agent_config,
            agent_model_config=model_config
        )

        print("\n正在初始化设备...")
        device = ubox.init_device(
            udid=device_config['default_udid'],
            os_type=OSType(device_config['default_os_type']),
            auth_code=device_config.get('auth_code', None),
            # force_proxy=True if ubox.mode == RunMode.NORMAL else False
        )
        print(f"设备初始化成功: {device.udid}")
        print(f"设备类型: {device.os_type.value}")
        print(f"Debug ID: {getattr(device, 'debugId', 'N/A')}")
        
        # 创建截图保存目录
        screenshot_dir = f"screenshots_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        os.makedirs(screenshot_dir, exist_ok=True)
        print(f"截图将保存到目录: {screenshot_dir}")
        
        def save_screenshot(base64_str: str, filepath: str) -> bool:
            """保存base64截图到文件。"""
            if not base64_str:
                return False
            try:
                # 处理base64字符串，移除前缀
                if base64_str.startswith('data:image'):
                    base64_str = base64_str.split(',', 1)[1]
                
                # 解码并保存
                img_data = base64.b64decode(base64_str)
                os.makedirs(os.path.dirname(filepath), exist_ok=True)
                with open(filepath, "wb") as f:
                    f.write(img_data)
                return True
            except Exception as e:
                print(f"保存截图失败 {filepath}: {e}")
                return False
        
        # 执行任务并处理更新
        total_prompt_tokens = 0
        total_completion_tokens = 0
        total_tokens = 0
        # 推理时长统计
        total_inference_time = 0.0
        total_time_to_first_token = 0.0
        total_time_to_thinking_end = 0.0
        step_count = 0
        task = """
                软件名称：腾讯视频
                包名：com.tencent.qqlive
                任务引导：针对个人中心的vip功能做详细测试
                """
        for update in device.explore("打开相册"):
            if update.update_type == "thinking":
                print(update.content, end="", flush=True)
            elif update.update_type == "model_complete":
                # 模型处理完成，显示token统计和推理时长
                content = update.content
                if isinstance(content, dict):
                    # 统计token
                    tokens = content.get("tokens")
                    if tokens:
                        step_prompt = tokens.get("prompt_tokens", 0) or 0
                        step_completion = tokens.get("completion_tokens", 0) or 0
                        step_total = tokens.get("total_tokens", 0) or 0
                        total_prompt_tokens += step_prompt
                        total_completion_tokens += step_completion
                        total_tokens += step_total
                    
                    # 统计推理时长
                    timing = content.get("timing")
                    timing_info = []
                    if timing:
                        if timing.get("time_to_first_token") is not None:
                            ttft = timing["time_to_first_token"]
                            total_time_to_first_token += ttft
                            timing_info.append(f"首Token延迟={ttft:.3f}s")
                        if timing.get("time_to_thinking_end") is not None:
                            tte = timing["time_to_thinking_end"]
                            total_time_to_thinking_end += tte
                            timing_info.append(f"思考完成={tte:.3f}s")
                        if timing.get("total_time") is not None:
                            total_time = timing["total_time"]
                            total_inference_time += total_time
                            timing_info.append(f"总时长={total_time:.3f}s")
                        step_count += 1
                    
                    # 打印统计信息
                    info_parts = []
                    if tokens:
                        info_parts.append(f"Token: 输入={step_prompt}, 输出={step_completion}, 总计={step_total}")
                    if timing_info:
                        info_parts.append(" | ".join(timing_info))
                    
                    if info_parts:
                        print(f"\n模型处理完成 - {' | '.join(info_parts)}")
            elif update.update_type == "action":
                content = update.content
                if isinstance(content, dict):
                    # 打印操作描述
                    message = content.get("message", "")
                    if message:
                        print(f"\n操作: {message}")
                    
                    # 保存操作前截图
                    screenshot_before = content.get("screenshot_before")
                    if screenshot_before:
                        filepath = os.path.join(
                            screenshot_dir,
                            f"step_{update.step:03d}_action_before.png"
                        )
                        if save_screenshot(screenshot_before, filepath):
                            print(f"  已保存操作前截图: {filepath}")
            elif update.update_type == "result":
                content = update.content
                if isinstance(content, dict):
                    # 保存操作后截图
                    screenshot_after = content.get("screenshot_after")
                    if screenshot_after:
                        filepath = os.path.join(
                            screenshot_dir,
                            f"step_{update.step:03d}_result_after.png"
                        )
                        if save_screenshot(screenshot_after, filepath):
                            print(f"  已保存操作后截图: {filepath}")
            elif update.update_type == "finished":
                content = update.content
                message = content.get("message") if isinstance(content, dict) else content
                print(f"\n{message}", end="", flush=True)
        
        # 打印总token统计和推理时长统计
        if total_tokens > 0 or total_inference_time > 0:
            print(f"\n\n{'='*60}")
            if total_tokens > 0:
                print(f"📊 总Token统计:")
                print(f"  输入Token总数: {total_prompt_tokens}")
                print(f"  输出Token总数: {total_completion_tokens}")
                print(f"  总Token数: {total_tokens}")
            if total_inference_time > 0:
                print(f"\n⏱️  推理时长统计 (共{step_count}轮):")
                if total_time_to_first_token > 0:
                    avg_ttft = total_time_to_first_token / step_count
                    print(f"  首Token延迟: 总计={total_time_to_first_token:.3f}s, 平均={avg_ttft:.3f}s")
                if total_time_to_thinking_end > 0:
                    avg_tte = total_time_to_thinking_end / step_count
                    print(f"  思考完成延迟: 总计={total_time_to_thinking_end:.3f}s, 平均={avg_tte:.3f}s")
                if total_inference_time > 0:
                    avg_total = total_inference_time / step_count
                    print(f"  总推理时间: 总计={total_inference_time:.3f}s, 平均={avg_total:.3f}s")
            print(f"{'='*60}")
        print("\n" + "=" * 80)
        print("注意：设备会自动释放，无需手动操作")
        print("=" * 80)

    except Exception as e:
        print(f"❌ 示例执行失败: {e}\n{traceback.format_exc()}")

    finally:
        # 关闭客户端
        try:
            ubox.close()
            print("SDK已关闭")
        except:
            pass
