# -*- coding: utf-8 -*-
# filename: prompt_enhancer.py
# @Time    : 2025-01-19
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""Prompt enhancer using ZhipuAI GLM-4.6V for intelligent prompt enhancement"""
import logging
from datetime import datetime
from typing import Any, Optional

from pydantic import BaseModel, ConfigDict

from tf_tools_mcp.tools.img_gen_tools.enhancers.base import BasePromptEnhancer
from tf_tools_mcp.tools.img_gen_tools.enhancers.templates import (
    get_enhancement_template,
)
from tf_tools_mcp.tools.img_gen_tools.types import PromptEnhancementResult

logger = logging.getLogger(__name__)


class PromptEnhancer(BaseModel, BasePromptEnhancer):
    """
    Prompt enhancer using ZhipuAI GLM-4.6V

    Uses GLM-4.6V to perform intelligent prompt enhancement based on specified modes.
    Supports both text-to-image and image-to-image scenarios.
    """

    glm_client: Any  # zhipuai.ZhipuAI client
    model_name: str = "glm-4.6v"
    timeout: int = 30
    max_retries: int = 3

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

    def enhance(
        self,
        prompt: str,
        mode: str,
        reference_image: Optional[str] = None,
    ) -> PromptEnhancementResult:
        """
        Enhance prompt

        Args:
            prompt: Original prompt
            mode: Enhancement mode (TextToImageMode or ImageToImageMode)
            reference_image: Reference image path (optional, for image-to-image)

        Returns:
            PromptEnhancementResult

        Raises:
            ValueError: If mode is unsupported and not custom
        """
        # custom mode returns directly without enhancement
        if mode == "custom":
            return PromptEnhancementResult(
                original_prompt=prompt,
                enhanced_prompt=prompt,
                mode=mode,
                enhancement_model=self.model_name,
                timestamp=datetime.now().isoformat(),
                token_usage={"total_tokens": 0, "input_tokens": 0, "output_tokens": 0},
                reference_image_used=False,
            )

        # Get enhancement template
        template = get_enhancement_template(mode)
        if not template:
            raise ValueError(f"Unsupported mode: {mode}")

        # Build system prompt
        system_prompt = template.format(user_prompt=prompt)

        # Try enhancement with retries
        for attempt in range(self.max_retries):
            try:
                if reference_image:
                    enhanced_prompt, token_usage = self._enhance_with_image(
                        reference_image, system_prompt, prompt
                    )
                else:
                    enhanced_prompt, token_usage = self._enhance_text_only(
                        system_prompt
                    )

                return PromptEnhancementResult(
                    original_prompt=prompt,
                    enhanced_prompt=enhanced_prompt,
                    mode=mode,
                    enhancement_model=self.model_name,
                    timestamp=datetime.now().isoformat(),
                    token_usage=token_usage,
                    reference_image_used=reference_image is not None,
                )

            except Exception as e:
                logger.warning(f"Prompt enhancement attempt {attempt + 1} failed: {e}")
                if attempt >= self.max_retries - 1:
                    # Last retry failed, fallback to original prompt
                    logger.error(
                        f"Prompt enhancement failed after {self.max_retries} attempts, using original prompt"
                    )
                    return PromptEnhancementResult(
                        original_prompt=prompt,
                        enhanced_prompt=prompt,
                        mode=mode,
                        enhancement_model=self.model_name,
                        timestamp=datetime.now().isoformat(),
                        token_usage={},
                        reference_image_used=False,
                        enhancement_failed=True,
                    )

    def _enhance_text_only(self, system_prompt: str) -> tuple[str, dict]:
        """Text-only enhancement

        Args:
            system_prompt: System prompt with enhancement template

        Returns:
            Tuple of (enhanced_prompt, token_usage)
        """
        response = self.glm_client.chat.completions.create(
            model=self.model_name,
            messages=[{"role": "user", "content": system_prompt}],
            max_tokens=2048,
            timeout=self.timeout,
        )

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

        return enhanced_prompt, token_usage

    def _enhance_with_image(
        self, reference_image: str, system_prompt: str, original_prompt: str
    ) -> tuple[str, dict]:
        """Enhancement with reference image (for image-to-image)

        Args:
            reference_image: Reference image path
            system_prompt: System prompt with enhancement template
            original_prompt: Original user prompt

        Returns:
            Tuple of (enhanced_prompt, token_usage)
        """
        # Prepare image content
        if reference_image.startswith(("http://", "https://")):
            image_url = reference_image
        else:
            # Convert local file to base64
            import base64

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

        response = self.glm_client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": image_url}},
                        {"type": "text", "text": system_prompt},
                    ],
                }
            ],
            max_tokens=2048,
            timeout=self.timeout,
        )

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

        return enhanced_prompt, token_usage
