Seedream 4.0 MCP Server

"""
下载管理模块
实现异步图片下载、超时控制、重试机制
"""

import asyncio
import aiohttp
import logging
from typing import Optional, Dict, Any
from pathlib import Path
import time

logger = logging.getLogger(__name__)


class DownloadError(Exception):
    """下载错误异常"""
    pass


class DownloadManager:
    """异步下载管理器"""
    
    def __init__(
        self,
        timeout: int = 30,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        max_file_size: int = 50 * 1024 * 1024  # 50MB
    ):
        """
        初始化下载管理器
        
        Args:
            timeout: 下载超时时间（秒）
            max_retries: 最大重试次数
            retry_delay: 重试延迟时间（秒）
            max_file_size: 最大文件大小（字节）
        """
        self.timeout = timeout
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.max_file_size = max_file_size
        
    async def download_image(
        self,
        url: str,
        save_path: Path,
        headers: Optional[Dict[str, str]] = None
    ) -> Dict[str, Any]:
        """
        异步下载图片
        
        Args:
            url: 图片URL
            save_path: 保存路径
            headers: 请求头
            
        Returns:
            下载结果信息
            
        Raises:
            DownloadError: 下载失败时抛出
        """
        if headers is None:
            headers = {
                'User-Agent': 'Seedream-MCP/1.0',
                'Accept': 'image/*'
            }
            
        start_time = time.time()
        last_error = None
        
        for attempt in range(self.max_retries + 1):
            try:
                logger.info(f"开始下载图片 (尝试 {attempt + 1}/{self.max_retries + 1}): {url}")
                
                async with aiohttp.ClientSession(
                    timeout=aiohttp.ClientTimeout(total=self.timeout)
                ) as session:
                    async with session.get(url, headers=headers) as response:
                        # 检查响应状态
                        if response.status != 200:
                            raise DownloadError(f"HTTP错误: {response.status}")
                        
                        # 检查内容类型
                        content_type = response.headers.get('content-type', '')
                        if not content_type.startswith('image/'):
                            logger.warning(f"内容类型可能不是图片: {content_type}")
                        
                        # 检查文件大小
                        content_length = response.headers.get('content-length')
                        if content_length and int(content_length) > self.max_file_size:
                            raise DownloadError(f"文件过大: {content_length} 字节")
                        
                        # 确保目录存在
                        save_path.parent.mkdir(parents=True, exist_ok=True)
                        
                        # 下载并保存文件
                        total_size = 0
                        with open(save_path, 'wb') as f:
                            async for chunk in response.content.iter_chunked(8192):
                                total_size += len(chunk)
                                if total_size > self.max_file_size:
                                    raise DownloadError(f"文件过大: {total_size} 字节")
                                f.write(chunk)
                        
                        download_time = time.time() - start_time
                        
                        result = {
                            'success': True,
                            'file_path': str(save_path),
                            'file_size': total_size,
                            'download_time': download_time,
                            'content_type': content_type,
                            'attempts': attempt + 1
                        }
                        
                        logger.info(f"图片下载成功: {save_path} ({total_size} 字节, {download_time:.2f}秒)")
                        return result
                        
            except asyncio.TimeoutError as e:
                last_error = DownloadError(f"下载超时: {e}")
                logger.warning(f"下载超时 (尝试 {attempt + 1}): {url}")
                
            except aiohttp.ClientError as e:
                last_error = DownloadError(f"网络错误: {e}")
                logger.warning(f"网络错误 (尝试 {attempt + 1}): {e}")
                
            except OSError as e:
                last_error = DownloadError(f"文件系统错误: {e}")
                logger.warning(f"文件系统错误 (尝试 {attempt + 1}): {e}")
                
            except Exception as e:
                last_error = DownloadError(f"未知错误: {e}")
                logger.warning(f"下载失败 (尝试 {attempt + 1}): {e}")
            
            # 如果不是最后一次尝试，等待后重试
            if attempt < self.max_retries:
                await asyncio.sleep(self.retry_delay * (attempt + 1))
        
        # 所有重试都失败了
        logger.error(f"图片下载失败，已重试 {self.max_retries} 次: {url}")
        raise last_error or DownloadError("下载失败")
    
    async def download_multiple_images(
        self,
        urls_and_paths: list[tuple[str, Path]],
        headers: Optional[Dict[str, str]] = None,
        max_concurrent: int = 5
    ) -> list[Dict[str, Any]]:
        """
        并发下载多个图片
        
        Args:
            urls_and_paths: URL和保存路径的元组列表
            headers: 请求头
            max_concurrent: 最大并发数
            
        Returns:
            下载结果列表
        """
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def download_with_semaphore(url: str, path: Path) -> Dict[str, Any]:
            async with semaphore:
                try:
                    return await self.download_image(url, path, headers)
                except Exception as e:
                    logger.error(f"下载失败: {url} -> {e}")
                    return {
                        'success': False,
                        'url': url,
                        'file_path': str(path),
                        'error': str(e)
                    }
        
        tasks = [
            download_with_semaphore(url, path)
            for url, path in urls_and_paths
        ]
        
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理异常结果
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                url, path = urls_and_paths[i]
                processed_results.append({
                    'success': False,
                    'url': url,
                    'file_path': str(path),
                    'error': str(result)
                })
            else:
                processed_results.append(result)
        
        return processed_results
    
    def validate_url(self, url: str) -> bool:
        """
        验证URL格式
        
        Args:
            url: 要验证的URL
            
        Returns:
            是否为有效URL
        """
        try:
            from urllib.parse import urlparse
            result = urlparse(url)
            return all([result.scheme, result.netloc])
        except Exception:
            return False
    
    def get_file_extension_from_url(self, url: str) -> str:
        """
        从URL获取文件扩展名
        
        Args:
            url: 图片URL
            
        Returns:
            文件扩展名（包含点号）
        """
        try:
            from urllib.parse import urlparse
            path = urlparse(url).path
            if '.' in path:
                return Path(path).suffix.lower()
            return '.jpg'  # 默认扩展名
        except Exception:
            return '.jpg'