from datetime import datetime
from pathlib import Path
from typing import Optional, Union
import cv2
from maa.controller import Controller
from maa.tasker import TaskDetail
from maa.pipeline import JRecognitionType, JOCR
from maa_mcp.core import mcp, object_registry, _saved_screenshots
from maa_mcp.resource import get_or_create_tasker
from maa_mcp.download import check_ocr_files_exist
from maa_mcp.paths import get_screenshots_dir
@mcp.tool(
name="screencap_and_ocr",
description="""
对当前设备屏幕进行截图，并执行光学字符识别（OCR）处理。
参数：
- controller_id: 控制器 ID，由 connect_adb_device() 或 connect_window() 返回
返回值：
- 成功：返回识别结果列表，包含识别到的文字、坐标信息、置信度等结构化数据
- OCR 资源不存在（首次使用）：返回字符串提示信息，需要调用 check_and_download_ocr() 下载资源后重试
- 失败：返回 None（截图失败或 OCR 识别失败）
说明：
识别结果可用于后续的坐标定位和自动化决策，通常包含文本内容、坐标等信息。
需根据坐标信息理解屏幕上文字的位置和布局，以便进行进一步的交互操作。
""",
)
def screencap_and_ocr(controller_id: str) -> Optional[Union[list, str]]:
# 先检查 OCR 资源是否存在，不存在则返回提示信息让 AI 主动调用下载
if not check_ocr_files_exist():
return "OCR 模型文件不存在，请先调用 check_and_download_ocr() 下载 OCR 资源后重试"
controller: Controller | None = object_registry.get(controller_id)
tasker = get_or_create_tasker(controller_id)
if not controller or not tasker:
return None
image = controller.post_screencap().wait().get()
info: TaskDetail | None = (
tasker.post_recognition(JRecognitionType.OCR, JOCR(), image).wait().get()
)
if not info:
return None
return info.nodes[0].recognition.all_results
@mcp.tool(
name="screencap_only",
description="""
对当前设备屏幕进行截图。
参数：
- controller_id: 控制器 ID，由 connect_adb_device() 或 connect_window() 返回
返回值：
- 成功：返回截图文件的绝对路径，可通过 read_file 工具读取图片内容
- 失败：返回 None
""",
)
def screencap_only(controller_id: str) -> Optional[str]:
controller = object_registry.get(controller_id)
if not controller:
return None
image = controller.post_screencap().wait().get()
if image is None:
return None
# 保存截图到跨平台用户数据目录，返回路径供大模型按需读取
screenshots_dir = get_screenshots_dir()
screenshots_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
filepath = screenshots_dir / f"screenshot_{timestamp}.png"
success = cv2.imwrite(str(filepath), image)
if not success:
return None
# 记录当前会话保存的截图文件路径，用于退出时清理
_saved_screenshots.append(filepath)
return str(filepath.absolute())
