Seedream 4.0 MCP

"""
Seedream 4.0 API 客户端模块

本模块提供 Seedream 4.0 API 的客户端封装，支持文生图、图生图、
多图融合和组图生成等功能。
"""

import asyncio
import base64
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import httpx

from .config import SeedreamConfig, get_global_config
from .utils.errors import SeedreamAPIError, SeedreamNetworkError, SeedreamTimeoutError
from .utils.logging import get_logger, log_function_call
from .utils.validation import (
    validate_image_list,
    validate_image_url,
    validate_max_images,
    validate_prompt,
    validate_response_format,
    validate_size,
    validate_watermark,
)


class SeedreamClient:
    """
    Seedream 4.0 API 客户端类
    
    提供异步 HTTP 客户端封装，支持多种图像生成功能：
    - 文生图（text_to_image）
    - 图生图（image_to_image）
    - 多图融合（multi_image_fusion）
    - 组图生成（sequential_generation）
    """

    def __init__(self, config: Optional[SeedreamConfig] = None):
        """
        初始化 Seedream API 客户端
        
        Args:
            config: 配置对象，若为 None 则使用全局默认配置
        """
        self.config = config or get_global_config()
        self.logger = get_logger(__name__)
        self._client = None

    async def __aenter__(self):
        """
        异步上下文管理器入口
        
        Returns:
            SeedreamClient: 当前客户端实例
        """
        await self._ensure_client()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """
        异步上下文管理器出口
        
        Args:
            exc_type: 异常类型
            exc_val: 异常值
            exc_tb: 异常追踪信息
        """
        await self.close()

    @log_function_call
    async def text_to_image(
        self,
        prompt: str,
        size: str = "1K",
        watermark: bool = True,
        response_format: str = "url"
    ) -> Dict[str, Any]:
        """
        文生图功能
        
        根据文本提示词生成图像。
        
        Args:
            prompt: 文本提示词，描述要生成的图像内容
            size: 图像尺寸，可选值为 "1K"、"2K"、"4K"，默认为 "1K"
            watermark: 是否添加水印，默认为 True
            response_format: 响应格式，可选值为 "url" 或 "b64_json"，默认为 "url"
        
        Returns:
            包含生成结果的字典，包括图像数据、使用信息、任务 ID 等
        
        Raises:
            SeedreamAPIError: API 调用失败
            SeedreamValidationError: 参数验证失败
        """
        # 参数验证
        prompt = validate_prompt(prompt)
        size = validate_size(size)
        watermark = validate_watermark(watermark)
        response_format = validate_response_format(response_format)

        self.logger.info(f"开始文生图任务: prompt='{prompt[:50]}...', size={size}")

        try:
            # 构建请求参数
            request_data = {
                "model": self.config.model_id,
                "prompt": prompt,
                "size": size,
                "watermark": watermark,
                "response_format": response_format
            }

            # 调用 API
            response = await self._call_api("text_to_image", request_data)

            self.logger.info("文生图任务完成")
            return response

        except Exception as e:
            self.logger.error(f"文生图任务失败: {str(e)}")
            handled_error = self._handle_api_error(e)
            raise handled_error

    @log_function_call
    async def image_to_image(
        self,
        prompt: str,
        image: str,
        size: str = "1K",
        watermark: bool = True,
        response_format: str = "url"
    ) -> Dict[str, Any]:
        """
        图生图功能
        
        基于输入图像和文本提示词生成新图像。
        
        Args:
            prompt: 文本提示词，描述要生成的图像内容
            image: 输入图像的 URL 或本地文件路径
            size: 图像尺寸，可选值为 "1K"、"2K"、"4K"，默认为 "1K"
            watermark: 是否添加水印，默认为 True
            response_format: 响应格式，可选值为 "url" 或 "b64_json"，默认为 "url"
        
        Returns:
            包含生成结果的字典，包括图像数据、使用信息、任务 ID 等
        
        Raises:
            SeedreamAPIError: API 调用失败或图像处理失败
            SeedreamValidationError: 参数验证失败
        """
        # 参数验证
        prompt = validate_prompt(prompt)
        image = validate_image_url(image)
        size = validate_size(size)
        watermark = validate_watermark(watermark)
        response_format = validate_response_format(response_format)

        self.logger.info(f"开始图生图任务: prompt='{prompt[:50]}...', size={size}")

        try:
            # 处理图像输入
            image_data = await self._prepare_image_input(image)

            # 构建请求参数
            request_data = {
                "model": self.config.model_id,
                "prompt": prompt,
                "image": image_data,
                "size": size,
                "watermark": watermark,
                "response_format": response_format
            }

            # 调用 API
            response = await self._call_api("image_to_image", request_data)

            self.logger.info("图生图任务完成")
            return response

        except Exception as e:
            self.logger.error(f"图生图任务失败: {str(e)}")
            handled_error = self._handle_api_error(e)
            raise handled_error

    @log_function_call
    async def multi_image_fusion(
        self,
        prompt: str,
        images: List[str],
        size: str = "1K",
        watermark: bool = True,
        response_format: str = "url"
    ) -> Dict[str, Any]:
        """
        多图融合功能
        
        将多张图像融合生成新图像。
        
        Args:
            prompt: 文本提示词，描述要生成的图像内容
            images: 输入图像的 URL 或本地文件路径列表，数量范围为 2-5 张
            size: 图像尺寸，可选值为 "1K"、"2K"、"4K"，默认为 "1K"
            watermark: 是否添加水印，默认为 True
            response_format: 响应格式，可选值为 "url" 或 "b64_json"，默认为 "url"
        
        Returns:
            包含生成结果的字典，包括图像数据、使用信息、任务 ID 等
        
        Raises:
            SeedreamAPIError: API 调用失败或图像处理失败
            SeedreamValidationError: 参数验证失败（如图像数量不符合要求）
        """
        # 参数验证
        prompt = validate_prompt(prompt)
        images = validate_image_list(images, min_count=2, max_count=5)
        size = validate_size(size)
        watermark = validate_watermark(watermark)
        response_format = validate_response_format(response_format)

        self.logger.info(f"开始多图融合任务: prompt='{prompt[:50]}...', images={len(images)}张, size={size}")

        try:
            # 处理图像输入
            image_data_list = []
            for image in images:
                image_data = await self._prepare_image_input(image)
                image_data_list.append(image_data)

            # 构建请求参数
            request_data = {
                "model": self.config.model_id,
                "prompt": prompt,
                "images": image_data_list,
                "size": size,
                "watermark": watermark,
                "response_format": response_format
            }

            # 调用 API
            response = await self._call_api("multi_image_fusion", request_data)

            self.logger.info("多图融合任务完成")
            return response

        except Exception as e:
            self.logger.error(f"多图融合任务失败: {str(e)}")
            handled_error = self._handle_api_error(e)
            raise handled_error

    @log_function_call
    async def sequential_generation(
        self,
        prompt: str,
        max_images: int = 4,
        size: str = "2K",
        watermark: bool = True,
        response_format: str = "url",
        image: Optional[Union[str, List[str]]] = None
    ) -> Dict[str, Any]:
        """
        组图生成功能（连续生成多张图像）
        
        支持三种输入类型：
        1. 文生组图：仅使用文本提示词
        2. 单图生组图：使用单张参考图像和文本提示词
        3. 多图生组图：使用多张参考图像和文本提示词
        
        Args:
            prompt: 文本提示词，描述要生成的图像内容
            max_images: 最大生成图像数量，范围为 1-15，默认为 4
            size: 图像尺寸，可选值为 "1K"、"2K"、"4K"，默认为 "2K"
            watermark: 是否添加水印，默认为 True
            response_format: 响应格式，可选值为 "url" 或 "b64_json"，默认为 "url"
            image: 可选的参考图像，支持单张图像 URL/路径或多张图像 URL/路径列表（最多 10 张）
        
        Returns:
            包含生成结果的字典，包括图像数据、使用信息、任务 ID 等
        
        Raises:
            SeedreamAPIError: API 调用失败或图像处理失败
            SeedreamValidationError: 参数验证失败
        """
        # 参数验证
        prompt = validate_prompt(prompt)
        max_images = validate_max_images(max_images)
        size = validate_size(size)
        watermark = validate_watermark(watermark)
        response_format = validate_response_format(response_format)

        # 处理图像输入
        processed_image = None
        if image is not None:
            if isinstance(image, str):
                # 单张图片
                processed_image = await self._prepare_image_input(image)
            elif isinstance(image, list):
                # 多张图片
                if len(image) > 10:
                    raise SeedreamAPIError("最多支持 10 张参考图片")
                processed_image = []
                for img in image:
                    processed_img = await self._prepare_image_input(img)
                    processed_image.append(processed_img)
            else:
                raise SeedreamAPIError("image 参数必须是字符串或字符串列表")

        self.logger.info(f"开始组图生成任务: prompt='{prompt[:50]}...', max_images={max_images}, size={size}")

        try:
            # 构建请求参数
            request_data = {
                "model": self.config.model_id,
                "prompt": prompt,
                "sequential_image_generation": "auto",
                "sequential_image_generation_options": {
                    "max_images": max_images
                },
                "size": size,
                "watermark": watermark,
                "response_format": response_format
            }

            # 添加图像参数（如果存在）
            if processed_image is not None:
                request_data["image"] = processed_image

            # 调用 API
            response = await self._call_api("sequential_generation", request_data)

            self.logger.info("组图生成任务完成")
            return response

        except Exception as e:
            self.logger.error(f"组图生成任务失败: {str(e)}")
            handled_error = self._handle_api_error(e)
            raise handled_error

    async def close(self):
        """
        关闭 HTTP 客户端连接
        
        释放客户端资源，关闭所有打开的连接。
        """
        if self._client:
            await self._client.aclose()
            self._client = None

    async def _ensure_client(self):
        """
        确保 HTTP 客户端已创建
        
        如果客户端未初始化，则创建新的 AsyncClient 实例。
        
        Raises:
            SeedreamAPIError: 客户端创建失败或配置无效
        """
        if self._client is None:
            try:
                headers = self._get_headers()
                if not headers:
                    raise SeedreamAPIError("无法生成请求头：配置可能无效")

                self._client = httpx.AsyncClient(
                    timeout=self.config.api_timeout,
                    headers=headers
                )

                # 验证客户端是否正确创建
                if self._client is None:
                    raise SeedreamAPIError("HTTP 客户端创建失败")

                self.logger.debug("HTTP 客户端创建成功")

            except Exception as e:
                self.logger.error(f"HTTP 客户端创建失败: {str(e)}")
                self._client = None
                raise SeedreamAPIError(f"HTTP 客户端初始化失败: {str(e)}")

    def _get_headers(self) -> Dict[str, str]:
        """
        获取 API 请求头
        
        构建包含认证信息的 HTTP 请求头。
        
        Returns:
            包含 Authorization 和 Content-Type 的请求头字典
        
        Raises:
            SeedreamAPIError: 配置对象为空或 API 密钥为空
        """
        if not self.config:
            raise SeedreamAPIError("配置对象为空")

        if not self.config.api_key:
            raise SeedreamAPIError("API 密钥为空，请检查环境变量 ARK_API_KEY")

        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }

        self.logger.debug(f"生成请求头: Authorization=Bearer {self.config.api_key[:10]}...")
        return headers

    async def _call_api(
        self,
        endpoint: str,
        request_data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        调用 Seedream API
        
        执行 HTTP POST 请求，支持自动重试机制。
        
        Args:
            endpoint: API 端点标识（用于日志记录）
            request_data: 请求体数据
        
        Returns:
            包含成功标志、数据、使用信息等的响应字典
        
        Raises:
            SeedreamAPIError: API 调用失败或响应解析失败
            SeedreamTimeoutError: 请求超时
            SeedreamNetworkError: 网络连接失败
        """
        await self._ensure_client()

        # 验证客户端是否正确创建
        if self._client is None:
            raise SeedreamAPIError("HTTP 客户端未正确初始化")

        # 构建 URL（Seedream 4.0 API 仅有一个端点）
        url = f"{self.config.base_url}/images/generations"

        for attempt in range(self.config.max_retries):
            try:
                self.logger.debug(f"API 调用尝试 {attempt + 1}/{self.config.max_retries}")
                self.logger.debug(f"请求 URL: {url}")
                self.logger.debug(f"请求数据: {request_data}")

                # 确保客户端的 post 方法存在且可调用
                if not hasattr(self._client, 'post') or not callable(self._client.post):
                    raise SeedreamAPIError("HTTP 客户端的 post 方法不可用")

                response = await self._client.post(
                    url,
                    json=request_data,
                    timeout=self.config.api_timeout
                )

                # 验证响应对象
                if response is None:
                    raise SeedreamAPIError("API 响应为空")

                self.logger.debug(f"收到响应: 状态码={response.status_code}")

                # 检查 HTTP 状态码
                if response.status_code == 200:
                    # 解析 JSON 响应
                    try:
                        result = response.json()
                        self.logger.debug(f"解析 JSON 成功: {result}")
                    except Exception as json_error:
                        raise SeedreamAPIError(f"JSON 解析失败: {str(json_error)}")

                    return {
                        "success": True,
                        "data": result.get("data", []),
                        "usage": result.get("usage", {}),
                        "task_id": result.get("task_id"),
                        "status": result.get("status")
                    }
                else:
                    error_msg = f"HTTP {response.status_code}: {response.text}"
                    raise SeedreamAPIError(error_msg)

            except httpx.TimeoutException:
                self.logger.warning(f"API 调用超时 (尝试 {attempt + 1}/{self.config.max_retries})")
                if attempt == self.config.max_retries - 1:
                    raise SeedreamTimeoutError("API 调用超时")

            except httpx.NetworkError as e:
                self.logger.warning(f"网络错误 (尝试 {attempt + 1}/{self.config.max_retries}): {str(e)}")
                if attempt == self.config.max_retries - 1:
                    raise SeedreamNetworkError(f"网络连接失败: {str(e)}")

            except Exception as e:
                self.logger.warning(f"API 调用失败 (尝试 {attempt + 1}/{self.config.max_retries}): {str(e)}")
                if attempt == self.config.max_retries - 1:
                    raise

            # 指数退避重试
            await asyncio.sleep(2 ** attempt)

        raise SeedreamAPIError("API 调用重试次数已用尽")

    async def _prepare_image_input(self, image: str) -> str:
        """
        准备图像输入数据
        
        将图像 URL 或本地文件路径转换为 API 所需格式。
        对于 URL 直接返回，对于本地文件读取并转换为 Base64 编码。
        
        Args:
            image: 图像 URL 或本地文件路径
        
        Returns:
            处理后的图像数据（URL 或 Base64 Data URI）
        
        Raises:
            SeedreamAPIError: 图像文件不存在或处理失败
        """
        try:
            # 如果是 URL，直接返回
            if image.startswith(("http://", "https://")):
                return image

            # 如果是文件路径，读取并转换为 Base64
            path = Path(image)
            if not path.exists():
                raise SeedreamAPIError(f"图像文件不存在: {image}")

            with open(path, "rb") as f:
                image_bytes = f.read()

            # 转换为 Base64
            image_b64 = base64.b64encode(image_bytes).decode('utf-8')

            # 获取 MIME 类型
            suffix = path.suffix.lower()
            mime_type_map = {
                '.jpg': 'image/jpeg',
                '.jpeg': 'image/jpeg',
                '.png': 'image/png',
                '.gif': 'image/gif',
                '.bmp': 'image/bmp',
                '.webp': 'image/webp'
            }
            mime_type = mime_type_map.get(suffix, 'image/jpeg')

            return f"data:{mime_type};base64,{image_b64}"

        except Exception as e:
            raise SeedreamAPIError(f"图像处理失败: {str(e)}")

    def _handle_api_error(self, error: Exception) -> Exception:
        """
        处理 API 错误
        
        将通用异常转换为特定的 Seedream 错误类型。
        
        Args:
            error: 原始异常对象
        
        Returns:
            处理后的 Seedream 特定异常对象
        """
        if isinstance(error, (SeedreamAPIError, SeedreamTimeoutError, SeedreamNetworkError)):
            return error

        error_str = str(error)

        # 超时错误
        if "timeout" in error_str.lower():
            return SeedreamTimeoutError(f"API 调用超时: {error_str}")

        # 网络错误
        if any(keyword in error_str.lower() for keyword in ["connection", "network", "dns"]):
            return SeedreamNetworkError(f"网络连接失败: {error_str}")

        # 其他 API 错误
        return SeedreamAPIError(f"API 调用失败: {error_str}")