import asyncio
import hashlib
import shutil
import sys
import tempfile
import uuid
from collections.abc import Callable, Coroutine
from datetime import UTC, datetime
from io import BytesIO
from pathlib import Path
from shutil import which
from types import NoneType
from typing import Any, TypeVar, overload

import httpx
from adbutils import AdbClient
from dotenv import load_dotenv
from langchain_core.messages import AIMessage
from limrun_api import AsyncLimrun
from PIL import Image
from pydantic import BaseModel

from minitap.mobile_use.agents.outputter.outputter import outputter
from minitap.mobile_use.agents.planner.types import Subgoal
from minitap.mobile_use.clients.browserstack_client import BrowserStackClientWrapper
from minitap.mobile_use.clients.idb_client import IdbClientWrapper
from minitap.mobile_use.clients.ios_client import DeviceType, IosClientWrapper, get_ios_client
from minitap.mobile_use.clients.ui_automator_client import UIAutomatorClient
from minitap.mobile_use.clients.wda_client import WdaClientWrapper
from minitap.mobile_use.config import AgentNode, OutputConfig, record_events, settings
from minitap.mobile_use.context import (
    DeviceContext,
    DevicePlatform,
    ExecutionSetup,
    IsReplan,
    MobileUseContext,
)
from minitap.mobile_use.controllers.limrun_controller import (
    LimrunAndroidController,
    LimrunIosController,
)
from minitap.mobile_use.controllers.platform_specific_commands_controller import get_first_device
from minitap.mobile_use.graph.graph import get_graph
from minitap.mobile_use.graph.state import State
from minitap.mobile_use.sdk.builders.agent_config_builder import get_default_agent_config
from minitap.mobile_use.sdk.builders.task_request_builder import TaskRequestBuilder
from minitap.mobile_use.sdk.constants import DEFAULT_PROFILE_NAME
from minitap.mobile_use.sdk.services.cloud_mobile import CloudMobileService
from minitap.mobile_use.sdk.services.platform import PlatformService
from minitap.mobile_use.sdk.types.agent import AgentConfig, LimrunPlatform
from minitap.mobile_use.sdk.types.exceptions import (
    AgentError,
    AgentNotInitializedError,
    AgentProfileNotFoundError,
    AgentTaskRequestError,
    CloudMobileServiceUninitializedError,
    DeviceNotFoundError,
    ExecutableNotFoundError,
    PlatformServiceUninitializedError,
    ServerStartupError,
)
from minitap.mobile_use.sdk.types.platform import TaskRunPlanResponse, TaskRunStatus
from minitap.mobile_use.sdk.types.task import (
    AgentProfile,
    PlatformTaskInfo,
    PlatformTaskRequest,
    Task,
    TaskRequest,
)
from minitap.mobile_use.services.telemetry import telemetry
from minitap.mobile_use.utils.app_launch_utils import _handle_initial_app_launch
from minitap.mobile_use.utils.logger import get_logger
from minitap.mobile_use.utils.media import (
    create_gif_from_trace_folder,
    create_steps_json_from_trace_folder,
    remove_images_from_trace_folder,
    remove_steps_json_from_trace_folder,
)
from minitap.mobile_use.utils.recorder import log_agent_thought

logger = get_logger(__name__)

TOutput = TypeVar("TOutput", bound=BaseModel | None)

load_dotenv()


