import asyncio
import os
from enum import StrEnum
from shutil import which
from typing import Annotated

import typer
from adbutils import AdbClient
from langchain.callbacks.base import Callbacks
from rich.console import Console

from minitap.mobile_use.clients.ios_client_config import (
    IdbClientConfig,
    IosClientConfig,
    WdaClientConfig,
)
from minitap.mobile_use.clients.limrun_factory import (
    LimrunInstanceConfig,
    LimrunPlatform,
    create_limrun_android_instance,
    create_limrun_ios_instance,
    delete_limrun_android_instance,
    delete_limrun_ios_instance,
)
from minitap.mobile_use.config import initialize_llm_config, settings
from minitap.mobile_use.sdk import Agent
from minitap.mobile_use.sdk.builders import Builders
from minitap.mobile_use.sdk.types.task import AgentProfile
from minitap.mobile_use.services.telemetry import telemetry
from minitap.mobile_use.utils.cli_helpers import display_device_status
from minitap.mobile_use.utils.logger import get_logger
from minitap.mobile_use.utils.video import check_ffmpeg_available

app = typer.Typer(add_completion=False, pretty_exceptions_enable=False)
logger = get_logger(__name__)


class DeviceType(StrEnum):
    """Device type for mobile-use agent."""

    LOCAL = "local"
    LIMRUN = "limrun"


async def run_automation(
    goal: str,
    locked_app_package: str | None = None,
    test_name: str | None = None,
    traces_output_path_str: str = "traces",
    output_description: str | None = None,
    graph_config_callbacks: Callbacks = [],
    video_recording_tools_enabled: bool = False,
    wda_url: str | None = None,
    wda_timeout: float | None = None,
    wda_auto_start_iproxy: bool | None = None,
    wda_auto_start_wda: bool | None = None,
    wda_project_path: str | None = None,
    wda_startup_timeout: float | None = None,
    idb_host: str | None = None,
    idb_port: int | None = None,
    device_type: DeviceType = DeviceType.LOCAL,
    limrun_platform: LimrunPlatform | None = None,
):
    llm_config = initialize_llm_config()
    agent_profile = AgentProfile(name="default", llm_config=llm_config)
    config = Builders.AgentConfig.with_default_profile(profile=agent_profile)
    if video_recording_tools_enabled:
        config.with_video_recording_tools()

    # Limrun device provisioning
    limrun_instance_id: str | None = None
    limrun_controller = None
    limrun_config: LimrunInstanceConfig | None = None

    if device_type == DeviceType.LIMRUN:
        if limrun_platform is None:
            raise ValueError("--limrun-platform is required when using --device-type limrun")

        logger.info(f"Provisioning Limrun {limrun_platform.value} device...")
        limrun_config = LimrunInstanceConfig()

        if limrun_platform == LimrunPlatform.ANDROID:
            instance, limrun_controller = await create_limrun_android_instance(limrun_config)
            limrun_instance_id = instance.metadata.id
            await limrun_controller.connect()
            config.with_limrun_android_controller(limrun_controller)
        else:
            instance, _, limrun_controller = await create_limrun_ios_instance(limrun_config)
            limrun_instance_id = instance.metadata.id
            # Connection is done in the factory, no need to call connect()
            config.with_limrun_ios_controller(limrun_controller)

        logger.info(f"Limrun {limrun_platform.value} device ready: {limrun_instance_id}")
    else:
        # Build iOS client config from CLI options (local device)
        wda_config = WdaClientConfig.with_overrides(
            wda_url=wda_url,
            timeout=wda_timeout,
            auto_start_iproxy=wda_auto_start_iproxy,
            auto_start_wda=wda_auto_start_wda,
            wda_project_path=wda_project_path,
            wda_startup_timeout=wda_startup_timeout,
        )
        idb_config = IdbClientConfig.with_overrides(host=idb_host, port=idb_port)
        config.with_ios_client_config(IosClientConfig(wda=wda_config, idb=idb_config))

        if settings.ADB_HOST:
            config.with_adb_server(host=settings.ADB_HOST, port=settings.ADB_PORT)

    if graph_config_callbacks:
        config.with_graph_config_callbacks(graph_config_callbacks)

    agent: Agent | None = None
    try:
        agent = Agent(config=config.build())
        await agent.init(
            retry_count=int(os.getenv("MOBILE_USE_HEALTH_RETRIES", 5)),
            retry_wait_seconds=int(os.getenv("MOBILE_USE_HEALTH_DELAY", 2)),
        )

        task = agent.new_task(goal)
        if locked_app_package:
            task.with_locked_app_package(locked_app_package)
        if test_name:
            task.with_name(test_name).with_trace_recording(path=traces_output_path_str)
        if output_description:
            task.with_output_description(output_description)

        agent_thoughts_path = os.getenv("EVENTS_OUTPUT_PATH", None)
        llm_result_path = os.getenv("RESULTS_OUTPUT_PATH", None)
        if agent_thoughts_path:
            task.with_thoughts_output_saving(path=agent_thoughts_path)
        if llm_result_path:
            task.with_llm_output_saving(path=llm_result_path)

        await agent.run_task(request=task.build())
    finally:
        if agent is not None:
            await agent.clean()

        # Cleanup Limrun device
        if limrun_instance_id and limrun_config:
            logger.info(f"Cleaning up Limrun device: {limrun_instance_id}")
            if limrun_controller:
                await limrun_controller.cleanup()
            if limrun_platform == LimrunPlatform.ANDROID:
                await delete_limrun_android_instance(limrun_config, limrun_instance_id)
            else:
                await delete_limrun_ios_instance(limrun_config, limrun_instance_id)


