"""MobileAgent v3.5 async agent loop.

Adapted from X-PLUG/MobileAgent Mobile-Agent-v3.5/mobile_use/run_gui_owl_1_5_for_mobile.py.
Converted from synchronous script to an async generator that yields step results,
integrated with DeviceRuntime instead of raw ADB subprocess calls.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
import tempfile
from dataclasses import dataclass, field
from typing import Any, AsyncGenerator

from PIL import Image

from ...runtime.base import DeviceRuntime
from .packages import NAME_PACKAGE_DICT, PACKAGES_NAME_DICT
from .vlm import (
    annotate_screenshot,
    build_messages,
    call_vlm,
    convert_to_openai_messages,
    detect_model_profile,
    parse_uitars_action,
    resolve_app_name,
    smart_resize,
)

logger = logging.getLogger(__name__)


@dataclass
class StepResult:
    index: int
    action: str
    action_params: dict[str, Any] = field(default_factory=dict)
    reasoning: str = ""
    screenshot_path: str | None = None
    annotated_path: str | None = None
    status: str = "running"


# ---------------------------------------------------------------------------
# Action parsing
# ---------------------------------------------------------------------------


def parse_action(output_text: str) -> dict[str, Any]:
    """Extract the action dict from VLM output containing a <tool_call> block."""
    try:
        after_tag = output_text.split("<tool_call>")[1]
        # Find the closing tag or end of JSON
        if "</tool_call>" in after_tag:
            json_str = after_tag.split("</tool_call>")[0].strip()
        elif "}}" in after_tag:
            json_str = after_tag.split("}}")[0] + "}}"
        else:
            json_str = after_tag.strip()

        # Handle newline after <tool_call>\n
        json_str = json_str.strip()
        return json.loads(json_str)
    except (IndexError, json.JSONDecodeError) as e:
        raise ValueError(f"Failed to parse <tool_call> from output: {e}")


def rescale_coordinates(
    params: dict[str, Any],
    device_w: int,
    device_h: int,
    vlm_img_w: int,
    vlm_img_h: int,
) -> dict[str, Any]:
    """Map VLM output coordinates to device screen coordinates.

    Auto-detects coordinate system:
    - If coords are in 0-1000 range (GUI-Owl style) → normalize from 1000-grid
    - If coords exceed 1000 (GPT-4o style pixel coords) → map from VLM image space
    """
    for key in ("coordinate", "coordinate1", "coordinate2"):
        if key not in params or not isinstance(params[key], list) or len(params[key]) < 2:
            continue
        cx, cy = params[key][0], params[key][1]

        if cx <= 1000 and cy <= 1000:
            params[key][0] = int(cx / 1000 * device_w)
            params[key][1] = int(cy / 1000 * device_h)
        else:
            params[key][0] = int(cx / vlm_img_w * device_w) if vlm_img_w else int(cx)
            params[key][1] = int(cy / vlm_img_h * device_h) if vlm_img_h else int(cy)

    return params


# ---------------------------------------------------------------------------
# Main agent loop
# ---------------------------------------------------------------------------


async def run_mobile_agent(
    runtime: DeviceRuntime,
    instruction: str,
    api_key: str,
    base_url: str,
    model: str,
    *,
    max_steps: int = 30,
    add_info: str = "",
    app_resolver_model: str | None = None,
    save_screenshots: bool = False,
) -> AsyncGenerator[StepResult, None]:
    """Run the MobileAgent v3.5 perception-action loop.

    Yields a ``StepResult`` for each step. The caller can cancel by
    breaking out of the generator (triggers cleanup).
    """
    full_instruction = f"{instruction} ({add_info})" if add_info else instruction
    resolver_model = app_resolver_model or model
    profile = detect_model_profile(model)

    work_dir = tempfile.mkdtemp(prefix="beeos_ma_")
    anno_dir = os.path.join(work_dir, "annotated")
    os.makedirs(anno_dir, exist_ok=True)

    history: list[dict[str, str]] = []

    dev_info = await runtime.get_info()
    device_w, device_h = dev_info.width, dev_info.height

    logger.info("Agent loop: model=%s profile=%s device=%dx%d",
                model, profile, device_w, device_h)

    try:
        for step_id in range(max_steps):
            # --- 1. Screenshot -----------------------------------------------
            screenshot_bytes = await runtime.screenshot()
            ss_path = os.path.join(work_dir, f"ss_{step_id}.png")
            with open(ss_path, "wb") as f:
                f.write(screenshot_bytes)

            # --- 2. Build messages & call VLM --------------------------------
            messages = build_messages(ss_path, full_instruction, history, model)
            openai_msgs = convert_to_openai_messages(messages)

            try:
                output_text = await call_vlm(api_key, base_url, model, openai_msgs)
            except RuntimeError as exc:
                yield StepResult(
                    index=step_id, action="error",
                    reasoning=f"VLM call failed: {exc}", status="failed",
                )
                return

            logger.debug("Step %d VLM output: %s", step_id, output_text[:300])

            # --- 3. Parse action (model-specific) ----------------------------
            try:
                if profile == "uitars":
                    action_params = parse_uitars_action(output_text)
                else:
                    parsed = parse_action(output_text)
                    action_params = parsed.get("arguments", parsed)
            except ValueError as exc:
                yield StepResult(
                    index=step_id, action="error",
                    reasoning=f"Parse failed: {exc}",
                    status="failed",
                )
                return

            # --- 4. Rescale coordinates --------------------------------------
            img = Image.open(ss_path)
            vlm_h, vlm_w = smart_resize(
                img.height, img.width,
                factor=28, min_pixels=3136, max_pixels=10035200,
            )
            action_params = rescale_coordinates(
                action_params, device_w, device_h, vlm_w, vlm_h,
            )

            action_type = action_params.get("action", "unknown")

            # Extract reasoning from output text (model-specific)
            reasoning = ""
            if profile == "uitars":
                for line in output_text.strip().splitlines():
                    if line.strip().startswith("Thought:"):
                        reasoning = line.strip()[8:].strip()
                        break
            elif "<tool_call>" in output_text:
                reasoning = output_text.split("<tool_call>")[0].strip()
                if reasoning.startswith("Action:"):
                    reasoning = reasoning[7:].strip()

            # Annotate screenshot for debug
            anno_path = os.path.join(anno_dir, f"anno_{step_id}.png")
            annotate_screenshot(ss_path, action_params, anno_path)

            # --- 5. Execute action -------------------------------------------
            step_status = "running"

            if action_type == "click":
                coord = action_params.get("coordinate", [0, 0])
                await runtime.tap(coord[0], coord[1])

            elif action_type == "long_press":
                coord = action_params.get("coordinate", [0, 0])
                dur = int(action_params.get("time", 0.8) * 1000)
                await runtime.long_press(coord[0], coord[1], dur)

            elif action_type == "type":
                text = action_params.get("text", "")
                if text:
                    coord = action_params.get("coordinate")
                    if coord and len(coord) >= 2:
                        await runtime.tap(coord[0], coord[1])
                        await asyncio.sleep(0.5)
                    await runtime.type_text(text)

            elif action_type in ("swipe", "scroll"):
                c1 = action_params.get("coordinate", [0, 0])
                c2 = action_params.get("coordinate2", c1)
                await runtime.swipe(c1[0], c1[1], c2[0], c2[1], 800)

            elif action_type == "system_button":
                btn = action_params.get("button", "")
                key_map = {
                    "Back": "KEYCODE_BACK",
                    "Home": "KEYCODE_HOME",
                    "Menu": "KEYCODE_APP_SWITCH",
                    "Enter": "KEYCODE_ENTER",
                }
                if btn in key_map:
                    await runtime.press_key(key_map[btn])

            elif action_type == "key":
                key_text = action_params.get("text", "")
                code_map = {
                    "volume_up": "KEYCODE_VOLUME_UP",
                    "volume_down": "KEYCODE_VOLUME_DOWN",
                    "power": "KEYCODE_POWER",
                    "camera": "KEYCODE_CAMERA",
                    "clear": "KEYCODE_CLEAR",
                    "enter": "KEYCODE_ENTER",
                }
                kc = code_map.get(key_text.lower(), f"KEYCODE_{key_text.upper()}")
                await runtime.press_key(kc)

            elif action_type == "open":
                app_name = action_params.get("text", "")
                opened = await _handle_open_app(
                    runtime, app_name, full_instruction,
                    api_key, base_url, resolver_model,
                )
                if not opened:
                    reasoning = f"App '{app_name}' not found on device"

            elif action_type == "wait":
                wait_secs = action_params.get("time", 2)
                await asyncio.sleep(wait_secs)

            elif action_type == "terminate":
                status_val = action_params.get("status", "unknown")
                step_status = "completed" if status_val == "success" else "failed"
                yield StepResult(
                    index=step_id, action=action_type,
                    action_params=action_params, reasoning=reasoning,
                    screenshot_path=ss_path if save_screenshots else None,
                    annotated_path=anno_path if save_screenshots else None,
                    status=step_status,
                )
                return

            elif action_type == "answer":
                if "<tool_call>" in output_text:
                    conclusion = output_text.split("<tool_call>")[0].strip()
                else:
                    conclusion = action_params.get("text", reasoning)
                yield StepResult(
                    index=step_id, action=action_type,
                    action_params=action_params, reasoning=conclusion,
                    screenshot_path=ss_path if save_screenshots else None,
                    annotated_path=anno_path if save_screenshots else None,
                    status="completed",
                )
                return

            elif action_type in ("call_user", "calluser", "interact"):
                reasoning = f"Skipped user interaction: {action_params.get('text', '')}"

            else:
                logger.warning("Unsupported action type: %s", action_type)

            # --- 6. Record history & yield -----------------------------------
            history.append({"output": output_text, "image": ss_path})

            yield StepResult(
                index=step_id, action=action_type,
                action_params=action_params, reasoning=reasoning,
                screenshot_path=ss_path if save_screenshots else None,
                annotated_path=anno_path if save_screenshots else None,
                status=step_status,
            )

            await asyncio.sleep(1.5)

    finally:
        if not save_screenshots:
            shutil.rmtree(work_dir, ignore_errors=True)


# ---------------------------------------------------------------------------
# App launcher helper
# ---------------------------------------------------------------------------


async def _handle_open_app(
    runtime: DeviceRuntime,
    app_name: str,
    instruction: str,
    api_key: str,
    base_url: str,
    resolver_model: str,
) -> bool:
    """Resolve app name → package and launch via ADB monkey."""
    # Get installed packages
    result = await runtime.execute_shell("pm list packages")
    installed: set[str] = set()
    for line in result.splitlines():
        line = line.strip()
        if line.startswith("package:"):
            installed.add(line[8:])

    normalized = app_name.lower().strip().replace(" ", "").replace("-", "")

    # Direct lookup
    for pkg in NAME_PACKAGE_DICT.get(normalized, []):
        if pkg in installed:
            await runtime.execute_shell(
                f"monkey -p {pkg} -c android.intent.category.LAUNCHER 1"
            )
            return True

    # LLM-based resolution
    installed_names: list[str] = []
    for pkg in installed:
        if pkg in PACKAGES_NAME_DICT:
            installed_names.append(PACKAGES_NAME_DICT[pkg][0])

    if installed_names:
        resolved = await resolve_app_name(
            instruction, installed_names, api_key, base_url, resolver_model,
        )
        if resolved:
            resolved_norm = resolved.lower().strip().replace(" ", "").replace("-", "")
            for pkg in NAME_PACKAGE_DICT.get(resolved_norm, []):
                if pkg in installed:
                    await runtime.execute_shell(
                        f"monkey -p {pkg} -c android.intent.category.LAUNCHER 1"
                    )
                    return True

    logger.warning("App '%s' not found on device", app_name)
    return False