class Agent:
    _config: AgentConfig
    _tasks: list[Task] = []
    _tmp_traces_dir: Path
    _initialized: bool = False
    _device_context: DeviceContext
    _adb_client: AdbClient | None
    _ui_adb_client: UIAutomatorClient | None
    _ios_client: IosClientWrapper | None
    _ios_device_type: DeviceType | None
    _current_task: asyncio.Task | None = None
    _task_lock: asyncio.Lock
    _cloud_mobile_id: str | None = None
    _limrun_instance_id: str | None = None
    _limrun_controller: Any = None

    def __init__(self, *, config: AgentConfig | None = None):
        self._config = config or get_default_agent_config()
        self._tasks = []
        self._tmp_traces_dir = Path(tempfile.gettempdir()) / "mobile-use-traces"
        self._initialized = False
        self._task_lock = asyncio.Lock()

        # Initialize platform service if API key is available in environment
        # Note: Can also be initialized later with API key at agent .init()
        if settings.MINITAP_API_KEY:
            self._platform_service = PlatformService()
            self._cloud_mobile_service = CloudMobileService()
        else:
            self._platform_service = None
            self._cloud_mobile_service = None

    async def init(
        self,
        api_key: str | None = None,
        server_restart_attempts: int = 3,
        retry_count: int = 5,
        retry_wait_seconds: int = 5,
    ):
        # Start telemetry session for SDK usage (if not already started by CLI)
        if not telemetry._session_id:
            telemetry.start_session({"source": "sdk"})

        try:
            return await self._init_internal(
                api_key=api_key,
                server_restart_attempts=server_restart_attempts,
                retry_count=retry_count,
                retry_wait_seconds=retry_wait_seconds,
            )
        except Exception as e:
            session_id = telemetry._session_id
            telemetry.capture_exception(e, {"phase": "agent_init"})
            telemetry.end_session(success=False, error=str(e))
            if session_id:
                logger.info(f"If you need support, please include this session ID: {session_id}")
            raise

    async def _init_internal(
        self,
        api_key: str | None = None,
        server_restart_attempts: int = 3,
        retry_count: int = 5,
        retry_wait_seconds: int = 5,
    ):
        if api_key:
            self._platform_service = PlatformService(api_key=api_key)
            self._cloud_mobile_service = CloudMobileService(api_key=api_key)

        # Skip initialization for cloud devices - no local setup required
        if self._config.cloud_mobile_id_or_ref:
            if not self._cloud_mobile_service:
                raise CloudMobileServiceUninitializedError()
            self._cloud_mobile_id = await self._cloud_mobile_service.resolve_cloud_mobile_id(
                cloud_mobile_id_or_ref=self._config.cloud_mobile_id_or_ref,
            )
            logger.info("Cloud device configured - skipping local initialization")
            self._initialized = True
            return True

        # Handle BrowserStack initialization
        if self._config.browserstack_config:
            logger.info("Initializing BrowserStack session...")
            self._ios_client = BrowserStackClientWrapper(config=self._config.browserstack_config)
            session_started = await self._ios_client.init_client()
            if not session_started:
                raise ServerStartupError(
                    message="Failed to create BrowserStack session. "
                    "Please check your credentials and device configuration."
                )
            self._ios_device_type = DeviceType.BROWSERSTACK
            self._adb_client = None
            self._ui_adb_client = None
            logger.success("BrowserStack session created successfully")

            self._device_context = await self._get_device_context(
                device_id="browserstack", platform=DevicePlatform.IOS
            )
            logger.info(self._device_context.to_str())
            logger.info("✅ Mobile-use agent initialized with BrowserStack.")
            self._initialized = True
            telemetry.capture_agent_initialized(
                platform=DevicePlatform.IOS.value,
                device_id="browserstack",
            )
            return True

        # Handle Limrun cloud device initialization
        # Check for pre-configured controllers first, then fall back to limrun_config
        if self._config.limrun_android_controller or self._config.limrun_ios_controller:
            return await self._init_limrun_device(
                android_controller=self._config.limrun_android_controller,
                ios_controller=self._config.limrun_ios_controller,
            )
        if self._config.limrun_config:
            return await self._init_limrun_device()

        if not which("adb") and not which("xcrun"):
            raise ExecutableNotFoundError("cli_tools")

        if self._initialized:
            logger.warning("Agent is already initialized. Skipping...")
            return True

        # Get first available device ID
        if not self._config.device_id or not self._config.device_platform:
            device_id, platform, ios_device_type = get_first_device(logger=logger)
        else:
            device_id, platform = self._config.device_id, self._config.device_platform
            ios_device_type = None  # Will be auto-detected in _init_clients

        if not device_id or not platform:
            error_msg = "No device found. Exiting."
            logger.error(error_msg)
            raise DeviceNotFoundError(error_msg)

        # Initialize clients
        self._init_clients(
            device_id=device_id,
            platform=platform,
            ios_device_type=ios_device_type,
            retry_count=retry_count,
            retry_wait_seconds=retry_wait_seconds,
        )

        # Initialize iOS client (IDB companion for simulators, WDA already running for physical)
        if self._ios_client:
            if isinstance(self._ios_client, IdbClientWrapper):
                logger.info("Starting IDB companion for iOS simulator...")
                companion_started = await self._ios_client.init_companion()
                if not companion_started:
                    raise ServerStartupError(
                        message="Failed to start IDB companion for iOS simulator. "
                        "Please ensure fb-idb is installed: https://fbidb.io/docs/installation/"
                    )
                logger.success("IDB companion started successfully")
            elif isinstance(self._ios_client, WdaClientWrapper):
                logger.info("Connecting to WebDriverAgent for physical iOS device...")
                wda_connected = await self._ios_client.init_client()
                if not wda_connected:
                    raise ServerStartupError(
                        message="Failed to connect to WebDriverAgent. "
                        "Please ensure WDA is running on your device. "
                        "See the setup instructions above."
                    )
                logger.success("WDA client connected for physical device")

        # Start necessary servers
        restart_attempt = 0
        while restart_attempt < server_restart_attempts:
            success = self._run_servers(
                device_id=device_id,
                platform=platform,
            )
            if success:
                break

            restart_attempt += 1
            if restart_attempt < server_restart_attempts:
                logger.warning(
                    f"Server start failed, attempting restart "
                    f"{restart_attempt}/{server_restart_attempts}"
                )
            else:
                error_msg = "Mobile-use servers failed to start after all restart attempts."
                logger.error(error_msg)
                raise ServerStartupError(message=error_msg)

        self._device_context = await self._get_device_context(
            device_id=device_id, platform=platform
        )
        logger.info(self._device_context.to_str())
        logger.info("✅ Mobile-use agent initialized.")
        self._initialized = True
        telemetry.capture_agent_initialized(
            platform=platform.value,
            device_id=device_id,
        )
        return True

    async def install_apk(self, apk_path: str | Path) -> None:
        """
        Install an APK on the connected device.
        For cloud mobiles, the APK must be x86_64 compatible.

        Args:
            apk_path: Path to the local APK file to install

        Raises:
            AgentNotInitializedError: If the agent is not initialized
            AgentError: If attempting to install on non-Android device or ADB operations fail
            FileNotFoundError: If the APK file doesn't exist
            CloudMobileServiceUninitializedError: If cloud service is unavailable
        """
        try:
            await self._install_apk_internal(apk_path)
        except Exception as e:
            telemetry.capture_exception(e, {"phase": "install_apk"})
            raise

    async def _install_apk_internal(self, apk_path: str | Path) -> None:
        if isinstance(apk_path, str):
            apk_path = Path(apk_path)

        if not apk_path.exists():
            raise FileNotFoundError(f"APK file not found: {apk_path}")

        if self._config.cloud_mobile_id_or_ref:
            await self._install_apk_on_cloud_mobile(apk_path)
        else:
            if not self._initialized:
                raise AgentNotInitializedError()

            if self._device_context.mobile_platform != DevicePlatform.ANDROID:
                raise AgentError(
                    "APK can only be installed on Android devices but got "
                    f"'{self._device_context.mobile_platform.value}'"
                )

            device_id = self._device_context.device_id
            logger.info(f"Installing APK on Android device '{device_id}'")
            if not self._adb_client:
                raise AgentError("ADB client not initialized")

            device = self._adb_client.device(serial=device_id)
            await asyncio.to_thread(device.install, apk_path)
            logger.info(f"APK installed successfully on Android device '{device_id}'")

    async def _install_apk_on_cloud_mobile(self, apk_path: Path) -> None:
        """
        Install an APK on a cloud mobile device.

        This method starts the cloud mobile if needed, then uploads and installs the APK.
        """
        if not self._cloud_mobile_id:
            raise AgentTaskRequestError("Cloud mobile ID is not configured")

        if not self._cloud_mobile_service:
            raise CloudMobileServiceUninitializedError()

        # Check platform before starting - fail early if not Android
        vm_info = await self._cloud_mobile_service._get_virtual_mobile_status(self._cloud_mobile_id)
        if vm_info.platform and vm_info.platform != "android":
            raise AgentError(
                f"APK can only be installed on Android cloud mobiles but got '{vm_info.platform}'"
            )

        # Start cloud mobile if not already started
        logger.info(f"Starting cloud mobile '{self._cloud_mobile_id}' for APK installation...")
        await self._cloud_mobile_service.start_and_wait_for_ready(
            cloud_mobile_id=self._cloud_mobile_id,
        )

        # Install APK
        logger.info(f"Installing APK '{apk_path.name}' on cloud mobile '{self._cloud_mobile_id}'")
        await self._cloud_mobile_service.install_apk(
            cloud_mobile_id=self._cloud_mobile_id,
            apk_path=apk_path,
        )
        logger.success(f"APK '{apk_path.name}' installed successfully")

    async def install_app(self, app_path: str | Path) -> str | None:
        """
        Install an app on the connected device.

        For Android: Installs an APK file using ADB.
        For iOS (Limrun): Uploads and installs a .app folder using diff-based
                         patch syncing for fast updates.

        Args:
            app_path: Path to the app to install:
                      - Android: Path to an APK file
                      - iOS: Path to a .app folder (simulator build)

        Returns:
            The bundle ID of the installed app (iOS only), or None for Android.

        Raises:
            AgentNotInitializedError: If the agent is not initialized
            AgentError: If installation fails or platform is unsupported
            FileNotFoundError: If the app file/folder doesn't exist
        """
        try:
            return await self._install_app_internal(app_path)
        except Exception as e:
            telemetry.capture_exception(e, {"phase": "install_app"})
            raise

    async def _install_app_internal(self, app_path: str | Path) -> str | None:
        if isinstance(app_path, str):
            app_path = Path(app_path)

        if not app_path.exists():
            raise FileNotFoundError(f"App not found: {app_path}")

        if not self._initialized:
            raise AgentNotInitializedError()

        platform = self._device_context.mobile_platform

        if platform == DevicePlatform.ANDROID:
            await self._install_apk_internal(app_path)
            return None

        elif platform == DevicePlatform.IOS:
            return await self._install_ios_app(app_path)

        else:
            raise AgentError(f"App installation not supported for platform: {platform}")

    async def _install_ios_app(self, app_path: Path) -> str:
        """
        Install an iOS .app bundle on a Limrun iOS device.

        Uses diff-based patch syncing for fast updates - only changed parts
        of the app are uploaded.

        Args:
            app_path: Path to the .app folder (simulator build)

        Returns:
            The bundle ID of the installed app

        Raises:
            AgentError: If not connected to a Limrun iOS device
            FileNotFoundError: If the .app folder doesn't exist
        """

        if not app_path.is_dir() or not app_path.suffix == ".app":
            raise AgentError(
                f"Expected a .app folder for iOS, got: {app_path}. "
                "Please provide the path to your built .app folder "
                "(e.g., build/Debug-iphonesimulator/MyApp.app)"
            )

        # Check if we have a Limrun iOS controller
        if not isinstance(self._ios_client, LimrunIosController):
            raise AgentError(
                "iOS app installation is only supported for Limrun iOS devices. "
                "For local simulators, use 'xcrun simctl install' directly."
            )

        limrun_controller: LimrunIosController = self._ios_client

        # Get API key from settings or platform service
        api_key_value: str | None = None
        if settings.MINITAP_API_KEY:
            api_key_raw = settings.MINITAP_API_KEY
            api_key_value = (
                api_key_raw.get_secret_value()
                if hasattr(api_key_raw, "get_secret_value")
                else str(api_key_raw)
            )
        if not api_key_value and self._platform_service:
            api_key_value = self._platform_service._api_key
        if not api_key_value:
            raise AgentError(
                "API key is required for iOS app installation. "
                "Set MINITAP_API_KEY environment variable."
            )

        base_url = settings.MINITAP_BASE_URL or "https://platform.minitap.ai"

        app_name = app_path.stem
        asset_name = f"{app_name}.zip"

        logger.info(f"Preparing iOS app for upload: {app_name}")

        # Create a temporary zip file
        with tempfile.TemporaryDirectory() as temp_dir:
            zip_path = Path(temp_dir) / asset_name
            zip_base = str(zip_path.with_suffix(""))
            shutil.make_archive(zip_base, "zip", app_path.parent, app_path.name)

            # Calculate MD5 hash
            md5 = hashlib.md5()
            with open(zip_path, "rb") as f:
                for chunk in iter(lambda: f.read(8192), b""):
                    md5.update(chunk)
            md5_hash = md5.hexdigest()

            logger.info(f"Created zip archive: {zip_path.name}, MD5: {md5_hash}")

            # Upload to Limrun using assets API
            client = AsyncLimrun(api_key=api_key_value, base_url=f"{base_url}/api/v1/limrun")
            try:
                logger.info("Getting upload URL from Limrun assets API...")
                asset_response = await client.assets.get_or_create(name=asset_name)

                # Check if we need to upload (MD5 mismatch or no existing file)
                existing_md5 = asset_response.md5
                if existing_md5 == md5_hash:
                    logger.info("App already uploaded with matching MD5, skipping upload")
                else:
                    upload_url = asset_response.signed_upload_url
                    if not upload_url:
                        raise AgentError("No upload URL returned from Limrun assets API")

                    logger.info("Uploading iOS app to Limrun storage...")
                    async with httpx.AsyncClient(timeout=300.0) as http_client:
                        file_content = zip_path.read_bytes()
                        response = await http_client.put(
                            upload_url,
                            content=file_content,
                            headers={"Content-Type": "application/zip"},
                        )
                        response.raise_for_status()
                    logger.info("iOS app uploaded successfully")

                download_url = asset_response.signed_download_url
                if not download_url:
                    raise AgentError("No download URL returned from Limrun assets API")

            finally:
                await client.close()

        # Install app via WebSocket
        logger.info("Installing app on Limrun iOS device...")
        install_result = await limrun_controller.client.install_app(url=download_url, md5=md5_hash)
        bundle_id = install_result.get("bundleId", "")

        if bundle_id:
            logger.success(f"iOS app installed successfully: {bundle_id}")
        else:
            logger.success("iOS app installed successfully")

        return bundle_id

    def new_task(self, goal: str):
        """
        Create a new task request builder.

        Args:
            goal: Natural language description of what to accomplish

        Returns:
            TaskRequestBuilder that can be configured with:
            - .with_output_format() for structured output
            - .with_output_description() for output description
            - .with_locked_app_package() to restrict execution to a specific app
            - .using_profile() to specify an LLM profile
            - .with_max_steps() to set maximum execution steps
            - .with_trace_recording() to enable trace recording
            - .with_name() to set a custom task name
        """
        return TaskRequestBuilder[None].from_common(
            goal=goal,
            common=self._config.task_request_defaults,
        )

    @overload
    async def run_task(
        self,
        *,
        goal: str,
        output: type[TOutput],
        profile: str | AgentProfile | None = None,
        name: str | None = None,
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> TOutput | None: ...

    @overload
    async def run_task(
        self,
        *,
        goal: str,
        output: str,
        profile: str | AgentProfile | None = None,
        name: str | None = None,
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> str | dict | None: ...

    @overload
    async def run_task(
        self,
        *,
        goal: str,
        output=None,
        profile: str | AgentProfile | None = None,
        name: str | None = None,
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> str | None: ...

    @overload
    async def run_task(
        self,
        *,
        request: TaskRequest[None],
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> str | dict | None: ...

    @overload
    async def run_task(
        self,
        *,
        request: TaskRequest[TOutput],
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> TOutput | None: ...

    @overload
    async def run_task(
        self,
        *,
        request: PlatformTaskRequest[None],
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> str | dict | None: ...

    @overload
    async def run_task(
        self,
        *,
        request: PlatformTaskRequest[TOutput],
        locked_app_package: str | None = None,
        app_path: str | Path | None = None,
    ) -> TOutput | None: ...

    async def run_task(
        self,
        *,
        goal: str | None = None,
        output: type[TOutput] | str | None = None,
        profile: str | AgentProfile | None = None,
        locked_app_package: str | None = None,
        name: str | None = None,
        app_path: str | Path | None = None,
        request: TaskRequest[TOutput] | PlatformTaskRequest[TOutput] | None = None,
    ) -> str | dict | TOutput | None:
        # Check if cloud mobile is configured
        if self._config.cloud_mobile_id_or_ref:
            if request is None or not isinstance(request, PlatformTaskRequest):
                raise AgentTaskRequestError(
                    "When using a cloud mobile, only PlatformTaskRequest is supported. "
                    "Use AgentConfigBuilder.for_cloud_mobile() only with PlatformTaskRequest."
                )
            # Use cloud mobile execution path
            return await self._run_cloud_mobile_task(
                request=request, locked_app_package=locked_app_package
            )

        # Normal local execution path
        if request is not None:
            task_info = None
            if isinstance(request, PlatformTaskRequest):
                if not self._platform_service:
                    raise PlatformServiceUninitializedError()
                task_info = await self._platform_service.create_task_run(
                    request=request,
                    locked_app_package=locked_app_package,
                    enable_video_tools=self._config.video_recording_enabled,
                )
                request.task_run_id = task_info.task_run.id
                request.task_run_id_available_event.set()
                self._config.agent_profiles[task_info.llm_profile.name] = task_info.llm_profile
                request = task_info.task_request
            elif locked_app_package is not None:
                if request.locked_app_package:
                    logger.warning(
                        "Locked app package specified both in the request and as a parameter. "
                        "Using the parameter value."
                    )
                request.locked_app_package = locked_app_package
            # Handle app_path parameter override
            if app_path is not None:
                if request.app_path:
                    logger.warning(
                        "App path specified both in the request and as a parameter. "
                        "Using the parameter value."
                    )
                request.app_path = Path(app_path) if isinstance(app_path, str) else app_path
            return await self._run_task(
                request=request, task_info=task_info, platform_service=self._platform_service
            )
        if goal is None:
            raise AgentTaskRequestError("Goal is required")
        task_request = self.new_task(goal=goal)
        if output is not None:
            if isinstance(output, str):
                task_request.with_output_description(description=output)
            elif output is not NoneType:
                task_request.with_output_format(output_format=output)
        if profile is not None:
            task_request.using_profile(profile=profile)
        if name is not None:
            task_request.with_name(name=name)
        if locked_app_package is not None:
            task_request.with_locked_app_package(package_name=locked_app_package)
        if app_path is not None:
            task_request.with_app_path(app_path=app_path)
        return await self._run_task(task_request.build())

    async def _run_cloud_mobile_task(
        self,
        request: PlatformTaskRequest[TOutput],
        locked_app_package: str | None = None,
    ) -> str | dict | TOutput | None:
        """
        Execute a task on a cloud mobile.

        This method triggers the task execution on the Platform and polls
        for completion without running any agentic logic locally.
        """
        if not self._cloud_mobile_id:
            raise AgentTaskRequestError("Cloud mobile ID is not configured")

        if not self._cloud_mobile_service:
            raise CloudMobileServiceUninitializedError()

        if not self._platform_service:
            raise PlatformServiceUninitializedError()

        if self._config.video_recording_enabled:
            profile_name = request.profile or DEFAULT_PROFILE_NAME
            _, profile = await self._platform_service.get_profile(profile_name)
            if not profile.llm_config.utils.video_analyzer:
                raise AgentTaskRequestError(
                    f"video_recording_enabled: profile '{profile_name}' "
                    "must have a video_analyzer agent configured"
                )

        # Start cloud mobile if not already started
        logger.info(f"Starting cloud mobile '{self._cloud_mobile_id}'...")
        await self._cloud_mobile_service.start_and_wait_for_ready(
            cloud_mobile_id=self._cloud_mobile_id,
        )
        logger.info(
            f"Starting cloud mobile task execution '{self._cloud_mobile_id}'",
        )

        def log_callback(message: str):
            """Callback for logging timeline updates."""
            logger.info(message)

        def status_callback(
            status: TaskRunStatus,
            status_message: str | None,
        ):
            """Callback for status updates."""
            logger.info(f"Task status update: [{status}] {status_message}")

        async def _execute_cloud(cloud_mobile_service: CloudMobileService, cloud_mobile_id: str):
            try:
                # Execute task on cloud mobile and wait for completion
                final_status, error, output = await cloud_mobile_service.run_task_on_cloud_mobile(
                    cloud_mobile_id=cloud_mobile_id,
                    request=request,
                    on_status_update=status_callback,
                    on_log=log_callback,
                    locked_app_package=locked_app_package,
                    enable_video_tools=self._config.video_recording_enabled,
                )
                if final_status == "completed":
                    logger.success("Cloud mobile task completed successfully")
                    return output
                if final_status == "failed":
                    logger.error(f"Cloud mobile task failed: {error}")
                    raise AgentTaskRequestError(
                        f"Task execution failed on cloud mobile: {error}",
                    )
                if final_status == "cancelled":
                    logger.warning("Cloud mobile task was cancelled")
                    raise AgentTaskRequestError("Task execution was cancelled")
                logger.error(f"Unknown cloud mobile task status: {final_status}")
                raise AgentTaskRequestError(f"Unknown task status: {final_status}")
            except asyncio.CancelledError:
                # Propagate cancellation to parent coroutine.
                logger.info("Task cancelled during execution, re-raising CancelledError")
                raise
            except AgentTaskRequestError as e:
                # Capture and re-raise known exceptions
                telemetry.capture_exception(e, {"phase": "cloud_mobile_task"})
                raise
            except Exception as e:
                logger.error(f"Unexpected error during cloud mobile task execution: {e}")
                telemetry.capture_exception(e, {"phase": "cloud_mobile_task"})
                raise AgentTaskRequestError(f"Unexpected error: {e}") from e

        async with self._task_lock:
            if self._current_task and not self._current_task.done():
                logger.warning(
                    "Another cloud task is running; cancelling it before starting new one.",
                )
                self.stop_current_task()
                try:
                    await self._current_task
                except asyncio.CancelledError:
                    pass
            try:
                self._current_task = asyncio.create_task(
                    _execute_cloud(
                        cloud_mobile_service=self._cloud_mobile_service,
                        cloud_mobile_id=self._cloud_mobile_id,
                    ),
                )
                return await self._current_task
            finally:
                self._current_task = None

    async def _run_task(
        self,
        request: TaskRequest[TOutput],
        task_info: PlatformTaskInfo | None = None,
        platform_service: PlatformService | None = None,
    ) -> str | dict | TOutput | None:
        if not self._initialized:
            raise AgentNotInitializedError()

        if request.profile:
            agent_profile = self._config.agent_profiles.get(request.profile)
            if agent_profile is None:
                raise AgentProfileNotFoundError(request.profile)
        else:
            agent_profile = self._config.default_profile

        if (
            self._config.video_recording_enabled
            and agent_profile.llm_config.utils.video_analyzer is None
        ):
            raise ValueError(
                f"with_video_recording_tools() requires 'video_analyzer' in utils for "
                f"profile '{agent_profile.name}'. Add 'video_analyzer' with a "
                f"video-capable model (e.g., gemini-3-flash-preview)."
            )

        logger.info(str(agent_profile))

        on_status_changed = None
        on_agent_thought = None
        on_plan_changes = None
        task_id = str(uuid.uuid4())
        if task_info:
            on_status_changed = self._get_task_status_change_callback(
                task_info=task_info, platform_service=platform_service
            )
            on_agent_thought = self._get_new_agent_thought_callback(
                task_info=task_info, platform_service=platform_service
            )
            on_plan_changes = self._get_plan_changes_callback(
                task_info=task_info, platform_service=platform_service
            )
            task_id = task_info.task_run.id

        task = Task(
            id=task_id,
            device=self._device_context,
            status="pending",
            request=request,
            created_at=datetime.now(),
            on_status_changed=on_status_changed,
        )
        self._tasks.append(task)
        task_name = task.get_name()

        # Extract API key from platform service if available
        api_key = None
        if platform_service:
            api_key = platform_service._api_key

        # Get limrun controller if available
        limrun_android_ctrl = None
        if hasattr(self, "_limrun_controller") and self._limrun_controller is not None:
            from minitap.mobile_use.controllers.limrun_controller import LimrunAndroidController

            if isinstance(self._limrun_controller, LimrunAndroidController):
                limrun_android_ctrl = self._limrun_controller

        context = MobileUseContext(
            trace_id=task.id,
            device=self._device_context,
            adb_client=self._adb_client,
            ui_adb_client=self._ui_adb_client,
            ios_client=self._ios_client,
            limrun_android_controller=limrun_android_ctrl,
            llm_config=agent_profile.llm_config,
            on_agent_thought=on_agent_thought,
            on_plan_changes=on_plan_changes,
            minitap_api_key=api_key,
            video_recording_enabled=(
                self._config.video_recording_enabled
                and agent_profile.llm_config.utils.video_analyzer is not None
            ),
        )

        self._prepare_tracing(task=task, context=context)
        await self._prepare_app_installation(task=task)
        await self._prepare_app_lock(task=task, context=context)
        self._prepare_output_files(task=task)

        output_config = None
        if request.output_description or request.output_format:
            output_config = OutputConfig(
                output_description=request.output_description,
                structured_output=request.output_format,  # type: ignore
            )
            logger.info(str(output_config))

        logger.info(f"[{task_name}] Starting graph with goal: `{request.goal}`")
        state = self._get_graph_state(task=task)
        graph_input = state.model_dump()
        task_start_time = datetime.now(UTC)

        telemetry.capture_task_started(
            task_id=task_id,
            platform=self._device_context.mobile_platform.value,
            has_locked_app=request.locked_app_package is not None,
        )

        async def _execute_task_logic():
            last_state: State | None = None
            last_state_snapshot: dict | None = None
            output = None
            try:
                logger.info(f"[{task_name}] Invoking graph with input: {graph_input}")
                await task.set_status(status="running", message="Invoking graph...")
                async for chunk in (await get_graph(context)).astream(
                    input=graph_input,
                    config={
                        "recursion_limit": task.request.max_steps,
                        "callbacks": self._config.graph_config_callbacks,
                    },
                    stream_mode=["messages", "custom", "updates", "values"],
                ):
                    stream_mode, payload = chunk
                    if stream_mode == "values":
                        last_state_snapshot = payload  # type: ignore
                        last_state = State(**last_state_snapshot)  # type: ignore
                        if task.request.thoughts_output_path:
                            record_events(
                                output_path=task.request.thoughts_output_path,
                                events=last_state.agents_thoughts,
                            )

                    if stream_mode == "updates":
                        for _, value in payload.items():  # type: ignore node name, node output
                            if value and "agents_thoughts" in value:
                                new_thoughts = value["agents_thoughts"]
                                last_item = new_thoughts[-1] if new_thoughts else None
                                if last_item:
                                    log_agent_thought(
                                        agent_thought=last_item,
                                    )

                if not last_state:
                    err = f"[{task_name}] No result received from graph"
                    logger.warning(err)
                    await task.finalize(content=output, state=last_state_snapshot, error=err)
                    return None

                print_ai_response_to_stderr(graph_result=last_state)
                output = await self._extract_output(
                    task_name=task_name,
                    ctx=context,
                    request=request,
                    output_config=output_config,
                    state=last_state,
                )
                logger.info(f"✅ Automation '{task_name}' is success ✅")
                await task.finalize(content=output, state=last_state_snapshot)
                duration = (datetime.now(UTC) - task_start_time).total_seconds()
                steps_count = len(last_state.agents_thoughts) if last_state else 0
                telemetry.capture_task_completed(
                    task_id=task_id,
                    success=True,
                    steps_count=steps_count,
                    duration_seconds=duration,
                )
                return output
            except asyncio.CancelledError:
                err = f"[{task_name}] Task cancelled"
                logger.warning(err)
                await task.finalize(
                    content=output,
                    state=last_state_snapshot,
                    error=err,
                    cancelled=True,
                )
                duration = (datetime.now(UTC) - task_start_time).total_seconds()
                steps_count = len(last_state.agents_thoughts) if last_state else 0
                telemetry.capture_task_completed(
                    task_id=task_id,
                    success=False,
                    steps_count=steps_count,
                    duration_seconds=duration,
                    cancelled=True,
                )
                raise
            except Exception as e:
                err = f"[{task_name}] Error running automation: {e}"
                logger.error(err)
                await task.finalize(
                    content=output,
                    state=last_state_snapshot,
                    error=err,
                )
                duration = (datetime.now(UTC) - task_start_time).total_seconds()
                steps_count = len(last_state.agents_thoughts) if last_state else 0
                telemetry.capture_task_completed(
                    task_id=task_id,
                    success=False,
                    steps_count=steps_count,
                    duration_seconds=duration,
                )
                telemetry.capture_exception(e, {"task_id": task_id})
                if telemetry._session_id:
                    logger.info(
                        "If you need support, please include this session ID: "
                        f"{telemetry._session_id}"
                    )
                raise
            finally:
                await self._finalize_tracing(task=task, context=context)

        async with self._task_lock:
            if self._current_task and not self._current_task.done():
                logger.warning(
                    "Another automation task is already running. "
                    "Stopping it before starting the new one."
                )
                self.stop_current_task()
                try:
                    await self._current_task
                except asyncio.CancelledError:
                    pass

            try:
                self._current_task = asyncio.create_task(_execute_task_logic())
                return await self._current_task
            finally:
                self._current_task = None

    def stop_current_task(self):
        """Requests cancellation of the currently running automation task."""
        if self._current_task and not self._current_task.done():
            logger.info("Requesting to stop the current automation task...")
            was_cancelled = self._current_task.cancel()
            if was_cancelled:
                logger.success("Cancellation request for the current task was sent.")
            else:
                logger.warning(
                    "Could not send cancellation request for the current task "
                    "(it may already be completing)."
                )
        else:
            logger.info("No active automation task to stop.")

    async def get_screenshot(self) -> Image.Image:
        """
        Capture a screenshot from the mobile device.

        For cloud mobiles, this method calls the mobile-manager endpoint.
        For local mobiles, it uses ADB (Android) or xcrun (iOS) directly.

        Returns:
            Screenshot as PIL Image

        Raises:
            AgentNotInitializedError: If the agent is not initialized
            PlatformServiceUninitializedError: If cloud mobile service is not available
            Exception: If screenshot capture fails
        """
        # Check if cloud mobile is configured
        if self._cloud_mobile_id:
            if not self._cloud_mobile_service:
                raise CloudMobileServiceUninitializedError()
            screenshot = await self._cloud_mobile_service.get_screenshot(
                cloud_mobile_id=self._cloud_mobile_id,
            )
            return screenshot

        # Local device - use ADB or xcrun directly
        if not self._initialized:
            raise AgentNotInitializedError()

        if self._device_context.mobile_platform == DevicePlatform.ANDROID:
            # Use ADB to capture screenshot
            logger.info("Capturing screenshot from local Android device")
            if not self._adb_client:
                raise Exception("ADB client not initialized")

            device = self._adb_client.device(serial=self._device_context.device_id)
            screenshot = await asyncio.to_thread(device.screenshot)
            logger.info("Screenshot captured from local Android device")
            return screenshot

        elif self._device_context.mobile_platform == DevicePlatform.IOS:
            from io import BytesIO

            from minitap.mobile_use.controllers.limrun_controller import LimrunIosController

            # Check if using Limrun iOS controller
            if isinstance(self._ios_client, LimrunIosController):
                logger.info("Capturing screenshot from Limrun iOS device")
                screenshot_bytes = await self._ios_client.screenshot()
                if screenshot_bytes is None:
                    raise Exception("Failed to capture screenshot from Limrun iOS device")
                screenshot = Image.open(BytesIO(screenshot_bytes))
                logger.info("Screenshot captured from Limrun iOS device")
                return screenshot

            # Use xcrun to capture screenshot for local simulators
            import functools
            import subprocess

            logger.info("Capturing screenshot from local iOS device")
            try:
                # xcrun simctl io <device> screenshot --type=png -
                result = await asyncio.to_thread(
                    functools.partial(
                        subprocess.run,
                        [
                            "xcrun",
                            "simctl",
                            "io",
                            self._device_context.device_id,
                            "screenshot",
                            "--type=png",
                            "-",
                        ],
                        capture_output=True,
                        check=True,
                    )
                )
                # Convert bytes to PIL Image
                screenshot = Image.open(BytesIO(result.stdout))
                logger.info("Screenshot captured from local iOS device")
                return screenshot
            except subprocess.CalledProcessError as e:
                logger.error(f"Failed to capture screenshot: {e}")
                raise Exception(f"Failed to capture screenshot from iOS device: {e}")

        else:
            raise Exception(f"Unsupported platform: {self._device_context.mobile_platform}")

    async def clean(self, force: bool = False):
        if self._cloud_mobile_id:
            self._initialized = False
            logger.info("✅ Cloud-mode agent stopped.")
            # End telemetry session if started by SDK (not CLI)
            if telemetry._session_id and telemetry._session_context.get("source") == "sdk":
                telemetry.end_session(success=True)
            return

        # Cleanup Limrun device if provisioned by SDK
        if self._limrun_instance_id:
            await self._cleanup_limrun_device()
            self._initialized = False
            logger.info("✅ Limrun agent stopped.")
            if telemetry._session_id and telemetry._session_context.get("source") == "sdk":
                telemetry.end_session(success=True)
            return

        if not self._initialized and not force:
            return

        if self._ios_client:
            await self._ios_client.cleanup()
            self._ios_client = None

        self._initialized = False
        logger.info("✅ Mobile-use agent stopped.")

        # End telemetry session if started by SDK (not CLI)
        if telemetry._session_id and telemetry._session_context.get("source") == "sdk":
            telemetry.end_session(success=True)

    async def _init_limrun_device(
        self,
        android_controller: LimrunAndroidController | None = None,
        ios_controller: LimrunIosController | None = None,
    ) -> bool:
        """
        Initialize a Limrun cloud device.

        This method either uses pre-configured controllers or provisions a new
        Limrun device based on limrun_config, connects to it, and sets up the
        appropriate clients for device interaction.

        Args:
            android_controller: Pre-configured Limrun Android controller (optional).
            ios_controller: Pre-configured Limrun iOS controller (optional).
        """

        # Use pre-configured Android controller
        if android_controller is not None:
            logger.info("Using pre-configured Limrun Android controller")
            self._limrun_controller = android_controller
            self._limrun_instance_id = android_controller.instance_id

            self._adb_client = android_controller._adb_client
            self._ui_adb_client = android_controller._ui_client
            self._ios_client = None
            self._ios_device_type = None

            self._device_context = DeviceContext(
                host_platform="LINUX",
                mobile_platform=DevicePlatform.ANDROID,
                device_id=android_controller._adb_serial or self._limrun_instance_id,
                device_width=android_controller.device_width,
                device_height=android_controller.device_height,
            )

            logger.info(f"Limrun Android device ready: {self._limrun_instance_id}")
            logger.info(self._device_context.to_str())
            logger.info("✅ Mobile-use agent initialized with pre-configured Limrun Android.")
            self._initialized = True
            telemetry.capture_agent_initialized(
                platform=DevicePlatform.ANDROID.value,
                device_id=self._limrun_instance_id,
            )
            return True

        # Use pre-configured iOS controller
        if ios_controller is not None:
            logger.info("Using pre-configured Limrun iOS controller")
            self._limrun_controller = ios_controller
            self._limrun_instance_id = ios_controller.instance_id

            self._adb_client = None
            self._ui_adb_client = None
            self._ios_client = ios_controller
            self._ios_device_type = DeviceType.LIMRUN

            self._device_context = DeviceContext(
                host_platform="LINUX",
                mobile_platform=DevicePlatform.IOS,
                device_id=self._limrun_instance_id,
                device_width=ios_controller.device_width,
                device_height=ios_controller.device_height,
            )

            logger.info(f"Limrun iOS device ready: {self._limrun_instance_id}")
            logger.info(self._device_context.to_str())
            logger.info("✅ Mobile-use agent initialized with pre-configured Limrun iOS.")
            self._initialized = True
            telemetry.capture_agent_initialized(
                platform=DevicePlatform.IOS.value,
                device_id=self._limrun_instance_id,
            )
            return True

        # Fall back to provisioning via limrun_config
        from minitap.mobile_use.clients.limrun_factory import (
            LimrunInstanceConfig,
            create_limrun_android_instance,
            create_limrun_ios_instance,
        )

        limrun_config = self._config.limrun_config
        if limrun_config is None:
            raise ValueError("limrun_config is not set and no pre-configured controller provided")

        logger.info(f"Provisioning Limrun {limrun_config.platform.value} device...")

        instance_config = LimrunInstanceConfig(
            api_key=limrun_config.api_key,
            base_url=limrun_config.base_url,
            inactivity_timeout=limrun_config.inactivity_timeout,
            hard_timeout=limrun_config.hard_timeout,
            display_name=limrun_config.display_name,
            labels=limrun_config.labels,
        )

        if limrun_config.platform == LimrunPlatform.ANDROID:
            instance, controller = await create_limrun_android_instance(instance_config)
            self._limrun_instance_id = instance.metadata.id
            self._limrun_controller = controller
            try:
                await controller.connect()
            except Exception:
                await self._cleanup_limrun_device()
                raise

            self._adb_client = controller._adb_client
            self._ui_adb_client = controller._ui_client
            self._ios_client = None
            self._ios_device_type = None

            self._device_context = DeviceContext(
                host_platform="LINUX",
                mobile_platform=DevicePlatform.ANDROID,
                device_id=controller._adb_serial or self._limrun_instance_id,
                device_width=controller.device_width,
                device_height=controller.device_height,
            )
        else:
            instance, controller, limrun_ctrl = await create_limrun_ios_instance(instance_config)
            self._limrun_instance_id = instance.metadata.id
            self._limrun_controller = limrun_ctrl  # Store underlying controller for cleanup

            self._adb_client = None
            self._ui_adb_client = None
            self._ios_client = controller.ios_client  # Use the LimrunIosController as ios_client
            self._ios_device_type = DeviceType.LIMRUN

            self._device_context = DeviceContext(
                host_platform="LINUX",
                mobile_platform=DevicePlatform.IOS,
                device_id=self._limrun_instance_id,
                device_width=controller.device_width,
                device_height=controller.device_height,
            )

        logger.info(
            f"Limrun {limrun_config.platform.value} device ready: {self._limrun_instance_id}"
        )
        logger.info(self._device_context.to_str())
        logger.info("✅ Mobile-use agent initialized with Limrun.")
        self._initialized = True
        telemetry.capture_agent_initialized(
            platform=limrun_config.platform.value,
            device_id=self._limrun_instance_id,
        )
        return True

    async def _cleanup_limrun_device(self) -> None:
        """Cleanup Limrun device resources."""
        from minitap.mobile_use.clients.limrun_factory import (
            LimrunInstanceConfig,
            delete_limrun_android_instance,
            delete_limrun_ios_instance,
        )

        # Always cleanup controller if present
        if self._limrun_controller:
            logger.info("Cleaning up Limrun controller...")
            await self._limrun_controller.cleanup()
            self._limrun_controller = None

        # Only attempt instance deletion if limrun_config is present
        if self._config.limrun_config and self._limrun_instance_id:
            logger.info(f"Deleting Limrun instance: {self._limrun_instance_id}")
            limrun_config = self._config.limrun_config
            instance_config = LimrunInstanceConfig(
                api_key=limrun_config.api_key,
                base_url=limrun_config.base_url,
            )

            try:
                if limrun_config.platform == LimrunPlatform.ANDROID:
                    await delete_limrun_android_instance(instance_config, self._limrun_instance_id)
                else:
                    await delete_limrun_ios_instance(instance_config, self._limrun_instance_id)
            except Exception as e:
                logger.warning(f"Failed to delete Limrun instance: {e}")
        elif self._limrun_instance_id:
            logger.info(
                f"Skipping Limrun instance deletion (no limrun_config): {self._limrun_instance_id}"
            )

        self._limrun_instance_id = None

    async def _prepare_app_installation(self, task: Task) -> str | None:
        """Install app if app_path is specified in the task request.

        Returns:
            The bundle ID of the installed app (iOS only), or None.
        """
        if not task.request.app_path:
            return None

        task_name = task.get_name()
        logger.info(f"[{task_name}] Installing app from: {task.request.app_path}")

        bundle_id = await self.install_app(task.request.app_path)

        if bundle_id:
            logger.info(f"[{task_name}] App installed with bundle ID: {bundle_id}")
            # If locked_app_package is not set, automatically lock to the installed app
            if not task.request.locked_app_package:
                logger.info(f"[{task_name}] Auto-locking to installed app: {bundle_id}")
                task.request.locked_app_package = bundle_id

        return bundle_id

    async def _prepare_app_lock(self, task: Task, context: MobileUseContext):
        """Prepare app lock by launching the locked app if specified."""
        if not task.request.locked_app_package:
            return

        task_name = task.get_name()
        logger.info(f"[{task_name}] Preparing app lock for: {task.request.locked_app_package}")

        app_lock_status = await _handle_initial_app_launch(
            ctx=context, locked_app_package=task.request.locked_app_package
        )

        if context.execution_setup is None:
            context.execution_setup = ExecutionSetup(app_lock_status=app_lock_status)
        else:
            context.execution_setup.app_lock_status = app_lock_status

        if app_lock_status.locked_app_initial_launch_success is False:
            error = app_lock_status.locked_app_initial_launch_error
            logger.warning(f"[{task_name}] Failed to launch locked app: {error}")

    def _prepare_tracing(self, task: Task, context: MobileUseContext):
        """Prepare tracing setup if record_trace is enabled."""
        if not task.request.record_trace:
            return

        task_name = task.get_name()
        temp_trace_path = Path(self._tmp_traces_dir / task_name).resolve()
        traces_output_path = Path(task.request.trace_path).resolve()
        logger.info(f"[{task_name}] 📂 Traces output path: {traces_output_path}")
        logger.info(f"[{task_name}] 📄📂 Traces temp path: {temp_trace_path}")
        traces_output_path.mkdir(parents=True, exist_ok=True)
        temp_trace_path.mkdir(parents=True, exist_ok=True)

        context.execution_setup = ExecutionSetup(
            traces_path=self._tmp_traces_dir,
            trace_name=task_name,
            enable_remote_tracing=task.request.enable_remote_tracing,
        )

    async def _finalize_tracing(self, task: Task, context: MobileUseContext):
        exec_setup_ctx = context.execution_setup
        if not exec_setup_ctx:
            return

        if exec_setup_ctx.traces_path is None or exec_setup_ctx.trace_name is None:
            return

        task_name = task.get_name()
        status = "_PASS" if task.status == "completed" else "_FAIL"
        ts = task.created_at.strftime("%Y-%m-%dT%H-%M-%S")
        new_name = f"{exec_setup_ctx.trace_name}{status}_{ts}"

        temp_trace_path = (self._tmp_traces_dir / exec_setup_ctx.trace_name).resolve()
        traces_output_path = Path(task.request.trace_path).resolve()

        logger.info(f"[{task_name}] Compiling trace FROM FOLDER: " + str(temp_trace_path))
        create_gif_from_trace_folder(temp_trace_path)
        create_steps_json_from_trace_folder(temp_trace_path)

        if exec_setup_ctx.enable_remote_tracing:
            gif_path = temp_trace_path / "trace.gif"
            if gif_path.exists() and self._platform_service:
                try:
                    task_run_id = await self._platform_service.upload_trace_gif(
                        task_run_id=task.id, gif_path=gif_path
                    )
                    if task_run_id:
                        platform_url = f"{settings.MINITAP_BASE_URL}/task-runs/{task_run_id}"
                        logger.info(f"[{task_name}] 🌐 View on platform: {platform_url}")
                except Exception as e:
                    logger.warning(f"[{task_name}] Failed to upload trace GIF: {e}")

        logger.info(f"[{task_name}] Video created, removing dust...")
        remove_images_from_trace_folder(temp_trace_path)
        remove_steps_json_from_trace_folder(temp_trace_path)
        logger.info(f"[{task_name}] 📽️ Trace compiled, moving to output path 📽️")

        output_folder_path = temp_trace_path.rename(traces_output_path / new_name).resolve()
        logger.info(f"[{task_name}] 📂✅ Traces located in: {output_folder_path}")

    def _prepare_output_files(self, task: Task):
        if task.request.llm_output_path:
            _validate_and_prepare_file(file_path=task.request.llm_output_path)
        if task.request.thoughts_output_path:
            _validate_and_prepare_file(file_path=task.request.thoughts_output_path)

    async def _extract_output(
        self,
        task_name: str,
        ctx: MobileUseContext,
        request: TaskRequest[TOutput],
        output_config: OutputConfig | None,
        state: State,
    ) -> str | dict | TOutput | None:
        if output_config and output_config.needs_structured_format():
            logger.info(f"[{task_name}] Generating structured output...")
            try:
                structured_output = await outputter(
                    ctx=ctx,
                    output_config=output_config,
                    graph_output=state,
                )
                logger.info(f"[{task_name}] Structured output: {structured_output}")
                record_events(output_path=request.llm_output_path, events=structured_output)
                if request.output_format is not None and request.output_format is not NoneType:
                    return request.output_format.model_validate(structured_output)
                return structured_output
            except Exception as e:
                logger.error(f"[{task_name}] Failed to generate structured output: {e}")
                return None
        if state and state.agents_thoughts:
            last_msg = state.agents_thoughts[-1]
            logger.info(str(last_msg))
            record_events(output_path=request.llm_output_path, events=last_msg)
            return last_msg
        return None

    def _get_graph_state(self, task: Task):
        return State(
            messages=[],
            initial_goal=task.request.goal,
            subgoal_plan=[],
            latest_ui_hierarchy=None,
            latest_screenshot=None,
            focused_app_info=None,
            device_date=None,
            structured_decisions=None,
            complete_subgoals_by_ids=[],
            agents_thoughts=[],
            remaining_steps=task.request.max_steps,
            executor_messages=[],
            cortex_last_thought=None,
            scratchpad={},
        )

    def _init_clients(
        self,
        device_id: str,
        platform: DevicePlatform,
        ios_device_type: DeviceType | None,
        retry_count: int,
        retry_wait_seconds: int,
    ):
        self._adb_client = (
            AdbClient(host=self._config.servers.adb_host, port=self._config.servers.adb_port)
            if platform == DevicePlatform.ANDROID
            else None
        )
        self._ui_adb_client = (
            UIAutomatorClient(device_id=device_id) if platform == DevicePlatform.ANDROID else None
        )

        # Initialize iOS client using factory (auto-detects device type if not provided)
        if platform == DevicePlatform.IOS:
            self._ios_client = get_ios_client(
                udid=device_id,
                config=self._config.ios_client_config,
            )
            self._ios_device_type = ios_device_type or (
                DeviceType.PHYSICAL
                if isinstance(self._ios_client, WdaClientWrapper)
                else DeviceType.SIMULATOR
            )
        else:
            self._ios_client = None
            self._ios_device_type = None

    def _run_servers(self, device_id: str, platform: DevicePlatform) -> bool:
        if platform == DevicePlatform.ANDROID:
            if not self._ui_adb_client:
                error_msg = (
                    "UIAutomator client is required for Android but not available. "
                    "Please ensure UIAutomator2 is properly installed and configured."
                )
                logger.error(error_msg)
                raise ValueError(error_msg)
            logger.info("✓ UIAutomator client available for Android")
        elif platform == DevicePlatform.IOS:
            if not self._ios_client:
                error_msg = (
                    "iOS client is required but not available. "
                    "Ensure idb (simulators) or WDA (physical) is available."
                )
                logger.error(error_msg)
                raise ValueError(error_msg)
            if isinstance(self._ios_client, WdaClientWrapper):
                client_type = "WDA"
            elif isinstance(self._ios_client, BrowserStackClientWrapper):
                client_type = "BrowserStack"
            else:
                client_type = "IDB"
            logger.info(f"✓ iOS client available ({client_type})")

        return True

    async def _get_device_context(
        self,
        device_id: str,
        platform: DevicePlatform,
    ) -> DeviceContext:
        from platform import system

        host_platform = system()

        # Get real device dimensions from the device
        if platform == DevicePlatform.ANDROID:
            if self._ui_adb_client:
                try:
                    # Use UIAutomator to get actual screen dimensions
                    screen_data = self._ui_adb_client.get_screen_data()
                    device_width = screen_data.width
                    device_height = screen_data.height
                    logger.info(
                        f"Retrieved Android screen dimensions: {device_width}x{device_height}"
                    )
                except Exception as e:
                    logger.warning(f"Failed to get Android screen dimensions: {e}, using defaults")
                    device_width, device_height = 1080, 2340
            else:
                logger.warning("UIAutomator client not available, using default dimensions")
                device_width, device_height = 1080, 2340
        else:  # iOS
            if self._ios_client:
                try:
                    # Use iOS client to take a screenshot and get dimensions
                    screenshot_data = await self._ios_client.screenshot()  # type: ignore[call-arg]
                    if screenshot_data:
                        img = Image.open(BytesIO(screenshot_data))
                        device_width = img.width
                        device_height = img.height
                        logger.info(
                            f"Retrieved iOS screen dimensions: {device_width}x{device_height}"
                        )
                    else:
                        logger.warning("IDB screenshot returned None, using default dimensions")
                        device_width, device_height = 375, 812
                except Exception as e:
                    logger.warning(f"Failed to get iOS screen dimensions: {e}, using defaults")
                    device_width, device_height = 375, 812
            else:
                logger.warning("IDB client not available, using default dimensions")
                device_width, device_height = 375, 812

        return DeviceContext(
            host_platform="WINDOWS" if host_platform == "Windows" else "LINUX",
            mobile_platform=platform,
            device_id=device_id,
            device_width=device_width,
            device_height=device_height,
        )

    def _get_task_status_change_callback(
        self,
        task_info: PlatformTaskInfo,
        platform_service: PlatformService | None = None,
    ) -> Callable[[TaskRunStatus, str | None, Any | None], Coroutine]:
        service = platform_service or self._platform_service

        async def change_status(
            status: TaskRunStatus,
            message: str | None = None,
            output: Any | None = None,
        ):
            if not service:
                raise PlatformServiceUninitializedError()
            try:
                await service.update_task_run_status(
                    task_run_id=task_info.task_run.id,
                    status=status,
                    message=message,
                    output=output,
                )
            except Exception as e:
                logger.error(f"Failed to update task run status: {e}")

        return change_status

    def _get_plan_changes_callback(
        self,
        task_info: PlatformTaskInfo,
        platform_service: PlatformService | None = None,
    ) -> Callable[[list[Subgoal], IsReplan], Coroutine]:
        service = platform_service or self._platform_service
        current_plan: TaskRunPlanResponse | None = None

        async def update_plan(plan: list[Subgoal], is_replan: IsReplan):
            nonlocal current_plan

            if not service:
                raise PlatformServiceUninitializedError()
            try:
                if is_replan and current_plan:
                    # End previous plan
                    await service.upsert_task_run_plan(
                        task_run_id=task_info.task_run.id,
                        started_at=current_plan.started_at,
                        plan=plan,
                        ended_at=datetime.now(UTC),
                        plan_id=current_plan.id,
                    )
                    current_plan = None

                current_plan = await service.upsert_task_run_plan(
                    task_run_id=task_info.task_run.id,
                    started_at=current_plan.started_at if current_plan else datetime.now(UTC),
                    plan=plan,
                    ended_at=current_plan.ended_at if current_plan else None,
                    plan_id=current_plan.id if current_plan else None,
                )
            except Exception as e:
                logger.error(f"Failed to update plan: {e}")

        return update_plan

    def _get_new_agent_thought_callback(
        self,
        task_info: PlatformTaskInfo,
        platform_service: PlatformService | None = None,
    ) -> Callable[[AgentNode, str], Coroutine]:
        service = platform_service or self._platform_service

        async def add_agent_thought(agent: AgentNode, thought: str):
            if not service:
                raise PlatformServiceUninitializedError()
            try:
                await service.add_agent_thought(
                    task_run_id=task_info.task_run.id,
                    agent=agent,
                    thought=thought,
                )
            except Exception as e:
                logger.error(f"Failed to add agent thought: {e}")

        return add_agent_thought


def _validate_and_prepare_file(file_path: Path):
    path_obj = Path(file_path)
    if path_obj.exists() and path_obj.is_dir():
        raise AgentTaskRequestError(f"Error: Path '{file_path}' is a directory, not a file.")
    try:
        path_obj.parent.mkdir(parents=True, exist_ok=True)
        path_obj.touch(exist_ok=True)
    except OSError as e:
        raise AgentTaskRequestError(f"Error creating file '{file_path}': {e}")


def print_ai_response_to_stderr(graph_result: State):
    for msg in reversed(graph_result.messages):
        if isinstance(msg, AIMessage):
            print(msg.content, file=sys.stderr)
            return
