"""ExplorationAgent主类，用于编排手机自动化探索测试操作。"""

import base64
import json
from dataclasses import dataclass
from io import BytesIO
from typing import Any, Callable, Generator, Literal, Optional, TYPE_CHECKING

from PIL import Image

from ubox_py_sdk.logger import get_logger

from .actions import ActionHandler
from .actions.handler import finish, parse_action
from .config import SYSTEM_PROMPT
from .device_interface import IDeviceOperations
from .model import ModelClient, ModelConfig, MessageDict
from .model.client import MessageBuilder

if TYPE_CHECKING:
    from ubox_py_sdk.device import Device

logger = get_logger("ubox_py_sdk.phone_agent")

# 流方式输出的消息类型
UpdateType = Literal["thinking", "model_complete", "action", "result", "finished"]
ActionDict = dict[str, Any]


@dataclass
class AgentConfig:
    """ExplorationAgent的配置类。"""

    """最大执行步数。Agent在执行任务时最多执行多少步，超过此步数将停止执行。"""
    max_steps: int = 100
    system_prompt: Optional[str] = None
    verbose: bool = True
    def __post_init__(self):
        if self.system_prompt is None:
            self.system_prompt = SYSTEM_PROMPT


@dataclass
class StepUpdate:
    """Agent步骤的实时更新信息。"""

    step: int
    """步骤编号。"""

    update_type: UpdateType
    """更新类型：'thinking'（思考过程）、'action'（动作）、'result'（执行结果）、'finished'（完成）。"""

    content: Any
    """更新内容。根据update_type不同而不同：
    - thinking: str - 思考过程的文本片段
    - model_complete: dict - 模型处理完成信息，包含 {
        "tokens": dict - Token统计，包含 {
          "prompt_tokens": int - 输入token数,
          "completion_tokens": int - 输出token数,
          "total_tokens": int - 总token数
        },
        "timing": dict - 推理时长统计，包含 {
          "time_to_first_token": float - 首Token延迟（秒）,
          "time_to_thinking_end": float - 思考完成延迟（秒）,
          "total_time": float - 总推理时间（秒）
        }
      }
    - action: dict - 动作信息，包含 {
        "action": ActionDict - 要执行的动作字典,
        "description": str - 操作描述（简要说明本次操作）,
        "screenshot_before": str - 操作前的截图（base64，已标注操作位置）
      }
    - result: dict - 执行结果，包含 {
        "success": bool - 执行是否成功,
        "message": str - 执行结果消息,
        "screenshot_after": str - 操作后的截图（base64）
      }
    - finished: dict - 完成信息，包含 {"success": bool, "message": str}
    """


@dataclass
class Screenshot:
    """表示一个截图的类。"""

    base64_data: str
    width: int
    height: int
    is_sensitive: bool = False


