# -*- coding: utf-8 -*-
# filename: image_to_image.py
# @Time    : 2025-01-19
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""Image-to-image generation tool using OpenRouter API"""
import base64
import io
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Any, Literal

import requests
from PIL import Image
from pydantic import ConfigDict, PrivateAttr
import zhipuai

from tf_tools_mcp.tools.img_gen_tools.constants import (
    MODEL_PARAMETER_CONFIG,
    SUPPORTED_GENERATION_MODELS,
)
from tf_tools_mcp.tools.img_gen_tools.enhancers.prompt_enhancer import PromptEnhancer
from tf_tools_mcp.tools.img_gen_tools.types import (
    ImageGenerationResult,
    ImageToImageMode,
)
from tf_tools_mcp.tools.img_gen_tools.utils import image_to_base64

logger = logging.getLogger(__name__)


class ImageToImageTool:
    """
    Image-to-image generation tool

    Uses OpenRouter API to access multiple image generation models.
    """

    name: str = "image_to_image"
    description: str = (
        "图生图工具，基于参考图片生成新图片，支持多种模型：Gemini 3 Pro、Seedream 4.5、GPT-5 Image、FLUX.2 等。"
    )
    version: str = "0.1.0"
    tags: set[str] = {"image", "generation", "vision", "image_to_image"}

    # Configuration parameters
    model: str  # OpenRouter model ID
    storage_dir: str | None  # Directory to save generated images (None = don't save)
    api_key: str | None = None  # OpenRouter API key (overrides env var)
    generation_timeout: int = 120  # Timeout in seconds
    enable_retries: bool = True
    max_retries: int = 3

    # Prompt enhancement configuration
    enable_prompt_enhancement: bool = True
    prompt_enhancer: PromptEnhancer | None = None

    _api_key: str = PrivateAttr(default="")

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

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

        Args:
            model: OpenRouter model ID (must be in SUPPORTED_GENERATION_MODELS)
            storage_dir: Directory to save generated images (None = don't save locally)
            api_key: OpenRouter API key (optional, defaults to OPENROUTER_API_KEY env var)
            generation_timeout: Request timeout in seconds
            enable_retries: Enable automatic retry on network errors
            max_retries: Maximum number of retries
        """
        # Set attributes from data
        for key, value in data.items():
            setattr(self, key, value)

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

        # Get and validate API key
        api_key = self.api_key or os.environ.get("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError(
                """❌ 缺少 OpenRouter API Key！

image_to_image 工具需要 OpenRouter API Key 来生成图片。

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

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

获取 API Key: https://openrouter.ai/keys"""
            )

        self._api_key = api_key

        # Ensure storage directory exists (if specified)
        if self.storage_dir:
            Path(self.storage_dir).mkdir(parents=True, exist_ok=True)

        # Initialize enhancer if enabled
        if self.enable_prompt_enhancement and self.prompt_enhancer is None:
            try:
                zhipu_key = os.environ.get("ZHIPUAI_API_KEY")
                if zhipu_key:
                    glm_client = zhipuai.ZhipuAI(api_key=zhipu_key)
                    self.prompt_enhancer = PromptEnhancer(glm_client=glm_client)
                else:
                    logger.warning(
                        "ZHIPUAI_API_KEY not found, disabling prompt enhancement"
                    )
                    self.enable_prompt_enhancement = False
            except Exception as e:
                logger.warning(
                    f"Failed to initialize prompt enhancer: {e}. Disabling prompt enhancement."
                )
                self.enable_prompt_enhancement = False
                self.prompt_enhancer = None

    def _map_params_to_model(
        self,
        quality: Literal["low", "medium", "high"],
        layout: Literal["square", "landscape", "portrait"],
    ) -> dict:
        """
        Map quality and layout to model-specific parameters

        Args:
            quality: Image quality level
            layout: Image layout

        Returns:
            Model-specific API parameters
        """
        # Gemini 3 Pro Image Preview
        if self.model.startswith("google/"):
            image_size_map = {"low": "1K", "medium": "2K", "high": "4K"}
            aspect_ratio_map = {
                "square": "1:1",
                "landscape": "16:9",
                "portrait": "9:16",
            }
            return {
                "image_config": {
                    "image_size": image_size_map[quality],
                    "aspect_ratio": aspect_ratio_map[layout],
                }
            }

        # Seedream 4.5
        elif self.model.startswith("bytedance-seed/"):
            aspect_ratio_map = {
                "square": "1:1",
                "landscape": "16:9",
                "portrait": "9:16",
            }
            return {"image_config": {"aspect_ratio": aspect_ratio_map[layout]}}

        # OpenAI GPT-5 Image series
        elif self.model.startswith("openai/"):
            quality_map = {"low": "low", "medium": "medium", "high": "high"}

            if "mini" in self.model.lower():
                size_map = {
                    "square": "1024x1024",
                    "landscape": "1536x1024",
                    "portrait": "1024x1536",
                }
            else:
                size_map = {
                    "square": "1024x1024",
                    "landscape": "1792x1024",
                    "portrait": "1024x1792",
                }

            return {"quality": quality_map[quality], "size": size_map[layout]}

        # FLUX.2 Max/Pro
        elif self.model.startswith("black-forest-labs/"):
            aspect_ratio_map = {
                "square": "1:1",
                "landscape": "16:9",
                "portrait": "9:16",
            }
            return {"image_config": {"aspect_ratio": aspect_ratio_map[layout]}}

        # Default fallback
        return {}

    def _prepare_reference_images(self, reference_images: list[str]) -> list[str]:
        """
        Prepare reference images for image-to-image generation

        Args:
            reference_images: List of local paths or URLs to reference images (1-8 images)

        Returns:
            List of base64 data URLs or URL strings
        """
        return [
            (
                ref_image
                if ref_image.startswith(("http://", "https://"))
                else image_to_base64(ref_image, format="PNG")
            )
            for ref_image in reference_images
        ]

    def _generate_image(
        self,
        prompt: str,
        reference_images: list[str],
        quality: Literal["low", "medium", "high"],
        layout: Literal["square", "landscape", "portrait"],
        save_to_local: bool = False,
        save_dir: str = "/tmp/generate_img/",
    ) -> ImageGenerationResult:
        """
        Generate image using OpenRouter API

        Args:
            prompt: Generation prompt
            reference_images: Reference images for image-to-image (1-8 images)
            quality: Image quality
            layout: Image layout
            save_to_local: Whether to save image to local file
            save_dir: Directory to save image (used when save_to_local=True)

        Returns:
            ImageGenerationResult
        """
        # Map parameters to model-specific format
        model_params = self._map_params_to_model(quality, layout)

        # Prepare reference images
        ref_image_urls = self._prepare_reference_images(reference_images)

        # Build content list with all reference images
        content = []
        for ref_url in ref_image_urls:
            content.append({"type": "image_url", "image_url": {"url": ref_url}})

        # Add text prompt
        if len(ref_image_urls) == 1:
            text_prompt = f"Based on this image, {prompt}"
        else:
            text_prompt = f"Based on these {len(ref_image_urls)} images, {prompt}"
        content.append({"type": "text", "text": text_prompt})

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

        # Call OpenRouter API using raw HTTP request
        url = "https://openrouter.ai/api/v1/chat/completions"

        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com/turingfocus/tfrobotv2",
            "X-Title": "TFRobotV2",
        }

        payload = {
            "model": self.model,
            "messages": messages,
            "modalities": ["image"],
            **model_params,
        }

        try:
            logger.info(f"Calling OpenRouter API with model: {self.model}")
            response = requests.post(
                url, headers=headers, json=payload, timeout=self.generation_timeout
            )
            response.raise_for_status()
            response_data = response.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"Image generation API call failed: {e}")
            if hasattr(e, "response") and e.response is not None:
                logger.error(f"Response status: {e.response.status_code}")
                logger.error(f"Response body: {e.response.text}")
            raise Exception(f"Image generation API call failed: {e}")

        # Check for API error response
        if "error" in response_data:
            error_info = response_data["error"]
            error_msg = error_info.get("message", "Unknown error")
            error_code = error_info.get("code", "N/A")
            provider = error_info.get("metadata", {}).get(
                "provider_name", "Unknown provider"
            )
            raise Exception(
                f"Provider error (code: {error_code}, provider: {provider}): {error_msg}"
            )

        # Extract image from response
        try:
            images = response_data["choices"][0]["message"]["images"]
            if not images or len(images) == 0:
                raise Exception("No image returned from model")

            # Handle different response formats
            if "url" in images[0]:
                image_data = images[0]["url"]
            elif "image_url" in images[0] and "url" in images[0]["image_url"]:
                image_data = images[0]["image_url"]["url"]
            else:
                raise Exception(f"Unknown image response format: {images[0].keys()}")

            if not image_data:
                raise Exception("Image URL is empty")

        except (KeyError, IndexError, TypeError) as e:
            logger.error(f"Failed to extract image from response: {e}")
            logger.error(f"Response data: {response_data}")
            raise Exception(f"Failed to extract image from response: {e}")

        # Ensure image_data is in data URL format
        if not image_data.startswith("data:"):
            image_data = f"data:image/png;base64,{image_data}"

        # Extract image format from data URL (e.g., "image/png" -> "PNG")
        try:
            format_from_data = (
                image_data.split(":")[1].split(";")[0].split("/")[1].upper()
            )
        except (IndexError, AttributeError):
            format_from_data = "PNG"  # Default fallback

        # Get image dimensions and file size from base64 data
        try:
            # Extract base64 string from data URL
            base64_str = (
                image_data.split(",", 1)[1]
                if image_data.startswith("data:")
                else image_data
            )
            img_bytes = base64.b64decode(base64_str)
            img = Image.open(io.BytesIO(img_bytes))
            width, height = img.size
            file_size = len(img_bytes)
        except Exception as e:
            logger.warning(f"Failed to get image metadata: {e}")
            width, height, file_size = 0, 0, 0

        # Save to local if requested (runtime parameter takes precedence)
        saved_file_path = None

        # Priority: save_to_local runtime parameter > storage_dir from initialization
        if save_to_local:
            try:
                # Convert to absolute path and expand user (~)
                save_path = Path(save_dir).expanduser().resolve()
                save_path.mkdir(parents=True, exist_ok=True)

                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{self.model.replace('/', '_')}_img2img_{timestamp}.{format_from_data.lower()}"
                file_path = save_path / filename

                # Extract and decode base64 data
                base64_str = (
                    image_data.split(",", 1)[1]
                    if image_data.startswith("data:")
                    else image_data
                )
                img_bytes = base64.b64decode(base64_str)

                with open(file_path, "wb") as f:
                    f.write(img_bytes)

                saved_file_path = str(file_path)
                logger.info(f"Image saved to local file: {saved_file_path}")

                # Get metadata from saved file
                try:
                    img = Image.open(file_path)
                    width, height = img.size
                    file_size = file_path.stat().st_size
                except Exception as e:
                    logger.warning(f"Failed to get image metadata from saved file: {e}")

            except Exception as e:
                logger.warning(f"Failed to save image locally: {e}")
        elif self.storage_dir:
            # Backward compatibility: save to storage_dir if configured
            try:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{self.model.replace('/', '_')}_img2img_{timestamp}.{format_from_data.lower()}"
                file_path = Path(self.storage_dir) / filename
                file_path.parent.mkdir(parents=True, exist_ok=True)

                # Extract and decode base64 data
                base64_str = (
                    image_data.split(",", 1)[1]
                    if image_data.startswith("data:")
                    else image_data
                )
                img_bytes = base64.b64decode(base64_str)

                with open(file_path, "wb") as f:
                    f.write(img_bytes)

                logger.info(f"Image saved to local file: {file_path}")
            except Exception as e:
                logger.warning(f"Failed to save image locally: {e}")

        # Build result: base64_data or file_path based on save_to_local
        if save_to_local and saved_file_path:
            return ImageGenerationResult(
                base64_data=None,
                file_path=saved_file_path,
                width=width,
                height=height,
                file_size=file_size,
                format=format_from_data,
                generated_at=datetime.now().isoformat(),
                model=self.model,
                quality=quality,
                layout=layout,
            )
        else:
            return ImageGenerationResult(
                base64_data=image_data,
                file_path=None,
                width=width,
                height=height,
                file_size=file_size,
                format=format_from_data,
                generated_at=datetime.now().isoformat(),
                model=self.model,
                quality=quality,
                layout=layout,
            )

    def run(
        self,
        prompt: str,
        reference_images: list[str],
        mode: ImageToImageMode = "custom",
        quality: Literal["low", "medium", "high"] = "medium",
        layout: Literal["square", "landscape", "portrait"] = "square",
        save_to_local: bool = False,
        save_dir: str = "/tmp/generate_img/",
        **kwargs: Any,
    ) -> ImageGenerationResult:
        """
        Run the image-to-image generation tool

        Args:
            prompt: Generation prompt
            reference_images: Reference images for image-to-image (1-8 images)
            mode: Enhancement mode
            quality: Image quality (low, medium, high)
            layout: Image layout (square, landscape, portrait)
            save_to_local: Whether to save image to local file (default: False)
            save_dir: Directory to save image (used when save_to_local=True, default: /tmp/generate_img/)

        Returns:
            ImageGenerationResult

        Raises:
            ValueError: If enhancement mode is selected but ZHIPUAI_API_KEY is not configured
            ValueError: If reference_images list is empty or contains more than 8 images
        """
        # Validate reference images
        if not reference_images or len(reference_images) == 0:
            raise ValueError("reference_images 不能为空，至少需要 1 张图片")
        if len(reference_images) > 8:
            raise ValueError(
                f"reference_images 最多支持 8 张图片，当前传入了 {len(reference_images)} 张"
            )

        # Validate enhancement mode configuration
        if mode != "custom" and not self.prompt_enhancer:
            raise ValueError(
                f"❌ 无法使用 prompt 增强模式: '{mode}'\n\n"
                f"image_to_image 工具的 prompt 增强功能需要智谱 AI API Key。\n\n"
                f"请设置环境变量：\n"
                f"  export ZHIPUAI_API_KEY='your_api_key_here'\n\n"
                f"或在 MCP Inspector 配置中添加环境变量：\n"
                f'  "env": {{"ZHIPUAI_API_KEY": "your_api_key_here"}}\n\n'
                f"如不需要 prompt 增强，请使用 mode='custom'\n"
                f"获取 API Key: https://open.bigmodel.cn/usercenter/apikeys"
            )

        # Prompt enhancement (with reference images - use first image for enhancement)
        original_prompt_value = prompt  # Always set to user's original input
        enhanced_prompt_value = None  # Only set when enhancement is used

        if self.prompt_enhancer and mode != "custom":
            # Use first reference image for prompt enhancement
            enhancement_result = self.prompt_enhancer.enhance(
                prompt=prompt, mode=mode, reference_image=reference_images[0]
            )
            enhanced_prompt_value = enhancement_result.enhanced_prompt

        # Generate image
        # Use enhanced prompt if available, otherwise use original prompt
        prompt_to_use = enhanced_prompt_value if enhanced_prompt_value else prompt

        result = self._generate_image(
            prompt=prompt_to_use,
            reference_images=reference_images,
            quality=quality,
            layout=layout,
            save_to_local=save_to_local,
            save_dir=save_dir,
        )

        # Set enhancement tracking
        result.original_prompt = original_prompt_value
        result.enhanced_prompt = enhanced_prompt_value

        return result
