"""
Nano Banana图像生成与编辑工具组件实现

提供基于远程API的图像生成和编辑功能
"""

import json
import requests
from pydantic import BaseModel, Field

from src.components.tools.base import component_tool
from src.utils.logger import get_logger


class NanoBananaInput(BaseModel):
    """Nano Banana图像生成/编辑输入参数 schema"""
    prompt: str = Field(
        description="生图/图片编辑提示词，提示词要表现出你的意图是得到一张图片"
    )
    images: str = Field(
        default="",
        description="要编辑的图片url。多张图片使用空格进行分隔。如果是生图任务则不需要提供"
    )


@component_tool(
    name="nano_banana_generate_image",
    description="""Nano Banana是谷歌最新的图像生成与编辑模型。
    核心优点：非常好的角色一致性
    使用场景：
    1. 提示词描述生成图片。提示词尽量详细且具体
    2. 图片编辑。在提示中描述清楚需要对图片如何编辑，比如修改图片背景、替换图片中某个元素、调整图片中人物的动作等等
    3. 图片融合。将多张图片的元素融合成在一张图片中
    4. 图片元素提取。提示词中清晰描述需要提取的元素

    Args:
        prompt: 生图/图片编辑提示词，提示词要表现出你的意图是得到一张图片
        images: 要编辑的图片url。多张图片使用空格进行分隔。如果是生图任务则不需要提供

    Returns:
        生成或编辑后的图片URL
    """,
    category="image_processing",
    schema=NanoBananaInput
)
def nano_banana_generate_image(prompt: str, images: str = "") -> str:
    """
    Nano Banana图像生成与编辑工具

    Args:
        prompt: 生图/图片编辑提示词
        images: 要编辑的图片url(可选)

    Returns:
        生成或编辑后的图片URL或错误信息
    """
    logger = get_logger("NanoBanana")

    try:
        # API接口地址
        api_url = "http://47.84.182.56:8001/tools/call/nano_banana_generate_image"

        # 构建请求参数
        payload = {
            "prompt": prompt
        }

        # 如果提供了图片URL，添加到请求参数中
        if images:
            payload["images"] = images
            logger.info(f"Requesting image editing with prompt: {prompt}, images: {images}...")
        else:
            logger.info(f"Requesting image generation with prompt: {prompt}")

        # 发起POST请求
        response = requests.post(api_url, json=payload, timeout=120)
        response.raise_for_status()

        # 解析响应
        result = response.json()

        logger.info(f"NanoBanana Raw response: {result}")

        # 检查响应格式
        if isinstance(result, dict):
            # 标准响应格式: {"success": True/False, "result": "...", ...}
            if "success" in result:
                success = result.get("success", False)
                result_data = result.get("result", "")

                # 处理result_data可能是JSON字符串的情况
                if isinstance(result_data, str):
                    try:
                        # 尝试解析JSON字符串
                        parsed_data = json.loads(result_data)
                        # 如果是数组，取第一个元素
                        if isinstance(parsed_data, list) and len(parsed_data) > 0:
                            result_data = parsed_data[0]
                        elif isinstance(parsed_data, list) and len(parsed_data) == 0:
                            error_msg = "Nano Banana返回空结果"
                            logger.error(error_msg)
                            return error_msg
                    except json.JSONDecodeError:
                        # 如果不是JSON，保持原样
                        pass
                # 处理result_data可能是数组的情况
                elif isinstance(result_data, list) and len(result_data) > 0:
                    result_data = result_data[0]  # 取第一个元素

                if success and result_data and not result_data.startswith("tools运行失败"):
                    # 成功且有有效结果
                    logger.info(f"Nano Banana completed successfully: {result_data}")
                    return result_data
                else:
                    # 失败或错误
                    error_msg = f"Nano Banana执行失败: {result_data}"
                    logger.error(error_msg)
                    return error_msg

            # 直接返回image_url的格式
            elif "image_url" in result:
                image_url = result["image_url"]
                # 处理image_url可能是数组的情况
                if isinstance(image_url, list) and len(image_url) > 0:
                    image_url = image_url[0]
                logger.info(f"Nano Banana completed successfully: {image_url}")
                return image_url

            # 其他字典格式
            else:
                logger.warning(f"Unexpected response format: {result}")
                return str(result)

        elif isinstance(result, str):
            # 如果直接返回字符串，假设是图片URL或错误信息
            if result.startswith("http"):
                logger.info(f"Nano Banana completed: {result}")
                return result
            else:
                logger.warning(f"Received non-URL string response: {result}")
                return result

        else:
            logger.warning(f"Unexpected response type: {type(result)}, value: {result}")
            return str(result)

    except requests.exceptions.Timeout:
        error_msg = "Nano Banana请求超时，请稍后重试"
        logger.error(error_msg)
        return error_msg

    except requests.exceptions.RequestException as e:
        error_msg = f"Nano Banana请求失败: {str(e)}"
        logger.error(error_msg, exc_info=True)
        return error_msg

    except Exception as e:
        error_msg = f"Nano Banana处理异常: {str(e)}"
        logger.error(error_msg, exc_info=True)
        return error_msg