class ExplorationAgent:
    """
    使用视觉语言模型来理解屏幕内容，并决定执行哪些操作来完成探索测试任务。
    直接使用SDK的Device对象进行操作，支持Android、iOS和HM三端。

    Args:
        model_config: AI模型的配置。
        agent_config: Agent行为的配置。
        confirmation_callback: 敏感操作确认的回调函数（可选）。
        takeover_callback: 接管请求的回调函数（可选）。
        device: SDK的Device实例（必需）。

    Note:
        使用时不需要直接创建ExplorationAgent实例，直接通过device.explore()方法使用。
    """

    def __init__(
        self,
        model_config: Optional[ModelConfig] = None,
        agent_config: Optional[AgentConfig] = None,
        confirmation_callback: Optional[Callable[[str], bool]] = None,
        takeover_callback: Optional[Callable[[str], None]] = None,
        device: Optional[IDeviceOperations] = None,
    ):
        """
        初始化ExplorationAgent。

        Args:
            model_config: AI模型配置
            agent_config: Agent行为配置
            confirmation_callback: 敏感操作确认回调
            takeover_callback: 接管请求回调
            device: 实现IDeviceOperations接口的设备对象（必需）
        """
        self.model_config = model_config or ModelConfig()
        self.agent_config = agent_config or AgentConfig()
        
        # 必须提供 device 对象
        if device is None:
            raise ValueError("device 参数是必需的，请传入实现 IDeviceOperations 接口的设备对象")
        
        self.device = device
        self.model_client = ModelClient(self.model_config)
        self.action_handler = ActionHandler(
            device=device,
            confirmation_callback=confirmation_callback,
            takeover_callback=takeover_callback,
            original_user_prompt_getter=lambda: self._original_user_prompt,
            model_client=self.model_client,  # 使用同一个模型客户端用于提取操作描述
        )

        self._context: list[MessageDict] = []
        self._step_count: int = 0
        self._original_user_prompt: Optional[str] = None  # 保存原始用户输入，用于提取包名
        self._stop_flag: bool = False  # 中断标志，用于主动停止执行
        self._is_running: bool = False  # 标记是否正在执行

    def __call__(self, task: str) -> str:
        """
        使 ExplorationAgent 对象可以直接被调用，内部调用 run() 方法。

        Args:
            task: 任务的自然语言描述。

        Returns:
            Agent的最终消息。
        """
        return self.run(task)

    def run(self, task: str) -> str:
        """
        运行agent完成任务（简单接口，不返回实时更新）。

        如果需要实时监控思考和动作，请使用 run_stream() 方法。

        Args:
            task: 任务的自然语言描述。

        Returns:
            Agent的最终消息。
        """
        # 使用 run_stream() 但只返回最终结果，复用代码
        final_message = None
        for update in self.run_stream(task):
            if update.update_type == "finished":
                final_message = update.content.get("message") if isinstance(update.content, dict) else update.content
                break
        return final_message or "任务完成"

    def stop(self) -> None:
        """
        主动停止agent的执行。
        """
        self._stop_flag = True
        logger.info("Agent收到停止信号，将在下一个检查点停止执行")
    
    @property
    def is_running(self) -> bool:
        """
        检查agent是否正在执行。
        Returns:
            True表示正在执行，False表示未在执行。
        """
        return self._is_running

    def run_stream(self, task: str) -> Generator[StepUpdate, None, None]:
        """
        运行agent完成任务，返回生成器以实时获取执行状态。

        类似 Go 的 channel，可以实时获取 AI 的思考和动作。
        可以通过调用 stop() 方法来主动中断执行。
        如果上一次执行还未完成，会自动停止上一次执行后再开始新的执行。

        Args:
            task: 任务的自然语言描述。

        Yields:
            StepUpdate: 每个步骤的实时更新信息。

        Example:
            >>> agent = ExplorationAgent(...)
            >>> for update in agent.run_stream("打开微信"):
            ...     if update.update_type == "thinking":
            ...         print(f"思考: {update.content}")
            ...     elif update.update_type == "action":
            ...         print(f"动作: {update.content}")
            ...     elif update.update_type == "finished":
            ...         print(f"完成: {update.content.get('message') if isinstance(update.content, dict) else update.content}")
            ...     # 可以在任何时候调用 agent.stop() 来中断执行
        """
        # 如果上一次执行还在进行，先停止它
        if self._is_running:
            logger.warning("检测到上一次执行还未完成，正在停止上一次执行...")
            self._stop_flag = True
            # 等待一小段时间让上一次执行有机会停止
            import time
            time.sleep(0.1)
        
        # 重置状态，开始新的执行
        self._context = []
        self._step_count = 0
        self._stop_flag = False
        self._is_running = True
        
        try:
            for update in self._execute_step_stream(task, is_first=True):
                yield update
                if update.update_type == "finished" or self._stop_flag:
                    if self._stop_flag:
                        yield StepUpdate(
                            step=self._step_count,
                            update_type="finished",
                            content={"success": False, "message": "用户主动中断执行"},
                        )
                    return

            # 继续执行直到完成或达到最大步数
            while self._step_count < self.agent_config.max_steps and not self._stop_flag:
                for update in self._execute_step_stream(is_first=False):
                    yield update
                    if update.update_type == "finished" or self._stop_flag:
                        if self._stop_flag:
                            yield StepUpdate(
                                step=self._step_count,
                                update_type="finished",
                                content={"success": False, "message": "用户主动中断执行"},
                            )
                        return

            # 达到最大步数或被中断
            if self._stop_flag:
                yield StepUpdate(
                    step=self._step_count,
                    update_type="finished",
                    content={"success": False, "message": "用户主动中断执行"},
                )
            else:
                yield StepUpdate(
                    step=self._step_count,
                    update_type="finished",
                    content={"success": False, "message": "达到最大步数"},
                )
        finally:
            # 无论正常完成还是异常，都要重置运行标志
            self._is_running = False

    def step(self, task: Optional[str] = None) -> StepUpdate:
        """
        执行agent的单个步骤。

        用于手动控制或调试。返回包含完整步骤信息的StepUpdate。

        Args:
            task: 任务描述（仅在第一步时需要）。

        Returns:
            包含步骤详情的StepUpdate，包含完整的thinking、action和result信息。
        """
        is_first = len(self._context) == 0

        if is_first and not task:
            raise ValueError("第一步需要提供任务描述")

        # 收集所有更新，返回最后一个包含完整信息的更新
        last_update = None
        for update in self._execute_step_stream(task, is_first):
            last_update = update
            if update.update_type == "finished":
                break
        
        # 如果没有更新，返回一个默认的更新
        if last_update is None:
            return StepUpdate(
                step=self._step_count,
                update_type="finished",
                content={"success": False, "message": "步骤执行失败"},
            )
        
        return last_update

    def reset(self) -> None:
        """重置agent状态以开始新任务。"""
        self._context = []
        self._step_count = 0
        self._stop_flag = False  # 重置停止标志
        self._is_running = False  # 重置运行标志


    def _execute_step_stream(
        self, user_prompt: Optional[str] = None, is_first: bool = False
    ) -> Generator[StepUpdate, None, None]:
        """
        执行agent循环的单个步骤，返回生成器以实时输出状态。

        Args:
            user_prompt: 用户提示（仅在第一步时需要）。
            is_first: 是否为第一步。

        Yields:
            StepUpdate: 每个步骤的实时更新信息。
        """
        self._step_count += 1

        # 获取截图和应用信息
        logger.debug("查看当前屏幕...")
        screenshot = self._get_screenshot()
        current_app = self._get_current_app()

        # 构建消息
        if is_first:
            # 保存原始用户输入，用于后续提取包名
            self._original_user_prompt = user_prompt
            
            self._context.append(
                MessageBuilder.create_system_message(self.agent_config.system_prompt)
            )

            screen_info = MessageBuilder.build_screen_info(current_app)
            text_content = f"{user_prompt}\n\n{screen_info}"

            self._context.append(
                MessageBuilder.create_user_message(
                    text=text_content, image_base64=screenshot.base64_data
                )
            )
        else:
            screen_info = MessageBuilder.build_screen_info(current_app)
            text_content = f"** Screen Info **\n\n{screen_info}"

            self._context.append(
                MessageBuilder.create_user_message(
                    text=text_content, image_base64=screenshot.base64_data
                )
            )

        # 检查是否被中断
        if self._stop_flag:
            yield StepUpdate(
                step=self._step_count,
                update_type="finished",
                content={"success": False, "message": "用户主动中断执行"},
            )
            return

        # 获取模型响应（思考过程通过流式输出）
        try:
            logger.debug("💭 开始思考...")
            # 使用流式请求，实时 yield 思考过程
            thinking_chunks = []
            for chunk in self.model_client.request_stream(self._context):
                # 检查是否被中断
                if self._stop_flag:
                    yield StepUpdate(
                        step=self._step_count,
                        update_type="finished",
                        content={"success": False, "message": "用户主动中断执行"},
                    )
                    return
                thinking_chunks.append(chunk)
                yield StepUpdate(
                    step=self._step_count,
                    update_type="thinking",
                    content=chunk,
                )
            if self._stop_flag:
                yield StepUpdate(
                    step=self._step_count,
                    update_type="finished",
                    content={"success": False, "message": "用户主动中断执行"},
                )
                return
            
            # 获取完整响应（包含token统计和推理时长）
            response = self.model_client.get_last_response()

        except Exception as e:
            if self.agent_config.verbose:
                logger.exception("模型请求失败")
            yield StepUpdate(
                step=self._step_count,
                update_type="finished",
                content={"success": False, "message": f"模型错误: {e}"},
            )
            return

        # 解析动作
        try:
            action = parse_action(response.action)
        except ValueError:
            if self.agent_config.verbose:
                logger.exception("解析动作失败")
            action = finish(message=response.action)

        # 如果是完成动作，直接走完成逻辑，不需要执行操作
        if action.get("_metadata") == "finish":
            # 先输出 model_complete
            if response:
                tokens_info = None
                if hasattr(response, 'total_tokens') and response.total_tokens is not None:
                    tokens_info = {
                        "prompt_tokens": response.prompt_tokens,
                        "completion_tokens": response.completion_tokens,
                        "total_tokens": response.total_tokens,
                    }
                
                timing_info = None
                if hasattr(response, 'total_time') and response.total_time is not None:
                    timing_info = {}
                    if hasattr(response, 'time_to_first_token') and response.time_to_first_token is not None:
                        timing_info["time_to_first_token"] = response.time_to_first_token
                    if hasattr(response, 'time_to_thinking_end') and response.time_to_thinking_end is not None:
                        timing_info["time_to_thinking_end"] = response.time_to_thinking_end
                    if response.total_time is not None:
                        timing_info["total_time"] = response.total_time
                
                content = {}
                if tokens_info:
                    content["tokens"] = tokens_info
                if timing_info:
                    content["timing"] = timing_info
                
                yield StepUpdate(
                    step=self._step_count,
                    update_type="model_complete",
                    content=content,
                )
            
            self._context.append(
                MessageBuilder.create_assistant_message(
                    f"<think>{response.thinking}</think><answer>{response.action}</answer>"
                )
            )
            finish_message = action.get("message", "任务完成")
            if self.agent_config.verbose:
                logger.info("\n" + "🎉 " + "=" * 48)
                logger.info(f"✅ 任务完成: {finish_message}")
                logger.info("=" * 50 + "\n")
            yield StepUpdate(
                step=self._step_count,
                update_type="finished",
                content={"success": True, "message": finish_message},
            )
            return

        if self.agent_config.verbose:
            logger.debug("🎯 执行动作:")
            logger.debug(json.dumps(action, ensure_ascii=False, indent=2))

        # 准备操作前截图并标注（在执行动作之前）
        screenshot_before, action_description, extract_token_info = self.action_handler.prepare_action_screenshot(
            action, screenshot.width, screenshot.height, thinking=response.thinking
        )
        
        # Yield 模型处理完成信息（包含累计后的token统计）
        if response:
            # 累计操作描述提取的token消耗到主响应的token统计中
            if extract_token_info:
                # 累计token消耗
                if response.prompt_tokens is not None:
                    response.prompt_tokens = (response.prompt_tokens or 0) + (extract_token_info.get("prompt_tokens") or 0)
                elif extract_token_info.get("prompt_tokens") is not None:
                    response.prompt_tokens = extract_token_info.get("prompt_tokens")
                
                if response.completion_tokens is not None:
                    response.completion_tokens = (response.completion_tokens or 0) + (extract_token_info.get("completion_tokens") or 0)
                elif extract_token_info.get("completion_tokens") is not None:
                    response.completion_tokens = extract_token_info.get("completion_tokens")
                
                if response.total_tokens is not None:
                    response.total_tokens = (response.total_tokens or 0) + (extract_token_info.get("total_tokens") or 0)
                elif extract_token_info.get("total_tokens") is not None:
                    response.total_tokens = extract_token_info.get("total_tokens")
            
            # 构建token统计信息
            tokens_info = None
            if hasattr(response, 'total_tokens') and response.total_tokens is not None:
                tokens_info = {
                    "prompt_tokens": response.prompt_tokens,
                    "completion_tokens": response.completion_tokens,
                    "total_tokens": response.total_tokens,
                }
            
            # 构建推理时长统计信息
            timing_info = None
            if hasattr(response, 'total_time') and response.total_time is not None:
                timing_info = {}
                if hasattr(response, 'time_to_first_token') and response.time_to_first_token is not None:
                    timing_info["time_to_first_token"] = response.time_to_first_token
                if hasattr(response, 'time_to_thinking_end') and response.time_to_thinking_end is not None:
                    timing_info["time_to_thinking_end"] = response.time_to_thinking_end
                if response.total_time is not None:
                    timing_info["total_time"] = response.total_time
            
            # 构建内容字典（可能包含tokens和timing，也可能为空）
            content = {}
            if tokens_info:
                content["tokens"] = tokens_info
            if timing_info:
                content["timing"] = timing_info
            
            # 无论是否有统计信息，都发送模型完成信号
            yield StepUpdate(
                step=self._step_count,
                update_type="model_complete",
                content=content,
            )

        # Yield 动作信息（包含操作描述和标注后的截图）
        yield StepUpdate(
            step=self._step_count,
            update_type="action",
            content={
                "action": action,
                "message": action_description,
                "screenshot_before": screenshot_before,
            },
        )

        self._context[-1] = MessageBuilder.remove_images_from_message(self._context[-1])

        # 执行动作
        try:
            result = self.action_handler.execute(
                action, screenshot.width, screenshot.height
            )
        except Exception as e:
            if self.agent_config.verbose:
                logger.exception("动作执行失败")
            result = self.action_handler.execute(
                finish(message=str(e)), screenshot.width, screenshot.height
            )

        # Yield 执行结果（包含操作后的截图）
        yield StepUpdate(
            step=self._step_count,
            update_type="result",
            content={
                "success": result.success,
                "message": result.message,
                "screenshot_after": result.screenshot_after,
            },
        )

        self._context.append(
            MessageBuilder.create_assistant_message(
                f"<think>{response.thinking}</think><answer>{response.action}</answer>"
            )
        )

        # 记录执行结果 模型可以根据执行结果决定下一步操作
        result_text = f"action执行结果: {'成功' if result.success else '失败'}"
        if result.message:
            result_text += f" - {result.message}"
        if not result.success:
            result_text += "\n请根据执行结果调整策略，重新尝试或选择其他操作。"
        self._context.append(
            MessageBuilder.create_user_message(text=result_text)
        )

    @property
    def context(self) -> list[MessageDict]:
        """获取当前对话上下文。"""
        return self._context.copy()

    @property
    def step_count(self) -> int:
        """获取当前步数。"""
        return self._step_count

    def _get_screenshot(self) -> Screenshot:
        """从设备获取截图。"""
        try:
            base64_data = self.device.screenshot_base64()
            if not base64_data:
                return self._create_fallback_screenshot()

            if base64_data.startswith("data:image"):
                base64_data = base64_data.split(",", 1)[1]

            img_data = base64.b64decode(base64_data)
            img = Image.open(BytesIO(img_data))
            width, height = img.size

            return Screenshot(
                base64_data=base64_data, width=width, height=height, is_sensitive=False
            )
        except Exception as e:
            if self.agent_config.verbose:
                logger.error(f"截图错误: {e}")
            return self._create_fallback_screenshot()

    def _get_current_app(self) -> str:
        """获取当前应用名称。"""
        try:
            package_name = self.device.current_app()
            return package_name if package_name else "System Home"
        except Exception as e:
            if self.agent_config.verbose:
                logger.error(f"获取当前应用失败: {e}")
            return "System Home"

    def _create_fallback_screenshot(self) -> Screenshot:
        """创建占位截图。"""
        default_width, default_height = 1080, 2400
        black_img = Image.new("RGB", (default_width, default_height), color="black")
        buffered = BytesIO()
        black_img.save(buffered, format="PNG")
        base64_data = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return Screenshot(
            base64_data=base64_data,
            width=default_width,
            height=default_height,
            is_sensitive=False,
        )
