"""
图像分割工具组件实现

提供基于远程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 SegmentImageInput(BaseModel):
    """图像分割输入参数 schema"""
    image: str = Field(description="要分割的图片URL")
    prompt: str = Field(
        description="分割对象描述，必须是简单的英文单词（1-3个单词），"
                    "应该是通俗易懂的描述，从小孩子的视角，不要使用具体的人名等；"
                    "要求描述能在图片中唯一确定一个对象"
    )


@component_tool(
    name="segment_image",
    description="""图片分割工具，用于:
    根据对象名称对图像的进行分割，抽取指定对象的切片
    图像对象识别与分离
    前景背景分离
    多目标实例分割
    精确边缘提取

    Args:
        image: 要分割的图片URL
        prompt: 分割对象描述，必须是简单的英文单词（1-3个单词），
               应该是通俗易懂的描述，从小孩子的视角，不要使用具体的人名等；
               要求描述能在图片中唯一确定一个对象

    Returns:
        分割后的图片URL
    """,
    category="image_processing",
    schema=SegmentImageInput
)
def segment_image(image: str, prompt: str) -> str:
    """
    图片分割工具，根据对象描述从图像中分割出指定对象

    Args:
        image: 要分割的图片URL
        prompt: 分割对象描述（必须是简单的英文单词，1-3个单词）

    Returns:
        分割后的图片URL或错误信息
    """
    logger = get_logger("SegmentImage")

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

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

        logger.info(f"Requesting image segmentation for: {image}... with prompt: {prompt}")

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

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

        logger.info(f"segmentation 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 = "图像分割返回空结果"
                            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"Image segmentation completed successfully: {result_data}")
                    return result_data
                else:
                    # 失败或错误
                    error_msg = f"图像分割失败: {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"Image segmentation 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"Image segmentation 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 = "图像分割请求超时，请稍后重试"
        logger.error(error_msg)
        return error_msg

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

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