pollo_service_single.py

import json
import os
import base64
import time
import uuid
import requests
from typing import Dict, Any, Optional
import logging
from PIL import Image
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# 简化的配置类，替代animation_app.config
import os

class SimpleConfig:
    """简化的配置类，从环境变量读取配置"""
    
    def __init__(self):
        self.POLLO_API_KEY = os.getenv('POLLO_API_KEY', '')
        self.PUBLIC_DOMAIN = os.getenv('PUBLIC_DOMAIN', '')
        self.API_PREFIX = os.getenv('API_PREFIX', '/api/v1')

def get_settings():
    """获取配置设置"""
    return SimpleConfig()

logger = logging.getLogger(__name__)

class PolloAIService:
    """Pollo AI 视频生成服务（支持 Seedance Lite 和 Pro 版本）"""
    
    def __init__(self, backend: str = "bytedance/seedance"):
        # 从settings配置中获取统一的Pollo API配置
        self.settings = get_settings()
        
        self.api_key = self.settings.POLLO_API_KEY
        
        # 动态设置endpoint和backend，不从配置文件读取
        self.endpoint = "https://pollo.ai"
        self.backend = backend  # 可以动态指定: "bytedance/seedance", "bytedance/seedance-pro", 等
        
        # 判断是否为 Pro 版本
        self.is_pro_version = "pro" in self.backend.lower()

        # 设置带重试的Session
        self.session = requests.Session()
        retries = Retry(
            total=3, 
            backoff_factor=1, 
            status_forcelist=[500, 502, 503, 504],
            allowed_methods={"POST", "GET"}
        )
        adapter = HTTPAdapter(max_retries=retries)
        self.session.mount("https://", adapter)
        self.session.mount("http://", adapter)
        
        # 默认请求参数
        self.default_headers = {
            "x-api-key": self.api_key,
            "Content-Type": "application/json",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Accept": "application/json, text/plain, */*",
            "Accept-Language": "en-US,en;q=0.9",
            "Accept-Encoding": "gzip, deflate, br",
            "Connection": "keep-alive",
            "Sec-Fetch-Dest": "empty",
            "Sec-Fetch-Mode": "cors",
            "Sec-Fetch-Site": "same-origin"
        }
        
        self.timeout = 60  # API请求超时时间
        
        logger.info(f"初始化 Pollo AI 服务: {self.backend} ({'Pro' if self.is_pro_version else 'Lite'} 版本, 使用统一API)")
    
    def _get_api_endpoint(self) -> str:
        """根据版本获取正确的 API 端点"""
        # 根据最新文档，所有版本都使用 /api/platform/ 前缀
        return f"{self.endpoint}/api/platform/generation/{self.backend}"
    
    def _get_status_endpoint(self, task_id: str) -> str:
        """根据版本获取正确的状态查询端点"""
        # 根据最新文档，所有版本都使用 /api/platform/ 前缀
        return f"{self.endpoint}/api/platform/generation/{task_id}/status"
    
    def _validate_image_aspect_ratio(self, image_path: str) -> bool:
        """
        验证图片宽高比是否符合要求（必须小于1:4或4:1）
        
        Args:
            image_path: 图片路径
            
        Returns:
            bool: 是否符合要求
        """
        try:
            with Image.open(image_path) as img:
                width, height = img.size
                aspect_ratio = width / height
                
                # 检查宽高比是否在允许范围内 (1:4 < ratio < 4:1)
                if aspect_ratio < 0.25 or aspect_ratio > 4.0:
                    logger.error(f"图片宽高比不符合要求: {aspect_ratio:.2f} (允许范围: 0.25-4.0)")
                    return False
                
                logger.info(f"图片宽高比验证通过: {width}x{height} (比例: {aspect_ratio:.2f})")
                return True
                
        except Exception as e:
            logger.error(f"验证图片宽高比失败: {e}")
            return False
    
    def generate_video(
        self,
        prompt: str,
        mode: str = "i2v",
        input_image_path: Optional[str] = None,
        end_image_path: Optional[str] = None,
        symlink_folder: str = "pollo_images",  # 新增参数，默认值为pollo_images
        **kwargs
    ) -> Dict[str, Any]:
        """
        生成视频的主要方法 - 只提交任务，不轮询
        
        Args:
            prompt: 提示词
            mode: 生成模式 ('i2v', 'transition', 't2v')
            input_image_path: 输入图片路径
            end_image_path: 结束图片路径（用于transition模式，Pollo暂不支持）
            symlink_folder: 用于构建公网URL的文件夹名称（默认: pollo_images）
            **kwargs: 其他参数
            
        Returns:
            包含Pollo AI任务ID的字典，用于后续轮询
        """
        try:
            logger.info(f"开始提交Pollo AI视频生成任务，模式: {mode}, 提示词: {prompt[:100]}...")
            
            start_time = time.time()
            
            # 检查API密钥配置
            if not self.api_key:
                raise ValueError(
                    'Pollo AI服务未配置API密钥。请设置环境变量 POLLO_API_KEY。'
                    '获取密钥请访问：https://pollo.ai'
                )
            
            # 验证提示词长度（最大500字符）
            if len(prompt) > 500:
                raise ValueError(f"提示词长度超出限制，当前{len(prompt)}字符，最大允许500字符")
            
            # 验证视频长度参数
            video_length = kwargs.get('video_length') or 5
            
            # sora/sora-2-pro 特殊处理
            if self.backend == "sora/sora-2-pro":
                if video_length not in [4, 8, 12]:
                    logger.warning(f"sora/sora-2-pro 模型视频长度 {video_length}s 不在支持范围内 [4, 8, 12]，将自动调整为 4s")
                    video_length = 4
            # minimax-hailuo-02 特殊处理
            elif self.backend == "minimax/minimax-hailuo-02":
                if video_length not in [6, 10]:
                    logger.warning(f"minimax 模型视频长度 {video_length}s 不在支持范围内 [6, 10]，将自动调整为 6s")
                    video_length = 6
            # google/veo3 特殊处理
            elif self.backend == "google/veo3":
                if video_length != 8:
                    logger.warning(f"google/veo3 模型视频长度 {video_length}s 不在支持范围内 [8]，将自动调整为 8s")
                    video_length = 8
            elif video_length not in [5, 10]:
                logger.warning(f"视频长度 {video_length} 不在支持范围内，将调整为5秒")
                video_length = 5
            
            # 验证种子值 - 安全处理None值
            seed = kwargs.get('seed') or 123
            try:
                seed = int(seed)
                if seed > 2147483647:
                    logger.warning(f"种子值 {seed} 超出范围，将调整为随机值")
                    seed = 123
            except (ValueError, TypeError):
                logger.warning(f"无效的种子值: {seed}, 使用默认值123")
                seed = 123
            
            # 处理图片输入 - 只支持URL，不支持base64
            image_data = None
            image_tail_data = None  # Lite版本的结束图片
            
            if input_image_path:
                # 检查是否是URL（以http://或https://开头）
                if input_image_path.startswith(('http://', 'https://')):
                    # 直接使用URL
                    image_data = input_image_path
                    logger.info(f"使用图片URL: {image_data}")
                elif os.path.exists(input_image_path):
                    # 本地文件路径，需要转换为URL
                    # 验证图片宽高比
                    if not self._validate_image_aspect_ratio(input_image_path):
                        raise Exception("图片宽高比不符合要求（必须小于1:4或4:1）")
                    
                    # 只尝试构建公网可访问的URL，不再使用base64
                    public_image_url = self._try_get_public_image_url(input_image_path, symlink_folder)
                    
                    if public_image_url:
                        image_data = public_image_url
                        logger.info(f"使用公网起始图片URL: {public_image_url}")
                    else:
                        raise Exception(
                            "Pollo AI只接受图片URL，无法生成公网可访问的图片URL。"
                            "请配置PUBLIC_DOMAIN设置或确保图片可通过URL访问。"
                        )
                else:
                    raise Exception(f"无效的图片路径或URL: {input_image_path}")
            
            if not image_data and mode != "t2v":
                raise Exception("Pollo AI需要输入图片URL，但未提供有效的图片路径或无法生成公网URL")
            
            # 处理转场模式 - Lite版本支持imageTail，Pro版本显示警告
            if mode == "transition" and end_image_path:
                if self.is_pro_version:
                    logger.warning("Pollo AI Pro版本不支持转场模式的双图片输入，将忽略结束图片")
                else:
                    # Lite版本支持imageTail参数
                    if os.path.exists(end_image_path):
                        # 验证结束图片宽高比
                        if not self._validate_image_aspect_ratio(end_image_path):
                            raise Exception("结束图片宽高比不符合要求（必须小于1:4或4:1）")
                        
                        # 生成结束图片的公网URL
                        public_end_image_url = self._try_get_public_image_url(end_image_path, symlink_folder)
                        
                        if public_end_image_url:
                            image_tail_data = public_end_image_url
                            logger.info(f"使用公网结束图片URL: {public_end_image_url}")
                            logger.info(f"Lite版本转场模式: 起始图片 -> 结束图片")
                        else:
                            logger.warning("无法生成结束图片的公网URL，将忽略结束图片")
                    else:
                        logger.warning(f"结束图片路径不存在: {end_image_path}")
            
            # 安全地提取宽度和高度参数，确保不是None
            width = kwargs.get('width') or 1280
            height = kwargs.get('height') or 720
            
            # 准备请求参数（按照Pollo API格式）
            payload = {
                "input": {
                    "prompt": prompt,
                    "resolution": self._get_resolution(width, height),
                    "length": video_length,  # 使用验证过的视频长度
                    "seed": seed,  # 使用验证过的种子值
                    "cameraFixed": kwargs.get('camera_fixed', False)
                }
            }
            
            # minimax/minimax-hailuo-02 模型需要显式传递 videoModel
            if "minimax" in self.backend:
                # Pollo API for minimax expects the model name without the prefix
                model_name_only = self.backend.split('/')[-1]
                logger.info(f"为 minimax 模型添加 videoModel 参数: {model_name_only}")
                payload["input"]["videoModel"] = model_name_only
            
            # 添加图片参数（如果有）
            if image_data:
                payload["input"]["image"] = image_data
            
            # 添加结束图片参数（仅Lite版本支持）
            if image_tail_data and not self.is_pro_version:
                payload["input"]["imageTail"] = image_tail_data
            
            # 添加webhook URL（可选）
            webhook_url = kwargs.get('webhook_url')
            if webhook_url:
                payload["webhookUrl"] = webhook_url
            
            # 构建完整的API URL
            request_url = self._get_api_endpoint()
            
            logger.info(f"发送Pollo AI提交请求: {request_url}")
            logger.debug(f"请求体参数: backend={self.backend}, prompt长度={len(prompt)}, resolution={payload['input']['resolution']}, length={video_length}, seed={seed}")
            
            # 发送请求
            try:
                response = self.session.post(
                    request_url, 
                    headers=self.default_headers, 
                    json=payload, 
                    timeout=self.timeout
                )
                
                # 详细记录响应信息
                logger.info(f"Pollo AI响应状态: {response.status_code}")
                
                # 检查HTTP状态码
                if response.status_code != 200:
                    error_text = response.text
                    logger.error(f"Pollo AI HTTP错误: {response.status_code} - {error_text}")
                    
                    # 检查是否是 Cloudflare 403 错误
                    if response.status_code == 403 and "cloudflare" in error_text.lower():
                        logger.error("检测到 Cloudflare 反机器人保护，请检查:")
                        logger.error("1. API 密钥是否正确")
                        logger.error("2. 请求频率是否过高")
                        logger.error("3. IP 地址是否被限制")
                        raise Exception(f"Pollo AI 服务被 Cloudflare 阻止 (403)，可能是反机器人保护触发")
                    
                    # 检查是否是 API 密钥错误
                    if response.status_code == 401:
                        raise Exception(f"Pollo AI API 密钥无效或过期 (401)")
                    
                    # 检查是否是速率限制
                    if response.status_code == 429:
                        raise Exception(f"Pollo AI API 请求频率过高 (429)，请稍后重试")
                    
                    raise Exception(f"Pollo AI服务HTTP错误: {response.status_code}")
                
                # 解析响应
                try:
                    result_data = response.json()
                except json.JSONDecodeError as e:
                    logger.error(f"Pollo AI响应JSON解析失败: {e}")
                    logger.error(f"原始响应: {response.text[:1000]}...")
                    raise Exception("Pollo AI响应格式错误")
                
                logger.debug(f"解析后的响应: {result_data}")
                
                processing_time = time.time() - start_time
                
                # 检查响应格式 - Pollo API 使用包装格式
                if 'code' in result_data and 'data' in result_data:
                    # 新的包装格式：{"code": "SUCCESS", "message": "success", "data": {...}}
                    code = result_data.get('code')
                    message = result_data.get('message', '')
                    data = result_data.get('data', {})
                    
                    if code != 'SUCCESS':
                        raise Exception(f"Pollo AI 请求失败: {code} - {message}")
                    
                    # 从 data 字段提取任务信息
                    task_id = data.get('taskId')
                    status = data.get('status')
                    
                    logger.info(f"Pollo AI 响应格式: 包装格式 (code={code}, message={message})")
                    
                else:
                    # 旧的直接格式（向后兼容）
                    task_id = result_data.get('taskId')
                    status = result_data.get('status')
                    data = result_data
                    
                    logger.info(f"Pollo AI 响应格式: 直接格式")
                
                if task_id:
                    logger.info(f"Pollo AI任务提交成功，任务ID: {task_id}, 状态: {status}")
                    
                    # 返回任务信息，不进行轮询
                    return {
                        'pollo_task_id': task_id,  # Pollo AI的任务ID
                        'submit_time': processing_time,
                        'status': 'submitted',  # 标记为已提交
                        'initial_status': status,  # Pollo返回的初始状态
                        'generation_info': {
                            'engine': 'pollo',
                            'backend': self.backend,
                            'mode': mode,
                            'prompt': prompt,
                            'input_image': input_image_path,
                            'end_image': end_image_path,
                            'submit_response': result_data,  # 完整响应
                            'submit_data': data,  # 数据部分
                            'validated_params': {
                                'prompt_length': len(prompt),
                                'video_length': video_length,
                                'seed': seed,
                                'resolution': payload['input']['resolution']
                            }
                        }
                    }
                else:
                    raise Exception("Pollo AI响应中未找到任务ID")
                
            except requests.exceptions.RequestException as e:
                logger.error(f"Pollo AI请求失败: {e}")
                raise Exception(f"Pollo AI服务请求失败: {str(e)}")
                
        except Exception as e:
            logger.error(f"Pollo AI视频生成任务提交失败: {str(e)}")
            raise Exception(f"Pollo AI视频生成任务提交失败: {str(e)}")
    
    def _get_resolution(self, width: int, height: int) -> str:
        """
        根据宽高确定分辨率字符串
        Lite版本: 支持480p, 720p, 1080p
        Pro版本: 只支持480p, 1080p
        Vidu-q1: 只支持1080p
        minimax/minimax-hailuo-02: 支持 768P, 1080P
        google/veo3: 支持720p, 1080p
        sora/sora-2-pro: 支持720p, 1080p
        """
        # sora/sora-2-pro 模型特殊处理
        if self.backend == "sora/sora-2-pro":
            logger.info("检测到 sora/sora-2-pro 模型，自动选择 720p 或 1080p 分辨率")
            
            _width = width or 1280
            _height = height or 720
            total_pixels = _width * _height
            
            # 720p 约 921,600 像素
            # 1080p 约 2,073,600 像素
            # 使用 1.5M 像素作为分界线
            if total_pixels > 1500000:
                return "1080p"
            else:
                return "720p"
        
        # minimax/minimax-hailuo-02 模型特殊处理
        if self.backend == "minimax/minimax-hailuo-02":
            logger.info("检测到 minimax/minimax-hailuo-02 模型，自动选择 768P 或 1080P 分辨率")
            
            _width = width or 1280
            _height = height or 720
            total_pixels = _width * _height

            # 1080P 约 2.07M 像素
            # 768P  约 0.78M 像素 (e.g., 1024x768)
            # 使用 1.5M 像素作为分界线
            if total_pixels > 1500000:
                return "1080P"
            else:
                return "768P"

        # Vidu-q1 模型特殊处理，强制使用1080p
        if self.backend == "vidu/vidu-q1":
            logger.info("检测到 vidu/vidu-q1 模型，强制使用 1080p 分辨率")
            return "1080p"
            
        # google/veo3 模型特殊处理，支持720p和1080p
        if self.backend == "google/veo3":
            logger.info("检测到 google/veo3 模型，自动选择 720p 或 1080p 分辨率")
            
            _width = width or 1280
            _height = height or 720
            total_pixels = _width * _height
            
            # 720p 约 921,600 像素
            # 1080p 约 2,073,600 像素
            # 使用 1.5M 像素作为分界线
            if total_pixels > 1500000:
                return "1080p"
            else:
                return "720p"
            
        # 处理None值，使用默认值
        if width is None:
            width = 1280
        if height is None:
            height = 720
            
        # 确保值是整数
        try:
            width = int(width)
            height = int(height)
        except (ValueError, TypeError):
            logger.warning(f"无效的宽高值: width={width}, height={height}, 使用默认值")
            width, height = 1280, 720
        
        # 根据像素数判断，选择最接近的支持分辨率
        total_pixels = width * height
        
        if self.is_pro_version:
            # Pro版本：只支持480p和1080p
            # 480p约为640x480 = 307,200像素
            # 1080p约为1920x1080 = 2,073,600像素
            threshold = (307200 + 2073600) / 2  # 约1,190,400像素
            
            if total_pixels <= threshold:
                return "480p"
            else:
                return "1080p"
        else:
            # Lite版本：支持480p, 720p, 1080p
            # 480p约为640x480 = 307,200像素
            # 720p约为1280x720 = 921,600像素  
            # 1080p约为1920x1080 = 2,073,600像素
            
            if total_pixels <= 600000:  # 约 600k 像素作为 480p 的上限
                return "480p"
            elif total_pixels <= 1500000:  # 约 1.5M 像素作为 720p 的上限
                return "720p"
            else:
                return "1080p"
    
    def _try_get_public_image_url(self, local_path: str, folder_name: str = "pollo_images") -> Optional[str]:
        """
        尝试将本地图片路径转换为公网可访问的URL
        
        Args:
            local_path: 本地图片文件路径
            folder_name: 用于构建URL的文件夹名称（默认: pollo_images）
            
        Returns:
            公网可访问的URL，如果无法生成则返回None
        """
        try:
            # 检查是否配置了公网访问域名
            public_domain = self.settings.PUBLIC_DOMAIN
            if not public_domain:
                logger.warning("未配置PUBLIC_DOMAIN设置，无法生成公网图片URL")
                return None
            
            # 提取文件名
            filename = os.path.basename(local_path)
            
            # 构建公网URL - 使用传入的文件夹名称
            public_url = f"{public_domain.rstrip('/')}/{folder_name}/{filename}"
            
            logger.info(f"生成公网图片URL: {public_url}")
            return public_url
            
        except Exception as e:
            logger.warning(f"无法生成公网图片URL: {e}")
            return None
    
    def poll_task_result(self, pollo_task_id: str) -> Dict[str, Any]:
        """
        轮询Pollo AI任务结果 - 单次查询，不循环
        
        Args:
            pollo_task_id: Pollo AI的任务ID
            
        Returns:
            包含任务状态和结果的字典
        """
        try:
            logger.info(f"轮询Pollo AI任务状态: {pollo_task_id}")
            
            # 构建查询URL
            request_url = self._get_status_endpoint(pollo_task_id)
            
            # 发送查询请求
            response = self.session.get(request_url, headers={"x-api-key": self.api_key}, timeout=self.timeout)
            response.raise_for_status()
            
            result_data = response.json()
            
            # 详细记录轮询响应（用于调试进度展示问题）
            logger.info(f"Pollo AI 轮询响应 (任务ID: {pollo_task_id}):")
            logger.info(f"  - HTTP状态码: {response.status_code}")
            logger.info(f"  - 响应头: {dict(response.headers)}")
            logger.info(f"  - 原始响应: {json.dumps(result_data, indent=2, ensure_ascii=False)}")
            
            # 检查查询响应格式 - Pollo API 使用包装格式
            if 'code' in result_data and 'data' in result_data:
                # 新的包装格式：{"code": "SUCCESS", "message": "success", "data": {...}}
                code = result_data.get('code')
                message = result_data.get('message', '')
                data = result_data.get('data', {})
                
                if code != 'SUCCESS':
                    logger.error(f"Pollo AI查询失败: {code} - {message}")
                    return {
                        'status': 'failed',
                        'error_message': f"Pollo AI查询失败: {code} - {message}"
                    }
                
                # 从 data 字段提取任务信息
                task_id = data.get('taskId')
                generations = data.get('generations', [])
                
                logger.info(f"Pollo AI 查询响应格式: 包装格式 (code={code})")
                
            else:
                # 旧的直接格式（向后兼容）
                task_id = result_data.get('taskId')
                generations = result_data.get('generations', [])
                data = result_data
                
                logger.info(f"Pollo AI 查询响应格式: 直接格式")
            
            if not task_id or not generations:
                logger.error(f"Pollo AI查询响应格式错误: {result_data}")
                return {
                    'status': 'failed',
                    'error_message': "Pollo AI查询响应格式错误"
                }
            
            # 获取第一个生成结果的状态
            generation = generations[0]
            status = generation.get('status')
            fail_msg = generation.get('failMsg')
            video_url = generation.get('url')
            media_type = generation.get('mediaType')
            
            if status == 'succeed' and video_url and media_type == 'video':
                # 任务完成
                logger.info(f"Pollo AI视频生成完成: {video_url}")
                
                return {
                    'status': 'completed',
                    'video_url': video_url,
                    'api_response': result_data
                }
                
            elif status == 'failed':
                error_msg = fail_msg or '任务失败'
                logger.error(f"Pollo AI任务失败: {error_msg}")
                return {
                    'status': 'failed',
                    'error_message': f"Pollo AI任务失败: {error_msg}"
                }
            
            elif status in ['waiting', 'processing']:
                # 任务还在处理中
                logger.info(f"Pollo AI任务处理中，状态: {status}")
                return {
                    'status': 'processing',
                    'task_status': status
                }
                
            else:
                logger.warning(f"Pollo AI未知任务状态: {status}")
                return {
                    'status': 'processing',
                    'task_status': status
                }
                
        except Exception as e:
            logger.error(f"轮询Pollo AI任务结果失败: {e}")
            return {
                'status': 'failed',
                'error_message': f"轮询任务失败: {str(e)}"
            }
    
    def get_model_status(self) -> Dict[str, Any]:
        """获取模型状态"""
        # 根据不同模型确定支持的分辨率
        if self.backend == "google/veo3":
            supported_resolutions = ['720p', '1080p']
        elif self.backend == "vidu/vidu-q1":
            supported_resolutions = ['1080p']
        elif self.backend == "minimax/minimax-hailuo-02":
            supported_resolutions = ['768P', '1080P']
        elif self.is_pro_version:
            supported_resolutions = ['480p', '1080p']
        else:
            supported_resolutions = ['480p', '720p', '1080p']
            
        return {
            'engine': 'pollo',
            'backend': self.backend,
            'version': 'Pro' if self.is_pro_version else 'Lite',
            'status': 'ready' if self.api_key else 'not_configured',
            'model_loaded': bool(self.api_key),
            'gpu_device_id': None,  # Pollo AI是云服务，无需GPU管理
            'memory_usage': 0,
            'config': {
                'api_key_configured': bool(self.api_key and self.api_key != ''),
                'endpoint': self.endpoint,
                'backend': self.backend,
                'version': 'Pro' if self.is_pro_version else 'Lite',
                'supports_transition': not self.is_pro_version,  # Lite版本支持转场
                'supported_resolutions': supported_resolutions,
                'is_configured': bool(self.api_key),
                'config_url': 'https://pollo.ai'
            }
        }

# 创建全局实例（默认Seedance Lite版本，保持向后兼容）
pollo_service = PolloAIService() 

def get_pollo_service(backend: str = "bytedance/seedance") -> PolloAIService:
    """
    工厂函数：获取指定backend的Pollo AI服务实例
    
    Args:
        backend: 支持的backend类型
            - "bytedance/seedance" (默认，Lite版本)
            - "bytedance/seedance-pro" (Pro版本)
            - 未来可支持更多backend
    
    Returns:
        PolloAIService实例
    """
    return PolloAIService(backend=backend)