"""Multi-provider VLM abstraction for AI agent perception.

Supports OpenAI-compatible (GPT-4o, Claude, Qwen-VL), Google Gemini,
and local models via litellm or raw HTTP.
"""

from __future__ import annotations

import base64
import json
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class VLMProviderConfig:
    provider: str = "openai"  # openai | gemini | qwen | local
    api_key: str = ""
    base_url: str = ""
    model: str = ""
    max_tokens: int = 1024
    temperature: float = 0.0
    extra: dict[str, Any] = field(default_factory=dict)


class VLMProvider(ABC):
    """Abstract base for multimodal LLM providers."""

    @abstractmethod
    async def analyze_screen(
        self,
        screenshot_b64: str,
        instruction: str,
        ui_tree: str | None = None,
        history: list[dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        """Analyze a screenshot and return action plan.

        Returns a dict with at minimum:
        - action: str (tap, swipe, type, back, home, scroll, done, error)
        - reasoning: str
        - parameters: dict (x, y, text, direction, etc.)
        """


class OpenAICompatibleProvider(VLMProvider):
    """Works with OpenAI, Azure OpenAI, Anthropic (via proxy), Qwen-VL."""

    def __init__(self, config: VLMProviderConfig) -> None:
        self._config = config

    async def analyze_screen(
        self,
        screenshot_b64: str,
        instruction: str,
        ui_tree: str | None = None,
        history: list[dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        try:
            import httpx
        except ImportError:
            raise ImportError("httpx is required for VLM provider: pip install httpx")

        base_url = self._config.base_url or "https://api.openai.com/v1"
        url = f"{base_url.rstrip('/')}/chat/completions"

        system_prompt = (
            "You are an AI agent controlling a mobile device. "
            "Analyze the screenshot and UI state, then decide the next action.\n\n"
            "Respond in JSON with:\n"
            '- "action": one of tap, swipe, type, back, home, scroll, long_press, done, error\n'
            '- "reasoning": brief explanation\n'
            '- "parameters": {x, y, text, direction, duration_ms} as needed\n\n'
            f"User instruction: {instruction}"
        )

        messages: list[dict] = [{"role": "system", "content": system_prompt}]

        if history:
            for h in history[-6:]:
                messages.append(h)

        content: list[dict] = []
        if ui_tree:
            content.append({"type": "text", "text": f"UI Tree:\n{ui_tree[:4000]}"})
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{screenshot_b64}", "detail": "high"},
        })
        content.append({"type": "text", "text": "What action should I take next?"})
        messages.append({"role": "user", "content": content})

        async with httpx.AsyncClient(timeout=60) as client:
            resp = await client.post(
                url,
                json={
                    "model": self._config.model,
                    "messages": messages,
                    "max_tokens": self._config.max_tokens,
                    "temperature": self._config.temperature,
                },
                headers={
                    "Authorization": f"Bearer {self._config.api_key}",
                    "Content-Type": "application/json",
                },
            )
            resp.raise_for_status()
            data = resp.json()

        text = data["choices"][0]["message"]["content"]
        text = text.strip()
        if text.startswith("```"):
            text = text.split("\n", 1)[1] if "\n" in text else text[3:]
            if text.endswith("```"):
                text = text[:-3]
            text = text.strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {"action": "error", "reasoning": text, "parameters": {}}


class OllamaProvider(VLMProvider):
    """VLM provider using local Ollama instance for GPU-accelerated inference."""

    def __init__(self, config: VLMProviderConfig) -> None:
        self._config = config

    async def analyze_screen(
        self,
        screenshot_b64: str,
        instruction: str,
        ui_tree: str | None = None,
        history: list[dict[str, Any]] | None = None,
    ) -> dict[str, Any]:
        try:
            import httpx
        except ImportError:
            raise ImportError("httpx is required: pip install httpx")

        base_url = self._config.base_url or "http://localhost:11434"
        url = f"{base_url.rstrip('/')}/api/chat"

        system_prompt = (
            "You are an AI agent controlling a device. "
            "Analyze the screenshot and decide the next action.\n"
            "Respond in JSON: {action, reasoning, parameters}\n"
            "Actions: tap, swipe, type, back, home, scroll, long_press, done, error\n"
            f"Instruction: {instruction}"
        )

        messages: list[dict] = [{"role": "system", "content": system_prompt}]

        content = f"What action should I take next?"
        if ui_tree:
            content = f"UI Tree:\n{ui_tree[:3000]}\n\n{content}"

        messages.append({
            "role": "user",
            "content": content,
            "images": [screenshot_b64],
        })

        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(url, json={
                "model": self._config.model,
                "messages": messages,
                "stream": False,
                "format": "json",
                "options": {
                    "temperature": self._config.temperature,
                    "num_predict": self._config.max_tokens,
                },
            })
            resp.raise_for_status()
            data = resp.json()

        text = data.get("message", {}).get("content", "")
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {"action": "error", "reasoning": text, "parameters": {}}


def create_vlm_provider(config: VLMProviderConfig) -> VLMProvider:
    """Factory to create the appropriate VLM provider."""
    if config.provider == "ollama":
        return OllamaProvider(config)
    if config.provider in ("openai", "azure", "qwen", "anthropic", "local"):
        return OpenAICompatibleProvider(config)
    raise ValueError(f"Unknown VLM provider: {config.provider}")
