# -*- coding: utf-8 -*-
# filename: understanding.py
# @Time    : 2025-01-16
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""Image understanding tool using ZhipuAI GLM-4.6V"""
import logging
import os
from datetime import datetime
from typing import Any, Optional, Literal

from pydantic import ConfigDict, PrivateAttr
import zhipuai

from tf_tools_mcp.tools.img_gen_tools.constants import (
    PROMPT_TEMPLATES,
    SUPPORTED_UNDERSTANDING_MODELS,
    UNDERSTANDING_MODEL_CONFIG,
)
from tf_tools_mcp.tools.img_gen_tools.types import (
    ImageUnderstandingParams,
    ImageUnderstandingResult,
)

logger = logging.getLogger(__name__)


class ImageUnderstandingTool:
    """
    Image understanding tool using ZhipuAI GLM-4.6V vision model

    Supports single and batch image analysis with multiple prompt templates.
    """

    name: str = "image_understanding"
    description: str = (
        "图片理解工具，支持单图或批量图片分析，基于智谱 GLM-4.6V 视觉模型。支持多种分析模板：描述、OCR、对象检测、图表分析。"
    )
    version: str = "0.1.0"
    tags: set[str] = {"image", "understanding", "vision", "ocr"}

    # Configuration parameters
    model: SUPPORTED_UNDERSTANDING_MODELS = "glm-4.6v-flashx"
    max_tokens: int = 4096
    timeout: int = 30

    _glm_client: Any = PrivateAttr(default=None)

    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=False)

    def __init__(self, **data: Any) -> None:
        """
        Initialize ImageUnderstandingTool

        Args:
            model: GLM model to use (glm-4.6v or glm-4.6v-flashx)
            max_tokens: Maximum output tokens
            timeout: Request timeout in seconds
        """
        # Set attributes from data
        for key, value in data.items():
            setattr(self, key, value)

        # Validate model
        if self.model not in UNDERSTANDING_MODEL_CONFIG:
            raise ValueError(
                f"Unsupported model: {self.model}. "
                f"Supported models: {list(UNDERSTANDING_MODEL_CONFIG.keys())}"
            )

        # Initialize ZhipuAI client
        api_key = os.environ.get("ZHIPUAI_API_KEY")
        if not api_key:
            raise ValueError(
                """❌ 缺少智谱 AI API Key！

image_understanding 工具需要智谱 AI API Key 来分析图片。

请设置环境变量：
  export ZHIPUAI_API_KEY='your_api_key_here'

或在 MCP Inspector 配置中添加环境变量：
  "env": {"ZHIPUAI_API_KEY": "your_api_key_here"}

获取 API Key: https://open.bigmodel.cn/usercenter/apikeys"""
            )

        self._glm_client = zhipuai.ZhipuAI(api_key=api_key)

    def _get_prompt(self, prompt_template: str, custom_prompt: Optional[str]) -> str:
        """
        Get the prompt based on template

        Args:
            prompt_template: Prompt template name
            custom_prompt: Custom prompt (overrides template)

        Returns:
            The final prompt string
        """
        if custom_prompt:
            return custom_prompt

        if prompt_template == "custom":
            raise ValueError(
                "custom_prompt is required when prompt_template is 'custom'"
            )

        if prompt_template not in PROMPT_TEMPLATES:
            raise ValueError(
                f"Unsupported prompt_template: {prompt_template}. "
                f"Supported: {list(PROMPT_TEMPLATES.keys())}"
            )

        return PROMPT_TEMPLATES[prompt_template]

    def _prepare_image_content(self, image: str) -> dict:
        """
        Prepare image for ZhipuAI API

        Args:
            image: Local file path or URL

        Returns:
            Image content dict for API
        """
        if image.startswith(("http://", "https://")):
            # Public URL - use directly
            return {"type": "image_url", "image_url": {"url": image}}
        else:
            # Local file - convert to base64
            import base64

            with open(image, "rb") as f:
                img_data = base64.b64encode(f.read()).decode("utf-8")
                return {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{img_data}"},
                }

    def _understand_single_image(self, image: str, prompt: str) -> tuple[str, dict]:
        """
        Understand a single image using GLM

        Args:
            image: Image path or URL
            prompt: Analysis prompt

        Returns:
            Tuple of (description, token_usage)
        """
        # Build messages
        image_content = self._prepare_image_content(image)

        messages = [
            {
                "role": "user",
                "content": [image_content, {"type": "text", "text": prompt}],
            }
        ]

        try:
            # Call ZhipuAI API
            response = self._glm_client.chat.completions.create(
                model=self.model,
                messages=messages,
                max_tokens=self.max_tokens,
                timeout=self.timeout,
            )

            description = response.choices[0].message.content
            token_usage = {
                "total_tokens": response.usage.total_tokens,
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
            }

            return description, token_usage

        except Exception as e:
            logger.error(f"Image understanding failed: {e}")
            raise Exception(f"Image understanding failed: {e}")

    def _understand_multiple_images(
        self, images: list[str], prompt: str
    ) -> tuple[str, dict]:
        """
        Understand multiple images using GLM

        Args:
            images: List of image paths or URLs
            prompt: Analysis prompt

        Returns:
            Tuple of (description, token_usage)
        """
        # Build messages with all images
        content = []
        for img in images:
            content.append(self._prepare_image_content(img))
        content.append({"type": "text", "text": prompt})

        messages = [{"role": "user", "content": content}]

        try:
            # Call ZhipuAI API
            response = self._glm_client.chat.completions.create(
                model=self.model,
                messages=messages,
                max_tokens=self.max_tokens,
                timeout=self.timeout,
            )

            description = response.choices[0].message.content
            token_usage = {
                "total_tokens": response.usage.total_tokens,
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
            }

            return description, token_usage

        except Exception as e:
            logger.error(f"Batch image understanding failed: {e}")
            raise Exception(f"Batch image understanding failed: {e}")

    def run(
        self,
        images: list[str],
        prompt_template: str = "describe",
        custom_prompt: str | None = None,
        **kwargs: Any,
    ) -> ImageUnderstandingResult:
        """
        Run the image understanding tool

        Args:
            images: List of image paths or URLs
            prompt_template: Analysis template (describe, ocr, detect_objects, analyze_chart, custom)
            custom_prompt: Custom prompt (required if prompt_template is 'custom')

        Returns:
            ImageUnderstandingResult
        """
        # Validate images
        if not images or len(images) == 0:
            raise ValueError("At least one image is required")

        # Get prompt
        prompt = self._get_prompt(prompt_template, custom_prompt)

        # Process images
        if len(images) == 1:
            description, token_usage = self._understand_single_image(images[0], prompt)
        else:
            description, token_usage = self._understand_multiple_images(images, prompt)

        # Build result
        return ImageUnderstandingResult(
            description=description,
            model=self.model,
            processed_at=datetime.now().isoformat(),
            image_count=len(images),
            prompt_template=prompt_template,
            used_prompt=prompt,
            total_tokens=token_usage["total_tokens"],
            input_tokens=token_usage["input_tokens"],
            output_tokens=token_usage["output_tokens"],
        )

    def format_output(self, result: ImageUnderstandingResult) -> str:
        """Format result for LLM consumption"""
        return (
            f"图片分析完成 (模型: {result.model}, 模板: {result.prompt_template}, "
            f"图片数量: {result.image_count}, "
            f"Tokens: {result.input_tokens}/{result.output_tokens})\n\n"
            f"使用的 Prompt: {result.used_prompt}\n\n"
            f"分析结果：\n{result.description}"
        )