@app.command()
def main(
    goal: Annotated[str, typer.Argument(help="The main goal for the agent to achieve.")],
    test_name: Annotated[
        str | None,
        typer.Option(
            "--test-name",
            "-n",
            help="A name for the test recording. If provided, a trace will be saved.",
        ),
    ] = None,
    traces_path: Annotated[
        str,
        typer.Option(
            "--traces-path",
            "-p",
            help="The path to save the traces.",
        ),
    ] = "traces",
    output_description: Annotated[
        str | None,
        typer.Option(
            "--output-description",
            "-o",
            help=(
                """
                A dict output description for the agent.
                Ex: a JSON schema with 2 keys: type, price
                """
            ),
        ),
    ] = None,
    wda_url: Annotated[
        str | None,
        typer.Option(
            "--wda-url",
            help="Override WebDriverAgent URL (e.g. http://localhost:8100).",
        ),
    ] = None,
    wda_timeout: Annotated[
        float | None,
        typer.Option(
            "--wda-timeout",
            help="Timeout (seconds) for WDA operations.",
        ),
    ] = None,
    wda_auto_start_iproxy: Annotated[
        bool | None,
        typer.Option(
            "--wda-auto-start-iproxy/--no-wda-auto-start-iproxy",
            help="Auto-start iproxy if not running.",
        ),
    ] = None,
    wda_auto_start_wda: Annotated[
        bool | None,
        typer.Option(
            "--wda-auto-start-wda/--no-wda-auto-start-wda",
            help="Auto-build and run WDA via xcodebuild if not responding.",
        ),
    ] = None,
    wda_project_path: Annotated[
        str | None,
        typer.Option(
            "--wda-project-path",
            help="Path to WebDriverAgent.xcodeproj.",
        ),
    ] = None,
    wda_startup_timeout: Annotated[
        float | None,
        typer.Option(
            "--wda-startup-timeout",
            help="Timeout (seconds) while waiting for WDA to start.",
        ),
    ] = None,
    idb_host: Annotated[
        str | None,
        typer.Option(
            "--idb-host",
            help="IDB companion host (for simulators).",
        ),
    ] = None,
    idb_port: Annotated[
        int | None,
        typer.Option(
            "--idb-port",
            help="IDB companion port (for simulators).",
        ),
    ] = None,
    with_video_recording_tools: Annotated[
        bool,
        typer.Option(
            "--with-video-recording-tools",
            help="Enable AI agents to use video recording tools"
            " to analyze dynamic content on the screen.",
        ),
    ] = False,
    device_type: Annotated[
        DeviceType,
        typer.Option(
            "--device-type",
            "-d",
            help="Device type: 'local' for connected devices, 'limrun' for cloud devices.",
        ),
    ] = DeviceType.LOCAL,
    limrun_platform: Annotated[
        LimrunPlatform | None,
        typer.Option(
            "--limrun-platform",
            help="Platform for Limrun cloud device: 'android' or 'ios'. "
            "Required when --device-type is 'limrun'.",
        ),
    ] = None,
):
    """
    Run the Mobile-use agent to automate tasks on a mobile device.
    """
    if with_video_recording_tools:
        check_ffmpeg_available()

    console = Console()

    if device_type == DeviceType.LOCAL:
        adb_client = None
        try:
            if which("adb"):
                adb_client = AdbClient(
                    host=settings.ADB_HOST or "localhost",
                    port=settings.ADB_PORT or 5037,
                )
        except Exception:
            pass  # ADB not available, will only support iOS devices

        display_device_status(console, adb_client=adb_client)
    else:
        if limrun_platform is None:
            console.print(
                "[red]Error: --limrun-platform is required when using --device-type limrun[/red]"
            )
            raise typer.Exit(1)
        console.print(f"[cyan]Using Limrun cloud device ({limrun_platform.value})...[/cyan]")

    # Start telemetry session with CLI context (only non-sensitive flags)
    session_id = telemetry.start_session(
        {
            "source": "cli",
            "has_output_description": output_description is not None,
        }
    )

    error_message = None
    cancelled = False
    try:
        asyncio.run(
            run_automation(
                goal=goal,
                test_name=test_name,
                traces_output_path_str=traces_path,
                output_description=output_description,
                wda_url=wda_url,
                wda_timeout=wda_timeout,
                wda_auto_start_iproxy=wda_auto_start_iproxy,
                wda_auto_start_wda=wda_auto_start_wda,
                wda_project_path=wda_project_path,
                wda_startup_timeout=wda_startup_timeout,
                idb_host=idb_host,
                idb_port=idb_port,
                video_recording_tools_enabled=with_video_recording_tools,
                device_type=device_type,
                limrun_platform=limrun_platform,
            )
        )
    except KeyboardInterrupt:
        cancelled = True
        error_message = "Task cancelled by user"
    except Exception as e:
        error_message = str(e)
        console.print(
            f"\n[dim]If you need support, please include this session ID: {session_id}[/dim]"
        )
        raise
    finally:
        telemetry.end_session(
            success=error_message is None,
            error=error_message,
        )
        if cancelled:
            raise SystemExit(130)


