"""Vision tool: analyze screenshots and images using a vision-capable LLM."""

from __future__ import annotations

import asyncio
import base64
import json
import mimetypes
from pathlib import Path
from typing import Any

from loguru import logger

from madpod_agent.agent.tools.base import Tool
from madpod_agent.agent.tools.browser import BrowserBackend
from madpod_agent.providers.base import LLMProvider

_SUPPORTED_IMAGE_TYPES = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
_MAX_IMAGE_SIZE = 20 * 1024 * 1024  # 20 MB


def _resolve_safe(path_str: str, workspace: Path, allowed_dir: Path | None) -> Path:
    p = Path(path_str).expanduser()
    if not p.is_absolute():
        p = workspace / p
    resolved = p.resolve()
    if allowed_dir:
        try:
            resolved.relative_to(allowed_dir.resolve())
        except ValueError:
            raise PermissionError(f"Path {path_str} is outside allowed directory")
    return resolved


def _image_to_data_url(path: Path) -> str:
    mime = mimetypes.guess_type(str(path))[0] or "image/png"
    data = path.read_bytes()
    b64 = base64.b64encode(data).decode("ascii")
    return f"data:{mime};base64,{b64}"


class ScreenshotAnalyzeTool(Tool):
    """Take a screenshot of the current browser page and analyze it with vision."""

    def __init__(
        self,
        backend: BrowserBackend,
        provider: LLMProvider | None,
        vision_model: str | None,
        workspace: Path,
    ):
        self._backend = backend
        self._provider = provider
        self._vision_model = vision_model
        self._workspace = workspace

    @property
    def name(self) -> str:
        return "screenshot_analyze"

    @property
    def description(self) -> str:
        return (
            "Take a screenshot of the current browser page and analyze it with a vision model. "
            "Use when text-based snapshot is insufficient — e.g. CAPTCHAs, charts, images, "
            "visual layouts, or pages that render content as images/canvas. "
            "Ask a specific question about what you see."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "question": {
                    "type": "string",
                    "description": "What to analyze in the screenshot (e.g. 'What CAPTCHA text is shown?', 'Describe the chart')",
                },
                "full_page": {
                    "type": "boolean",
                    "description": "Capture full page instead of viewport (default false)",
                },
            },
            "required": ["question"],
        }

    async def execute(self, question: str, full_page: bool = False, **kwargs: Any) -> str:
        if not self._provider or not self._vision_model:
            return json.dumps({"error": "No vision model configured. Set a vision-capable model."})

        # Take screenshot
        result_raw = await self._backend.take_screenshot(full_page=full_page)
        try:
            result = json.loads(result_raw)
            screenshot_path = Path(result.get("path", ""))
        except (json.JSONDecodeError, TypeError):
            return json.dumps({"error": "Failed to take screenshot"})

        if not screenshot_path.exists():
            return json.dumps({"error": f"Screenshot file not found: {screenshot_path}"})

        data_url = _image_to_data_url(screenshot_path)

        try:
            response = await self._provider.chat(
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": data_url}},
                            {"type": "text", "text": question},
                        ],
                    }
                ],
                model=self._vision_model,
                max_tokens=1500,
                temperature=0.1,
            )
            analysis = (response.content or "").strip()
            return json.dumps({
                "ok": True,
                "screenshot": str(screenshot_path),
                "question": question,
                "analysis": analysis,
            })
        except Exception as exc:
            logger.error("Vision analysis failed: {}", exc)
            return json.dumps({"error": f"Vision analysis failed: {exc}", "screenshot": str(screenshot_path)})


class ImageAnalyzeTool(Tool):
    """Analyze a local image file with a vision model."""

    def __init__(
        self,
        provider: LLMProvider | None,
        vision_model: str | None,
        workspace: Path,
        allowed_dir: Path | None = None,
    ):
        self._provider = provider
        self._vision_model = vision_model
        self._workspace = workspace
        self._allowed_dir = allowed_dir

    @property
    def name(self) -> str:
        return "image_analyze"

    @property
    def description(self) -> str:
        return (
            "Analyze a local image file using a vision model. "
            "Use for OCR (reading text from images), describing charts/diagrams, "
            "extracting data from screenshots, or analyzing any image content. "
            "Supports PNG, JPG, GIF, WebP, BMP."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path to the image file"},
                "question": {
                    "type": "string",
                    "description": "What to analyze (e.g. 'Extract all text', 'Describe this chart', 'What data is shown?')",
                },
            },
            "required": ["path", "question"],
        }

    async def execute(self, path: str, question: str, **kwargs: Any) -> str:
        if not self._provider or not self._vision_model:
            return json.dumps({"error": "No vision model configured."})

        try:
            file_path = _resolve_safe(path, self._workspace, self._allowed_dir)
        except PermissionError as exc:
            return json.dumps({"error": str(exc)})

        if not file_path.exists():
            return json.dumps({"error": f"File not found: {path}"})

        suffix = file_path.suffix.lower()
        if suffix not in _SUPPORTED_IMAGE_TYPES:
            return json.dumps({"error": f"Unsupported image type: {suffix}"})

        if file_path.stat().st_size > _MAX_IMAGE_SIZE:
            return json.dumps({"error": f"Image too large (max {_MAX_IMAGE_SIZE // 1024 // 1024}MB)"})

        data_url = _image_to_data_url(file_path)

        try:
            response = await self._provider.chat(
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": data_url}},
                            {"type": "text", "text": question},
                        ],
                    }
                ],
                model=self._vision_model,
                max_tokens=2000,
                temperature=0.1,
            )
            analysis = (response.content or "").strip()
            return json.dumps({
                "ok": True,
                "path": str(file_path),
                "question": question,
                "analysis": analysis,
            })
        except Exception as exc:
            logger.error("Image analysis failed: {}", exc)
            return json.dumps({"error": f"Analysis failed: {exc}"})
