#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author : aidenmo
# @Email : aidenmo@tencent.com
# @Time : 2025/5/23 15:31
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Optional, Union

from loguru import logger
from pydantic import TypeAdapter
from pydantic_ai import Agent, UserPromptNode, ModelRequestNode, CallToolsNode, RunContext, UnexpectedModelBehavior
from pydantic_ai.agent import AgentRunResult
from pydantic_ai.messages import ToolReturnPart, ToolCallPart
from pydantic_ai.usage import Usage

from .config import global_settings, model_settings
from .deps import AgentDeps, SimulateDeviceType, PlanningOutputType, StepOutputType, PlanningStep, ToolParams, StepInfo, \
    MarkFailedParams
from .device import AndroidDevice, WebDevice
from .prompt import SYSTEM_PROMPT, PLANNING_SYSTEM_PROMPT
from .tools import AndroidAgentTool, WebAgentTool, AgentDepsType
from .util.platform import Platform


@dataclass
class PlanningAgent:
    """PlanningAgent class for planning tasks."""
    model: Optional[str] = None
    deps: Optional[AgentDepsType] = None

    async def run(self, prompt: str) -> AgentRunResult[PlanningOutputType]:
        """Run the agent with the given prompt."""
        model = self.model or global_settings.model
        agent = Agent(
            model=model,
            system_prompt=PLANNING_SYSTEM_PROMPT,
            output_type=PlanningOutputType,
            model_settings=model_settings
        )
        return await agent.run(prompt.strip(), deps=self.deps)


@dataclass
class UiAgent:
    model: str
    deps: AgentDepsType
    agent: Agent[AgentDepsType]

    @classmethod
    async def create(cls, *args, **kwargs):
        """Async factory method to create an instance of this class."""
        raise NotImplementedError

    @staticmethod
    async def create_report(report_data: str, report_dir: Union[Path, str]) -> Path:
        """Create a report file based on the given data and directory."""
        logger.info('创建步骤报告...')
        logger.debug(f'report_data: {report_data}')

        report_dir = Path(report_dir)
        report_dir.mkdir(parents=True, exist_ok=True)

        template = Path(__file__).parent / "report_template.html"
        content = template.read_text(encoding='utf-8').replace('{reportData}', report_data)

        output_path = report_dir / f'report_{datetime.now():%Y%m%d%H%M%S}.html'
        output_path.write_text(content, encoding='utf-8')
        logger.info(f"报告：{output_path.resolve().as_uri()}")
        return output_path

    def handle_graph_node(self, node):
        """Format the logger node based on the given node type."""
        if isinstance(node, UserPromptNode):
            logger.log('DETAIL', f"🤖Agent start user task: {repr(node.user_prompt)}")

        elif isinstance(node, ModelRequestNode):
            for part in node.request.parts:
                if isinstance(part, ToolReturnPart):
                    logger.log('DETAIL', f"🤖Agent tool result: {part.tool_name} -> {part.content}")

        elif isinstance(node, CallToolsNode):
            parts = node.model_response.parts
            tool_parts = [part for part in parts if isinstance(part, ToolCallPart)]
            self.deps.context.current_step.parallel_tool_calls = False
            self.deps.context.current_step.parallel_tool_calls = len(tool_parts) > 1
            for part in tool_parts:
                logger.log('DETAIL', f"🤖Agent tool call: {part.tool_name}, args: {part.args}")

    async def _sub_agent_run(self, planning, usage) -> AgentRunResult:
        async with self.agent.iter(user_prompt=planning.instruction, deps=self.deps, usage=usage) as agent_run:
            async for node in agent_run:
                self.handle_graph_node(node)
            return agent_run.result

    async def run(self, prompt: str, system_prompt: Optional[str] = None, report_dir: str = "./report"):
        # TODO: 给用户添加额外的自定义系统提示词，某些场景需要，如：如果出现位置、权限、用户协议等弹窗，点击同意。如果出现登录页面，关闭它。
        logger.info(f"🤖Agent start planning...")

        planning_agent = PlanningAgent(model=self.model, deps=self.deps)
        planning_result = await planning_agent.run(prompt)
        planning_steps = planning_result.output.steps

        planning_steps = [*planning_steps, PlanningStep(instruction='结束任务')]
        logger.info(f"🤖Agent planning finished.")
        for index, step in enumerate(planning_steps, 1):
            logger.info(f'◽️step{index}. {step.instruction}')

        if system_prompt:
            self.agent.system_prompt(lambda: system_prompt)

        usage = planning_result.usage()
        ctx = RunContext(deps=self.deps, model=self.agent.model, usage=Usage(), prompt=None)

        logger.info(f"🤖Agent start executing steps...")
        for step, planning in enumerate(planning_steps, start=1):
            self.deps.context.add_step_info(StepInfo(step=step, planning=planning, description=planning.instruction))
            logger.info('')
            logger.info(f'▶️ step={step} {planning.instruction}')

            if planning.instruction != '结束任务':
                try:
                    result = await self._sub_agent_run(planning, usage)
                    usage = result.usage()
                    logger.info(f"💬 {result.output}")
                except UnexpectedModelBehavior as e:
                    await self.deps.tool.mark_failed(ctx, MarkFailedParams(
                        reason=str(e),
                    ))
                    logger.error(f'step={step} {planning.instruction}: {e}')

                logger.info(f'{"✅" if self.deps.context.current_step.is_success else "❌"} '
                            f'step={step} {planning.instruction}')
            else:
                await self.deps.tool.tear_down(ctx, params=ToolParams(action='tear_down', instruction='任务完成'))

            # 步骤执行后如果没有截图则自动补上，比如滑动、等待
            if not self.deps.context.current_step.image_url:
                await self.deps.tool.get_screen(ctx, parse_element=False)

            if not self.deps.context.current_step.is_success:
                break

        logger.debug(f"steps: {self.deps.context.steps}")
        logger.log('DETAIL', f"usage: {usage}")

        is_success_output = all([step.is_success for step in self.deps.context.steps.values()])

        report_data = {'is_success': is_success_output,
                       'device_size': self.deps.device.device_size,
                       'steps': self.deps.context.steps}
        report_json = TypeAdapter(dict).dump_json(report_data).decode(encoding='utf-8')
        report_path = await self.create_report(report_json, report_dir)

        steps_output = [
            step.model_dump(include={'step', 'description', 'action', 'is_success'})
            for step in self.deps.context.steps.values()
        ]

        return {
            'is_success': is_success_output,
            'steps': steps_output,
            'report_path': report_path
        }


