"""Task Engine — managed AI automation execution with lifecycle.

Uses MobileAgent v3.5 (bundled locally) as primary AI backend,
with a built-in VLM loop as legacy fallback.
"""

from __future__ import annotations

import asyncio
import base64
import logging
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Awaitable

from ..runtime.base import DeviceRuntime

logger = logging.getLogger(__name__)


class TaskStatus(str, Enum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class Task:
    id: str
    instruction: str
    status: TaskStatus = TaskStatus.QUEUED
    steps: list[dict[str, Any]] = field(default_factory=list)
    result: str = ""
    error: str = ""
    created_at: float = field(default_factory=time.time)
    started_at: float = 0.0
    completed_at: float = 0.0

    @property
    def elapsed(self) -> float:
        if self.started_at == 0:
            return 0
        end = self.completed_at if self.completed_at else time.time()
        return end - self.started_at


StepNotifier = Callable[[str, dict[str, Any]], Awaitable[None]]


class TaskEngine:
    """Manages AI automation tasks with queuing and lifecycle.

    Supports two execution backends:
    1. MobileAgent v3.5 (bundled, primary) — specialized prompts + multi-turn image history
    2. Built-in VLM loop (legacy fallback) — simpler single-turn approach
    """

    MAX_COMPLETED_TASKS = 100

    def __init__(
        self,
        runtime: DeviceRuntime,
        on_step: StepNotifier | None = None,
        max_concurrent: int = 1,
        max_steps: int = 30,
        vlm_api_key: str = "",
        vlm_base_url: str = "",
        vlm_model: str = "",
        vlm_provider: str = "openai",
        enable_ocr: bool = False,
    ) -> None:
        self._runtime = runtime
        self._on_step = on_step
        self._max_concurrent = max_concurrent
        self._max_steps = max_steps
        self._vlm_api_key = vlm_api_key
        self._vlm_base_url = vlm_base_url
        self._vlm_model = vlm_model
        self._vlm_provider = vlm_provider
        self._enable_ocr = enable_ocr
        self._tasks: dict[str, Task] = {}
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._cancel_events: dict[str, asyncio.Event] = {}
        self._workers: list[asyncio.Task] = []
        self._started = False

    async def start(self) -> None:
        if self._started:
            return
        self._started = True
        for i in range(self._max_concurrent):
            worker = asyncio.create_task(self._worker(i))
            self._workers.append(worker)

    async def stop(self) -> None:
        self._started = False
        for ce in self._cancel_events.values():
            ce.set()
        for w in self._workers:
            w.cancel()
        self._workers.clear()

    async def submit(self, instruction: str) -> Task:
        task_id = str(uuid.uuid4())
        task = Task(id=task_id, instruction=instruction)
        self._tasks[task_id] = task
        self._cancel_events[task_id] = asyncio.Event()
        await self._queue.put(task_id)
        return task

    async def cancel(self, task_id: str) -> bool:
        if task_id in self._cancel_events:
            self._cancel_events[task_id].set()
            return True
        return False

    def get_task(self, task_id: str) -> Task | None:
        return self._tasks.get(task_id)

    def list_tasks(self, limit: int = 50) -> list[Task]:
        tasks = sorted(self._tasks.values(), key=lambda t: t.created_at, reverse=True)
        return tasks[:limit]

    async def _worker(self, worker_id: int) -> None:
        logger.info("Task worker %d started", worker_id)
        while self._started:
            try:
                task_id = await asyncio.wait_for(self._queue.get(), timeout=5.0)
            except asyncio.TimeoutError:
                continue

            task = self._tasks.get(task_id)
            if not task:
                continue

            task.status = TaskStatus.RUNNING
            task.started_at = time.time()

            cancel_event = self._cancel_events.get(task_id, asyncio.Event())

            try:
                await self._execute_task(task, cancel_event)
                if cancel_event.is_set():
                    task.status = TaskStatus.CANCELLED
                else:
                    task.status = TaskStatus.COMPLETED
            except Exception as exc:
                logger.exception("Task %s failed", task_id)
                task.status = TaskStatus.FAILED
                task.error = str(exc)
            finally:
                task.completed_at = time.time()
                self._cancel_events.pop(task_id, None)
                self._evict_old_tasks()

    async def _execute_task(self, task: Task, cancel: asyncio.Event) -> None:
        """Execute a task using MobileAgent v3.5 (primary) or built-in VLM loop."""
        if not self._vlm_api_key:
            task.error = "VLM API key not configured"
            return
        if not self._vlm_model:
            task.error = "VLM model not configured (set --vlm-model or VLM_MODEL)"
            return

        try:
            await self._mobile_agent_loop(task, cancel)
        except Exception as exc:
            logger.warning("MobileAgent failed (%s), falling back to VLM loop", exc)
            task.steps.clear()
            task.error = ""
            await self._vlm_loop(task, cancel)

    async def _mobile_agent_loop(self, task: Task, cancel: asyncio.Event) -> None:
        """Run MobileAgent v3.5 perception-action loop."""
        from .mobile_agent.agent import run_mobile_agent

        async for step in run_mobile_agent(
            runtime=self._runtime,
            instruction=task.instruction,
            api_key=self._vlm_api_key,
            base_url=self._vlm_base_url,
            model=self._vlm_model,
            max_steps=self._max_steps,
        ):
            if cancel.is_set():
                break

            step_data = {
                "index": step.index,
                "action": step.action,
                "reasoning": step.reasoning,
                "parameters": step.action_params,
                "status": step.status,
            }
            task.steps.append(step_data)

            if self._on_step:
                await self._on_step(task.id, step_data)

            if step.status == "completed":
                task.result = step.reasoning or f"Completed in {step.index + 1} steps"
                return
            if step.status == "failed":
                task.error = step.reasoning or "Task failed"
                return

        if not task.result and not task.error:
            task.result = f"Completed in {len(task.steps)} steps"

    def _evict_old_tasks(self) -> None:
        """Remove oldest completed/failed/cancelled tasks if over the limit."""
        terminal = [
            tid for tid, t in self._tasks.items()
            if t.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED)
        ]
        if len(terminal) <= self.MAX_COMPLETED_TASKS:
            return
        terminal.sort(key=lambda tid: self._tasks[tid].completed_at)
        for tid in terminal[: len(terminal) - self.MAX_COMPLETED_TASKS]:
            del self._tasks[tid]

    async def _vlm_loop(self, task: Task, cancel: asyncio.Event) -> None:
        """Built-in VLM perception→action loop with hybrid perception."""
        try:
            from .vlm_provider import VLMProviderConfig, create_vlm_provider
            from .hybrid_perception import HybridPerceptionEngine
        except ImportError:
            task.error = "No AI backend available"
            return

        import os
        config = VLMProviderConfig(
            provider=self._vlm_provider or os.environ.get("VLM_PROVIDER", "openai"),
            api_key=self._vlm_api_key or os.environ.get("VLM_API_KEY", ""),
            base_url=self._vlm_base_url or os.environ.get("VLM_BASE_URL", ""),
            model=self._vlm_model or os.environ.get("VLM_MODEL", ""),
        )

        if config.provider != "ollama" and not config.api_key:
            task.error = "VLM_API_KEY not configured"
            return

        vlm = create_vlm_provider(config)
        perception = HybridPerceptionEngine(
            runtime=self._runtime,
            enable_ocr=self._enable_ocr or os.environ.get("ENABLE_OCR", "false") == "true",
        )
        history: list[dict[str, Any]] = []

        for step_idx in range(self._max_steps):
            if cancel.is_set():
                break

            state = await perception.perceive()
            ui_tree_str = state.to_prompt() if state.ui_elements else None

            result = await vlm.analyze_screen(
                screenshot_b64=state.screenshot_b64,
                instruction=task.instruction,
                ui_tree=ui_tree_str,
                history=history,
            )

            action = result.get("action", "error")
            reasoning = result.get("reasoning", "")
            params = result.get("parameters", {})

            step_data = {
                "index": step_idx,
                "action": action,
                "reasoning": reasoning,
                "parameters": params,
                "status": "running",
            }
            task.steps.append(step_data)

            if self._on_step:
                await self._on_step(task.id, step_data)

            history.append({
                "role": "assistant",
                "content": f"Step {step_idx}: {action} — {reasoning}",
            })

            if action == "done":
                task.result = reasoning or "Task completed"
                return

            if action == "error":
                task.error = reasoning or "Unknown error"
                return

            await self._execute_action(action, params)

            await asyncio.sleep(0.5)

        task.result = f"Reached max steps ({self._max_steps})"

    async def _execute_action(self, action: str, params: dict[str, Any]) -> None:
        x = params.get("x", 0)
        y = params.get("y", 0)

        if action == "tap":
            await self._runtime.tap(x, y)

        elif action == "long_press":
            duration = params.get("duration_ms", 1000)
            await self._runtime.long_press(x, y, duration)

        elif action == "swipe":
            x2 = params.get("x2", x)
            y2 = params.get("y2", y)
            dur = params.get("duration_ms", 300)
            await self._runtime.swipe(x, y, x2, y2, dur)

        elif action == "scroll":
            direction = params.get("direction", "down")
            dy = 500 if direction == "down" else -500
            await self._runtime.swipe(540, 960, 540, 960 + dy, 300)

        elif action == "type":
            text = params.get("text", "")
            if text:
                await self._runtime.type_text(text)

        elif action == "back":
            await self._runtime.press_key("KEYCODE_BACK")

        elif action == "home":
            await self._runtime.press_key("KEYCODE_HOME")

        else:
            logger.warning("Unknown action: %s", action)


def _simplify_ui_tree(nodes: list[dict]) -> str:
    """Convert raw UI tree nodes into a concise text representation."""
    lines = []
    for node in nodes[:100]:
        text = node.get("text", "") or node.get("content-desc", "")
        cls = node.get("class", "")
        bounds = node.get("bounds", "")
        clickable = node.get("clickable", False)
        if text or clickable:
            line = f"[{cls}] '{text}' bounds={bounds}"
            if clickable:
                line += " (clickable)"
            lines.append(line)
    return "\n".join(lines)
