"""
Unified Task Manager

统一任务管理器，整合了 TaskManagerV2 和 EnhancedTaskManager 的所有功能。

主要特性：
1. 任务队列管理（优先队列）
2. 并发控制（Semaphore）
3. 工作线程池
4. 任务优先级支持
5. 指数退避重试机制
6. 任务持久化和恢复
7. Prometheus 指标集成
8. 超时检测
9. 任务取消支持

Requirements: 2.1, 2.2, 2.3, 2.4
"""

import logging
import asyncio
import uuid
import time
from typing import Optional, Dict, Any, Callable, List
from datetime import datetime

from sqlmodel import Session, select
from prometheus_client import Counter, Gauge, Histogram, Summary

from src.config.database import engine, init_db
from src.models.entities import TaskDB
from src.models.schemas import Task
from src.mappers import TaskMapper
from src.services.provider.registry import ModelRegistry, ModelType
from src.services.provider.base import TaskStatus, map_provider_status, is_terminal_status
from src.utils.errors import (
    TaskQueueFullException,
    TaskTimeoutException,
    TaskNotFoundException,
    GenerationFailedException
)

from .models import TaskPriority, TaskConfig, TaskItem
from .metrics import (
    task_created_counter,
    task_completed_counter,
    task_failed_counter,
    task_retry_counter,
    task_queue_size_gauge,
    task_active_gauge,
    task_duration_histogram,
    task_wait_time_histogram,
    task_execution_summary,
)

logger = logging.getLogger(__name__)