class WebAgent(UiAgent):
    """WebAgent class for web automation."""

    @classmethod
    async def create(
            cls,
            model: Optional[str] = None,
            *,
            device: Optional[WebDevice] = None,
            simulate_device: Optional[SimulateDeviceType] = None,
            headless: Optional[bool] = None,
            tool_cls: Optional[type[WebAgentTool]] = None,
            debug: Optional[bool] = None,
    ):
        settings = global_settings.copy_and_update(
            model=model,
            simulate_device=simulate_device,
            headless=headless,
            debug=debug)

        logger.info(f'settings: {settings}')

        device = device or await WebDevice.create(settings.headless, settings.simulate_device)
        tool = WebAgentTool() if tool_cls is None else tool_cls()
        deps: AgentDeps[WebDevice, WebAgentTool] = AgentDeps(settings, device, tool)

        agent = Agent[AgentDeps](
            model=settings.model,
            system_prompt=SYSTEM_PROMPT,
            model_settings=model_settings,
            deps_type=AgentDeps,
            tools=tool.tools,
            retries=3
        )
        return cls(model, deps, agent)


class MobileAgent(UiAgent):
    """MobileAgent class for mobile device automation."""

    @classmethod
    async def create(
            cls, model: Optional[str] = None,
            *,
            serial: Optional[str] = None,
            platform: Optional[str | Platform] = None,
            tool_cls: Optional[type[AndroidAgentTool]] = None,
            debug: Optional[bool] = None,
    ):
        settings = global_settings.copy_and_update(model=model, debug=debug)

        logger.info(f'settings: {settings}')

        device = await AndroidDevice.create(serial=serial, platform=platform)

        tool = AndroidAgentTool() if tool_cls is None else tool_cls()
        deps: AgentDeps[AndroidDevice, AndroidAgentTool] = AgentDeps(settings, device, tool)

        agent = Agent[AgentDeps, StepOutputType](
            model=settings.model,
            system_prompt=SYSTEM_PROMPT,
            model_settings=model_settings,
            deps_type=AgentDeps,
            tools=tool.tools,
            output_type=StepOutputType,
            retries=2
        )
        return cls(model, deps, agent)