def _prompt_telemetry_consent(console: Console) -> None:
    """Prompt user for telemetry consent if not yet configured."""
    if not telemetry.needs_consent:
        return

    console.print()
    console.print("[bold]📊 Help improve mobile-use[/bold]")
    console.print(
        "We collect anonymous usage data to help debug and improve the SDK.\n"
        "No personal data, prompts, or device content is collected.\n"
        "You can change this anytime by setting MOBILE_USE_TELEMETRY_ENABLED=false\n"
    )

    try:
        import inquirer

        questions = [
            inquirer.Confirm(
                "consent",
                message="Enable anonymous telemetry?",
                default=True,
            )
        ]
        answers = inquirer.prompt(questions)
        if answers is not None:
            enabled = answers.get("consent", False)
            telemetry.set_consent(enabled)
            if enabled:
                console.print("[green]✓ Telemetry enabled. Thank you![/green]\n")
            else:
                console.print("[dim]Telemetry disabled.[/dim]\n")
        else:
            telemetry.set_consent(False)
    except (ImportError, KeyboardInterrupt):
        telemetry.set_consent(False)


def cli():
    console = Console()
    _prompt_telemetry_consent(console)
    telemetry.initialize()
    try:
        app()
    finally:
        telemetry.shutdown()


if __name__ == "__main__":
    cli()
