"""
Vision Model Client - Ollama API for image understanding.

Uses multimodal models like LLaVA or Qwen2.5-VL to interpret screenshots.
"""

import json
import re
import time
from typing import Optional, Tuple
import urllib.request
import urllib.error

from .models import Action, ActionType


SYSTEM_PROMPT = """You are a computer use agent. You see a screenshot and must decide what action to take.

Your task: {task}

Respond with a JSON object containing ONE action to take:
{{
  "action_type": "click|double_click|right_click|type|press|hotkey|scroll|move|wait|done|failed",
  "x": <x coordinate for click/move actions>,
  "y": <y coordinate for click/move actions>,
  "text": "<text to type for type action>",
  "key": "<key name for press action>",
  "keys": ["key1", "key2"] for hotkey action,
  "scroll_amount": <integer for scroll, positive=up, negative=down>,
  "duration": <seconds for wait action>,
  "reason": "<brief explanation of why you chose this action>"
}}

Guidelines:
- Look carefully at the screenshot to find the target element
- Use "click" to interact with buttons, links, input fields
- Use "type" after clicking an input field
- Use "hotkey" for keyboard shortcuts like ["ctrl", "c"] or ["alt", "tab"]
- Use "scroll" to see more content (negative to scroll down)
- Use "done" when the task is complete
- Use "failed" if you cannot complete the task
- Coordinates are in pixels from top-left corner

Only output the JSON object, nothing else."""


class OllamaVisionClient:
    """Client for Ollama vision models."""

    def __init__(
        self,
        model: str = "llava:7b",
        base_url: str = "http://localhost:11434",
        timeout: int = 120,
    ):
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    def analyze_screenshot(
        self,
        image_base64: str,
        task: str,
        context: str = "",
    ) -> Tuple[Action, str, float]:
        """
        Send screenshot to vision model and get action.

        Args:
            image_base64: PNG screenshot as base64
            task: What the agent is trying to accomplish
            context: Additional context from previous steps

        Returns:
            (Action, raw_response, inference_time_ms)
        """
        prompt = SYSTEM_PROMPT.format(task=task)
        if context:
            prompt += f"\n\nContext from previous steps:\n{context}"

        payload = {
            "model": self.model,
            "prompt": prompt,
            "images": [image_base64],
            "stream": False,
            "options": {
                "temperature": 0.1,  # Low temperature for deterministic actions
            }
        }

        url = f"{self.base_url}/api/generate"
        data = json.dumps(payload).encode("utf-8")

        req = urllib.request.Request(
            url,
            data=data,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        start_time = time.perf_counter()

        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                result = json.loads(resp.read().decode("utf-8"))
        except urllib.error.URLError as e:
            return Action(
                action_type=ActionType.FAILED,
                reason=f"Ollama connection failed: {e}",
            ), str(e), 0.0

        inference_time_ms = (time.perf_counter() - start_time) * 1000
        raw_response = result.get("response", "")

        # Parse action from response
        action = self._parse_action(raw_response)

        return action, raw_response, inference_time_ms

    def _parse_action(self, response: str) -> Action:
        """Parse JSON action from model response."""
        # Try to extract JSON from response
        try:
            # Find JSON in response (model might include extra text)
            json_match = re.search(r'\{[^{}]*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
            else:
                data = json.loads(response)

            # Map to ActionType enum
            action_type_str = data.get("action_type", "failed").lower()
            try:
                action_type = ActionType(action_type_str)
            except ValueError:
                action_type = ActionType.FAILED

            return Action(
                action_type=action_type,
                x=data.get("x"),
                y=data.get("y"),
                text=data.get("text"),
                key=data.get("key"),
                keys=data.get("keys"),
                scroll_amount=data.get("scroll_amount"),
                duration=data.get("duration"),
                reason=data.get("reason", ""),
            )

        except (json.JSONDecodeError, KeyError) as e:
            return Action(
                action_type=ActionType.FAILED,
                reason=f"Could not parse model response: {e}. Response: {response[:200]}",
            )

    def check_health(self) -> bool:
        """Check if Ollama is running and model is available."""
        try:
            url = f"{self.base_url}/api/tags"
            req = urllib.request.Request(url, method="GET")
            with urllib.request.urlopen(req, timeout=5) as resp:
                result = json.loads(resp.read().decode("utf-8"))
                models = [m.get("name", "") for m in result.get("models", [])]
                return self.model in models or any(self.model.split(":")[0] in m for m in models)
        except Exception:
            return False