class UnifiedTaskManager:
    """统一任务管理器

    整合了所有任务管理功能，提供单一、清晰的任务管理接口。

    特性：
    - 任务队列管理（优先队列）
    - 并发控制（每种任务类型独立的 Semaphore）
    - 工作线程池
    - 任务优先级
    - 指数退避重试机制
    - 任务持久化和系统重启后恢复
    - 超时检测
    - Prometheus 指标集成

    Requirements: 2.1, 2.2, 2.3, 2.4
    """

    # 默认轮询配置
    DEFAULT_POLL_INTERVAL = 2
    DEFAULT_MAX_POLLS = 150  # 5 分钟

    def __init__(self):
        # 任务配置
        self.task_configs: Dict[str, TaskConfig] = {
            "image": TaskConfig(max_retries=3, timeout=180, max_concurrent=10),
            "video": TaskConfig(max_retries=5, timeout=600, max_concurrent=5),
            "audio": TaskConfig(max_retries=3, timeout=240, max_concurrent=8),
            "music": TaskConfig(max_retries=3, timeout=600, max_concurrent=5),  # 音乐生成耗时长，10 分钟
            "script": TaskConfig(max_retries=2, timeout=120, max_concurrent=3)
        }

        # 并发控制（每种任务类型独立的信号量）
        self._semaphores: Dict[str, asyncio.Semaphore] = {}
        for task_type, config in self.task_configs.items():
            self._semaphores[task_type] = asyncio.Semaphore(config.max_concurrent)

        # 任务队列（使用优先队列）
        self._task_queue: asyncio.PriorityQueue = asyncio.PriorityQueue(maxsize=1000)

        # 工作线程
        self._workers: List[asyncio.Task] = []
        self._num_workers = 5
        self._running = False

        # 活跃任务跟踪
        self._active_tasks: Dict[str, asyncio.Task] = {}

        # 统计信息
        self._stats = {
            "total_tasks": 0,
            "completed_tasks": 0,
            "failed_tasks": 0,
            "queue_size": 0
        }

        # 初始化数据库
        init_db()

    @property
    def max_workers(self) -> int:
        """获取最大工作线程数"""
        return self._num_workers

    async def start(self):
        """启动任务管理器"""
        if self._running:
            logger.warning("Task manager is already running")
            return

        self._running = True

        # 恢复未完成的任务（系统重启后）
        await self._recover_pending_tasks()

        # 启动工作线程
        for i in range(self._num_workers):
            worker = asyncio.create_task(self._worker(f"worker-{i}"))
            self._workers.append(worker)

        logger.info(f"Task manager started with {self._num_workers} workers")

    async def _recover_pending_tasks(self):
        """从数据库恢复未完成的任务

        在系统重启后，将所有未完成的任务重新加入队列。
        这确保了任务的持久化和恢复能力。

        Requirements: 2.3
        """
        try:
            with Session(engine) as session:
                # 查询所有未完成的任务
                statement = select(TaskDB).where(
                    TaskDB.status.in_([
                        TaskStatus.PENDING.value,
                        TaskStatus.PROCESSING.value,
                        TaskStatus.RETRYING.value
                    ])
                )
                pending_tasks = session.exec(statement).all()

                recovered_count = 0
                for task_db in pending_tasks:
                    # 重置状态为 PENDING
                    task_db.status = TaskStatus.PENDING.value
                    task_db.updated_at = datetime.now().timestamp()
                    session.add(task_db)

                    # 加入队列（使用 NORMAL 优先级）
                    task_item = TaskItem(task_db.id, TaskPriority.NORMAL)
                    await self._task_queue.put(task_item)
                    recovered_count += 1

                session.commit()

                if recovered_count > 0:
                    logger.info(f"Recovered {recovered_count} pending tasks from database")
                    self._stats["total_tasks"] += recovered_count
                    self._stats["queue_size"] = self._task_queue.qsize()
                    task_queue_size_gauge.set(self._task_queue.qsize())

        except Exception as e:
            logger.error(f"Failed to recover pending tasks: {e}", exc_info=True)

    async def stop(self):
        """停止任务管理器"""
        if not self._running:
            return

        self._running = False

        # 等待队列中的任务完成
        await self._task_queue.join()

        # 取消所有工作线程
        for worker in self._workers:
            worker.cancel()

        # 等待工作线程结束
        await asyncio.gather(*self._workers, return_exceptions=True)

        self._workers.clear()
        logger.info("Task manager stopped")

    async def create_task(
        self,
        task_type: str,
        model: str,
        params: Dict[str, Any],
        config: Optional[TaskConfig] = None,
        user_id: Optional[str] = None,
        project_id: Optional[str] = None,
        priority: TaskPriority = TaskPriority.NORMAL
    ) -> Task:
        """创建任务并加入队列

        Args:
            task_type: 任务类型（image, video, script）
            model: 模型复合 ID（格式：provider/model_key）
            params: 任务参数
            config: 任务配置（可选）
            user_id: 用户 ID
            project_id: 项目 ID
            priority: 任务优先级

        Returns:
            Task: 创建的任务对象

        Raises:
            TaskQueueFullException: 队列已满
            ValueError: model 格式不正确
        """
        # 检查队列是否已满
        if self._task_queue.full():
            raise TaskQueueFullException(self._task_queue.maxsize)

        # 从复合 ID 中提取 provider
        # model 格式应该是 "provider/model_key"
        provider = None
        if model and '/' in model:
            provider = model.split('/', 1)[0]

        # 使用提供的配置或默认配置
        task_config = config if config else self.task_configs.get(task_type, TaskConfig())

        # 创建 TaskDB 对象
        # model 字段存储完整的复合 ID
        # provider 字段存储提取的 provider（用于索引和查询）
        task_db = TaskDB(
            id=str(uuid.uuid4()),
            type=task_type,
            model=model,
            provider=provider,
            params=params,
            status=TaskStatus.PENDING.value,
            retry_count=0,
            max_retries=task_config.max_retries,
            user_id=user_id,
            project_id=project_id
        )

        # 保存到数据库
        with Session(engine) as session:
            session.add(task_db)
            session.commit()
            session.refresh(task_db)

        # 加入队列
        task_item = TaskItem(task_db.id, priority)
        await self._task_queue.put(task_item)

        # 更新统计
        self._stats["total_tasks"] += 1
        self._stats["queue_size"] = self._task_queue.qsize()

        # 更新 Prometheus 指标
        task_created_counter.labels(type=task_type, priority=priority.name).inc()
        task_queue_size_gauge.set(self._task_queue.qsize())

        logger.info(
            f"Task created: {task_db.id} (type={task_type}, priority={priority.name}, queue_size={self._stats['queue_size']})"
        )

        return self._db_to_task(task_db)

    async def get_task(self, task_id: str) -> Optional[Task]:
        """获取任务信息"""
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                return None
            return self._db_to_task(task_db)

    def list_tasks(self, type: Optional[str] = None, limit: int = 50, offset: int = 0, user_id: Optional[str] = None) -> List[Task]:
        """列出任务

        Args:
            type: 任务类型过滤（可选）
            limit: 返回数量限制
            offset: 偏移量
            user_id: 用户ID过滤（可选）

        Returns:
            List[Task]: 任务列表
        """
        from sqlmodel import desc

        with Session(engine) as session:
            statement = select(TaskDB).order_by(desc(TaskDB.created_at))
            if type:
                statement = statement.where(TaskDB.type == type)
            if user_id:
                statement = statement.where(TaskDB.user_id == user_id)
            statement = statement.offset(offset).limit(limit)

            results = session.exec(statement).all()
            return [self._db_to_task(task) for task in results]

    def update_task(self, task_id: str, **updates) -> Optional[Task]:
        """更新任务信息（同步方法）

        Args:
            task_id: 任务 ID
            **updates: 要更新的字段

        Returns:
            Task: 更新后的任务对象

        Raises:
            TaskNotFoundException: 任务不存在
        """
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                raise TaskNotFoundException(task_id)

            # 更新字段
            for key, value in updates.items():
                if hasattr(task_db, key):
                    setattr(task_db, key, value)

            task_db.updated_at = datetime.now().timestamp()
            session.add(task_db)
            session.commit()
            session.refresh(task_db)

            return self._db_to_task(task_db)

    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                raise TaskNotFoundException(task_id)

            if task_db.status in [TaskStatus.SUCCEEDED.value, TaskStatus.FAILED.value]:
                return False  # 已完成的任务无法取消

            task_db.status = TaskStatus.CANCELLED.value
            task_db.completed_at = datetime.now().timestamp()
            session.commit()
            return True

    async def retry_task(self, task_id: str) -> bool:
        """手动重试失败的任务

        Args:
            task_id: 任务 ID

        Returns:
            bool: 是否成功加入重试队列

        Raises:
            TaskNotFoundException: 任务不存在
        """
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                raise TaskNotFoundException(task_id)

            # 只有失败或超时的任务可以重试
            if task_db.status not in [TaskStatus.FAILED.value, TaskStatus.TIMEOUT.value]:
                logger.warning(f"Task {task_id} cannot be retried (status: {task_db.status})")
                return False

            # 重置任务状态
            task_db.status = TaskStatus.PENDING.value
            task_db.retry_count = 0
            task_db.error = None
            task_db.updated_at = datetime.now().timestamp()

            # 清除完成时间
            task_db.completed_at = None

            session.commit()
            session.refresh(task_db)

            # 重新加入队列
            task_item = TaskItem(task_db.id, TaskPriority.NORMAL)
            await self._task_queue.put(task_item)

            logger.info(f"Task {task_id} added to retry queue")
            return True

    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        return {
            **self._stats,
            "queue_size": self._task_queue.qsize(),
            "workers": len(self._workers),
            "running": self._running
        }

    async def _worker(self, name: str):
        """工作线程

        从队列中获取任务并执行
        """
        logger.info(f"Worker {name} started")

        while self._running:
            try:
                # 从队列获取任务（带超时，避免阻塞）
                try:
                    task_item: TaskItem = await asyncio.wait_for(
                        self._task_queue.get(),
                        timeout=1.0
                    )
                except asyncio.TimeoutError:
                    continue

                # 跟踪活跃任务
                self._active_tasks[task_item.task_id] = asyncio.current_task()

                # 执行任务
                try:
                    await self._execute_task(task_item.task_id)
                except Exception as e:
                    logger.error(f"Worker {name} failed to execute task {task_item.task_id}: {e}", exc_info=True)
                finally:
                    # 移除活跃任务跟踪
                    self._active_tasks.pop(task_item.task_id, None)
                    self._task_queue.task_done()
                    self._stats["queue_size"] = self._task_queue.qsize()

            except asyncio.CancelledError:
                logger.info(f"Worker {name} cancelled")
                break
            except Exception as e:
                logger.error(f"Worker {name} error: {e}", exc_info=True)

        logger.info(f"Worker {name} stopped")

    async def _execute_task(self, task_id: str):
        """执行任务（带并发控制和重试）

        实现了完整的重试机制，包括：
        - 指数退避重试
        - 重试历史记录
        - 错误详情记录

        Requirements: 2.4
        """
        # 加载任务
        with Session(engine) as session:
            task_db = session.get(TaskDB, task_id)
            if not task_db:
                logger.error(f"Task {task_id} not found")
                return

        # 记录任务开始时间（用于计算等待时间）
        task_start_time = time.time()
        wait_time = task_start_time - task_db.created_at
        task_wait_time_histogram.labels(type=task_db.type).observe(wait_time)

        # 获取任务配置
        config = self.task_configs.get(task_db.type, TaskConfig())

        # 获取信号量（并发控制）
        semaphore = self._semaphores.get(task_db.type, asyncio.Semaphore(5))

        # 初始化重试历史（如果不存在）
        if not task_db.result:
            task_db.result = {}
        if 'retry_history' not in task_db.result:
            task_db.result['retry_history'] = []

        # 使用信号量控制并发
        async with semaphore:
            # 更新活跃任务计数
            task_active_gauge.labels(type=task_db.type).inc()

            # 记录执行开始时间
            execution_start_time = time.time()

            # 设置日志上下文，包含 user_id
            from src.utils.logging import set_log_context, clear_log_context
            log_context_set = False
            if task_db.user_id:
                set_log_context(user_id=task_db.user_id, task_id=str(task_db.id))
                log_context_set = True

            try:
                # 重试循环
                while task_db.retry_count <= task_db.max_retries:
                    retry_start_time = time.time()

                    try:
                        # 更新状态为处理中
                        task_db.status = TaskStatus.PROCESSING.value
                        task_db.started_at = datetime.now().timestamp()
                        await self._save_task(task_db)

                        # 执行任务（带超时）
                        result = await asyncio.wait_for(
                            self._call_provider(task_db),
                            timeout=config.timeout
                        )

                        # 成功
                        task_db.status = TaskStatus.SUCCEEDED.value
                        # 保存结果（保持与前端期望的格式一致）
                        # 前端期望: result.urls 或 result.data.urls
                        if isinstance(result, dict):
                            # 如果结果已经是字典格式，直接使用
                            task_db.result = result
                        else:
                            task_db.result = result
                        task_db.completed_at = datetime.now().timestamp()
                        await self._save_task(task_db)

                        # 记录成功指标
                        execution_duration = time.time() - execution_start_time
                        task_duration_histogram.labels(type=task_db.type).observe(execution_duration)
                        task_execution_summary.labels(type=task_db.type).observe(execution_duration)
                        task_completed_counter.labels(type=task_db.type, status='succeeded').inc()

                        self._stats["completed_tasks"] += 1
                        logger.info(f"Task {task_id} completed successfully in {execution_duration:.2f}s")
                        break

                    except asyncio.TimeoutError:
                        retry_duration = time.time() - retry_start_time
                        error_msg = f"Task timed out after {config.timeout} seconds"

                        # 记录重试历史
                        task_db.result['retry_history'].append({
                            'attempt': task_db.retry_count + 1,
                            'timestamp': datetime.now().timestamp(),
                            'duration': retry_duration,
                            'error': error_msg,
                            'error_type': 'timeout'
                        })

                        logger.error(f"Task {task_id} timed out after {config.timeout}s")
                        task_db.status = TaskStatus.TIMEOUT.value
                        task_db.error = error_msg
                        task_db.completed_at = datetime.now().timestamp()
                        await self._save_task(task_db)

                        # 记录失败指标
                        task_failed_counter.labels(type=task_db.type, reason='timeout').inc()
                        task_completed_counter.labels(type=task_db.type, status='timeout').inc()

                        self._stats["failed_tasks"] += 1
                        break

                    except Exception as e:
                        retry_duration = time.time() - retry_start_time
                        error_msg = str(e)

                        # 记录重试历史
                        task_db.result['retry_history'].append({
                            'attempt': task_db.retry_count + 1,
                            'timestamp': datetime.now().timestamp(),
                            'duration': retry_duration,
                            'error': error_msg,
                            'error_type': type(e).__name__
                        })

                        logger.error(f"Task {task_id} failed (attempt {task_db.retry_count + 1}): {e}")
                        task_db.retry_count += 1

                        if task_db.retry_count > task_db.max_retries:
                            # 重试次数用尽
                            task_db.status = TaskStatus.FAILED.value
                            task_db.error = error_msg
                            task_db.completed_at = datetime.now().timestamp()
                            await self._save_task(task_db)

                            # 记录失败指标
                            task_failed_counter.labels(type=task_db.type, reason='max_retries').inc()
                            task_completed_counter.labels(type=task_db.type, status='failed').inc()

                            self._stats["failed_tasks"] += 1
                            logger.error(
                                f"Task {task_id} failed after {task_db.max_retries} retries. "
                                f"Retry history: {len(task_db.result.get('retry_history', []))} attempts"
                            )
                            break
                        else:
                            # 等待后重试
                            task_db.status = TaskStatus.RETRYING.value
                            await self._save_task(task_db)

                            # 记录重试指标
                            task_retry_counter.labels(type=task_db.type).inc()

                            # 计算重试延迟（指数退避）
                            delay = config.retry_delay
                            if config.exponential_backoff:
                                delay = config.retry_delay * (2 ** (task_db.retry_count - 1))

                            logger.info(
                                f"Task {task_id} will retry in {delay}s "
                                f"(attempt {task_db.retry_count + 1}/{task_db.max_retries})"
                            )
                            await asyncio.sleep(delay)
            finally:
                # 更新活跃任务计数
                task_active_gauge.labels(type=task_db.type).dec()
                # 更新队列大小
                task_queue_size_gauge.set(self._task_queue.qsize())
                # 清除日志上下文
                if log_context_set:
                    clear_log_context()

    async def _call_provider(self, task_db: TaskDB) -> Dict[str, Any]:
        """调用 AI 服务提供商"""
        # 获取服务
        service = ModelRegistry.get(task_db.model)
        if not service:
            # 使用默认服务
            if task_db.type == "image":
                service = ModelRegistry.get_default(ModelType.IMAGE)
            elif task_db.type == "video":
                service = ModelRegistry.get_default(ModelType.VIDEO)
            elif task_db.type == "audio":
                service = ModelRegistry.get_default(ModelType.AUDIO)
            elif task_db.type == "music":
                service = ModelRegistry.get_default(ModelType.MUSIC)
            else:
                raise GenerationFailedException(f"No service found for model {task_db.model}")

        if not service:
            raise GenerationFailedException(f"No service available for task type {task_db.type}")

        # 清理参数：移除非 provider 需要的字段
        provider_params = self._prepare_provider_params(task_db.params, task_db.type)

        # 添加 user_id 到 provider_params，用于获取用户管理的 API key
        if task_db.user_id:
            provider_params['user_id'] = task_db.user_id

        # 调用服务 - 使用统一的 generate() 方法
        # 所有服务（图片、视频、音频）都实现了 generate() 方法
        if task_db.type == "image":
            # 图片服务：generate(prompt, **kwargs)
            response = await service.generate(**provider_params)
        elif task_db.type == "video":
            # 视频服务：generate(prompt, **kwargs)
            response = await service.generate(**provider_params)
        elif task_db.type == "audio":
            # 音频服务：generate(text, voice, **kwargs)
            response = await service.generate(**provider_params)
        elif task_db.type == "music":
            # 统一音乐服务：根据 generation_mode 分发到 lyrics/music
            generation_mode = provider_params.get("generation_mode")
            if generation_mode == "lyrics":
                if not hasattr(service, "generate_lyrics"):
                    raise GenerationFailedException(
                        f"Service {service.__class__.__name__} does not support lyrics generation"
                    )
                extras = {
                    k: v for k, v in provider_params.items()
                    if k not in {"generation_mode", "prompt", "lyrics_mode", "seed_lyrics", "title"}
                }
                lyrics_data = await service.generate_lyrics(
                    prompt=provider_params.get("prompt", "") or "",
                    mode=provider_params.get("lyrics_mode", "write_full_song"),
                    lyrics=provider_params.get("seed_lyrics"),
                    title=provider_params.get("title"),
                    **extras,
                )
                return {
                    "mode": "lyrics",
                    "song_title": (lyrics_data or {}).get("song_title"),
                    "style_tags": (lyrics_data or {}).get("style_tags"),
                    "lyrics": (lyrics_data or {}).get("lyrics", ""),
                }

            # 音乐服务：generate(lyrics, prompt, **kwargs)
            response = await service.generate(**provider_params)
        else:
            raise GenerationFailedException(f"Unsupported task type: {task_db.type}")

        # 检查是否是异步任务（需要轮询）
        # 不同 Provider 返回不同的初始状态：PENDING 或 QUEUED
        if hasattr(response, 'status') and response.status in [TaskStatus.PENDING, TaskStatus.QUEUED]:
            if hasattr(response, 'task_id') and response.task_id:
                # 异步任务，需要轮询
                logger.info(f"Task {task_db.id} is async, provider task_id: {response.task_id}, status: {response.status}")
                task_db.provider_task_id = response.task_id
                await self._save_task(task_db)

                # 轮询直到完成
                poll_interval = 5
                poll_attempts = 120
                service_module = getattr(service.__class__, "__module__", "")
                if task_db.type == "video" and "provider.minimax" in service_module:
                    # MiniMax 官方建议 10s 轮询，降低限频与状态抖动风险
                    poll_interval = 10
                    poll_attempts = 180  # 30 分钟窗口

                response = await self._poll_provider_task(
                    service,
                    response.task_id,
                    task_db.type,
                    max_attempts=poll_attempts,
                    interval=poll_interval,
                    user_id=task_db.user_id,
                )

        # 提取结果
        result = self._extract_result(response, task_db.type)

        # 检查是否有错误
        if result.get('error'):
            raise GenerationFailedException(result['error'])

        return result

    def _prepare_provider_params(self, params: Dict[str, Any], task_type: str) -> Dict[str, Any]:
        """
        准备传递给 Provider 的参数

        清理和转换参数：
        1. 移除非 provider 需要的元数据字段
        2. 确保 prompt 是字符串（不是 None）
        3. 合并 extra_params
        4. 标准化图片参数
        """
        # 需要移除的元数据字段（这些是请求级别的，不传给 provider）
        metadata_fields = {
            'model', 'provider', 'source', 'source_id', 'project_id',
            'storyboard_id'
        }

        # 复制参数，避免修改原始数据
        cleaned = {}
        extra_params = {}

        for key, value in params.items():
            # 跳过 None 值
            if value is None:
                continue

            # 提取 extra_params 以便后续合并
            if key == 'extra_params' and isinstance(value, dict):
                extra_params = value
                continue

            # 跳过元数据字段
            if key in metadata_fields:
                continue

            cleaned[key] = value

        # 合并 extra_params（这些是用户指定的额外参数）
        for key, value in extra_params.items():
            if value is not None and key not in metadata_fields:
                cleaned[key] = value

        # 确保 prompt 是字符串
        if 'prompt' not in cleaned or cleaned.get('prompt') is None:
            cleaned['prompt'] = ""

        # 特定类型的参数规范化
        if task_type == "image":
            # 图片生成参数规范化
            # ref_images 已经在 generations.py 中映射
            pass

        elif task_type == "video":
            # 视频生成参数规范化
            # 确保 duration 是字符串（某些 provider 需要）
            if 'duration' in cleaned and isinstance(cleaned['duration'], int):
                cleaned['duration'] = cleaned['duration']  # 保持 int，provider 会处理

            # 统一媒体输入参数：image_inputs / audio_inputs
            self._normalize_video_image_params(cleaned)
            self._normalize_video_audio_params(cleaned)

        return cleaned

    def _normalize_video_audio_params(self, params: Dict[str, Any]) -> None:
        """
        统一处理视频生成的音频参数
        """
        audios = []
        if params.get('audio_inputs') and isinstance(params['audio_inputs'], list):
            audios.extend(params['audio_inputs'])

        if audios:
            # 标准化字段
            params['audio_inputs'] = audios

    def _normalize_video_image_params(self, params: Dict[str, Any]) -> None:
        """
        统一处理视频生成的图片参数

        统一标准输出：
        - image_inputs: 图片输入数组（首帧/尾帧/参考图）
        """
        images = []

        # 只接受统一输入字段 image_inputs
        if params.get('image_inputs') and isinstance(params['image_inputs'], list):
            images.extend(params['image_inputs'])

        # 仅保留统一字段
        if images:
            params['image_inputs'] = images

    async def _poll_provider_task(self, service: Any, provider_task_id: str, task_type: str, max_attempts: int = 120, interval: int = 5, user_id: Optional[str] = None) -> Any:
        """轮询 Provider 的异步任务直到完成

        Args:
            service: Provider 服务实例
            provider_task_id: Provider 返回的任务 ID
            task_type: 任务类型
            max_attempts: 最大轮询次数（默认 120 次 = 10 分钟）
            interval: 轮询间隔（秒，默认 5 秒）
            user_id: 用户 ID，用于获取用户的 API key

        Returns:
            ServiceResponse: 完成后的响应

        Raises:
            GenerationFailedException: 任务失败或超时
        """
        attempts = 0

        while attempts < max_attempts:
            attempts += 1

            try:
                # 调用 Provider 的 check_status 方法
                response = await service.check_status(provider_task_id, user_id=user_id)

                # 检查状态
                if hasattr(response, 'status'):
                    status = response.status

                    if status == TaskStatus.SUCCEEDED:
                        logger.info(f"Provider task {provider_task_id} completed successfully after {attempts} attempts")
                        return response

                    elif status == TaskStatus.FAILED:
                        error_msg = getattr(response, 'error', 'Unknown error')
                        logger.error(f"Provider task {provider_task_id} failed: {error_msg}")
                        raise GenerationFailedException(f"Provider task failed: {error_msg}")

                    elif status in [TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.QUEUED]:
                        # 继续轮询
                        if attempts % 6 == 0:  # 每 30 秒记录一次
                            logger.info(f"Provider task {provider_task_id} still {status.value}, attempt {attempts}/{max_attempts}")
                        await asyncio.sleep(interval)
                        continue

                    else:
                        # 未知状态，继续轮询
                        logger.warning(f"Provider task {provider_task_id} has unknown status: {status}")
                        await asyncio.sleep(interval)
                        continue

                else:
                    # 响应没有 status 字段，可能是错误
                    logger.error(f"Provider task {provider_task_id} response has no status: {response}")
                    raise GenerationFailedException("Invalid response from provider")

            except GenerationFailedException:
                # 重新抛出业务异常
                raise

            except Exception as e:
                # 网络错误等，继续重试
                logger.warning(f"Error polling provider task {provider_task_id} (attempt {attempts}): {e}")
                await asyncio.sleep(interval)
                continue

        # 超时
        raise GenerationFailedException(f"Provider task {provider_task_id} timed out after {max_attempts * interval} seconds")

    def _extract_result(self, response: Any, task_type: str) -> Dict[str, Any]:
        """从响应中提取结果，转换为前端期望的格式"""
        # 如果是 ServiceResponse 对象，先检查错误状态
        if hasattr(response, 'status') and hasattr(response, 'error'):
            # 检查是否有错误
            error_msg = getattr(response, 'error', None)
            status = getattr(response, 'status', None)
            if error_msg or status == TaskStatus.FAILED:
                return {
                    "error": error_msg or "Generation failed",
                    "status": status.value if status else "failed",
                    "task_id": getattr(response, 'task_id', None),
                    "meta": getattr(response, 'meta', None)
                }

        # 如果是 ServiceResponse 对象且有结果
        if hasattr(response, 'results') and response.results:
            # 提取 URLs
            urls = [r.url for r in response.results if hasattr(r, 'url') and r.url]

            # 构建前端期望的格式
            result = {
                "urls": urls,
                "task_id": getattr(response, 'task_id', None),
                "meta": getattr(response, 'meta', None)
            }

            # 添加额外信息（如果有）
            if response.results and len(response.results) > 0:
                first_result = response.results[0]
                if hasattr(first_result, 'usage') and first_result.usage:
                    result["usage"] = first_result.usage
                if hasattr(first_result, 'actual_prompt'):
                    result["actual_prompt"] = first_result.actual_prompt

            return result

        # 如果已经是字典格式
        if isinstance(response, dict):
            # 检查是否已经有 urls 字段
            if 'urls' in response:
                return response

            # 检查是否有 results 字段（旧格式）
            if 'results' in response:
                urls = []
                for r in response['results']:
                    if isinstance(r, dict) and 'url' in r:
                        urls.append(r['url'])
                    elif hasattr(r, 'url'):
                        urls.append(r.url)

                return {
                    "urls": urls,
                    **{k: v for k, v in response.items() if k != 'results'}
                }

            # 直接返回字典
            return response

        # 其他情况，包装为字典
        return {"raw_response": str(response)}

    async def _save_task(self, task_db: TaskDB):
        """保存任务到数据库"""
        task_db.updated_at = datetime.now().timestamp()
        with Session(engine) as session:
            session.add(task_db)
            session.commit()
            session.refresh(task_db)

    def _db_to_task(self, task_db: TaskDB) -> Task:
        """将 TaskDB 转换为 Task"""
        return Task(
            id=task_db.id,
            type=task_db.type,
            status=task_db.status,
            created_at=task_db.created_at,
            updated_at=task_db.updated_at,
            model=task_db.model,
            provider=task_db.provider,
            params=task_db.params,
            provider_task_id=task_db.provider_task_id,
            result=task_db.result,
            error=task_db.error,
            retry_count=task_db.retry_count,
            max_retries=task_db.max_retries,
            started_at=task_db.started_at,
            completed_at=task_db.completed_at,
            user_id=task_db.user_id,
            project_id=task_db.project_id
        )


# 全局实例 - 统一任务管理器
task_manager = UnifiedTaskManager()
