"""
Batch Generation Endpoints - 批量生成模块

包含批量生成相关的端点:
- POST /generations/batch - 批量创建生成任务
"""
import logging
from typing import List, Dict, Any
from fastapi import APIRouter, Depends
from pydantic import BaseModel, Field, ConfigDict

from src.models.schemas import (
    ResponseModel,
    ImageGenerationRequest,
    VideoGenerationRequest,
    AudioGenerationRequest,
    MusicGenerationRequest
)
from src.services.provider.registry import ModelRegistry, ModelType
from src.services.task_manager import task_manager
from src.auth.dependencies import get_current_user, UserAuth
from src.utils.errors import ErrorCode, AppException

from .helpers import resolve_service, ensure_url, check_user_api_key

router = APIRouter()
logger = logging.getLogger(__name__)


class BatchGenerationRequest(BaseModel):
    """批量生成请求"""
    items: List[Dict[str, Any]] = Field(..., description="生成任务列表")

    model_config = ConfigDict(json_schema_extra={
        "example": {
            "items": [
                {
                    "type": "image",
                    "prompt": "A beautiful landscape",
                    "model": "wanx2.1-t2i-turbo"
                }
            ]
        }
    })


class BatchGenerationResponse(BaseModel):
    """批量生成响应"""
    task_ids: List[str] = Field(..., description="创建的任务ID列表")
    total: int = Field(..., description="总任务数")
    created: int = Field(..., description="成功创建的任务数")
    failed: int = Field(..., description="失败的任务数")
    errors: List[Dict[str, str]] = Field(default=[], description="错误信息列表")


@router.post("/generations/batch", response_model=ResponseModel)
async def generate_batch(
    request: BatchGenerationRequest,
    current_user: UserAuth = Depends(get_current_user)
):
    """批量创建生成任务
    
    支持同时创建多个图片、视频、音频、音乐生成任务
    
    Args:
        request: 批量生成请求，包含任务列表
        
    Returns:
        批量生成响应，包含创建的任务ID列表
    """
    logger.info(f"Batch generation request: {len(request.items)} items from user {current_user.id}")
    
    if not request.items:
        raise AppException(
            message="No items provided for batch generation",
            code=ErrorCode.INVALID_PARAMETER,
            status_code=400
        )
    
    if len(request.items) > 20:
        raise AppException(
            message="Batch size exceeds maximum limit of 20",
            code=ErrorCode.INVALID_PARAMETER,
            status_code=400
        )
    
    task_ids = []
    errors = []
    
    for idx, item in enumerate(request.items):
        try:
            item_type = item.get("type", "image")
            
            if item_type == "image":
                task_id = await _create_image_task(item, current_user)
            elif item_type == "video":
                task_id = await _create_video_task(item, current_user)
            elif item_type == "audio":
                task_id = await _create_audio_task(item, current_user)
            elif item_type == "music":
                task_id = await _create_music_task(item, current_user)
            else:
                raise ValueError(f"Unsupported generation type: {item_type}")
            
            task_ids.append(task_id)
            logger.info(f"Created {item_type} task {task_id} for batch item {idx}")
            
        except Exception as e:
            logger.error(f"Failed to create task for batch item {idx}: {e}")
            errors.append({
                "index": idx,
                "error": str(e)
            })
    
    result = BatchGenerationResponse(
        task_ids=task_ids,
        total=len(request.items),
        created=len(task_ids),
        failed=len(errors),
        errors=errors
    )
    
    logger.info(f"Batch generation completed: {result.created}/{result.total} tasks created")
    
    return ResponseModel(data=result)


async def _create_image_task(item: Dict[str, Any], current_user: UserAuth) -> str:
    """创建图片生成任务"""
    image_service = resolve_service(item.get("model"), ModelType.IMAGE)
    
    if not image_service:
        raise ValueError("Image service not configured")
    
    # 检查 API Key
    provider = getattr(image_service, 'provider_id', None)
    if provider:
        check_user_api_key(current_user.id, provider)
    
    # 解析尺寸
    aspect_ratio = item.get("aspect_ratio", "16:9")
    resolution = item.get("resolution", "1K")
    size = _resolve_image_size(image_service, aspect_ratio, resolution)
    
    # 构建参数
    task_params = {
        "prompt": item.get("prompt", ""),
        "size": size,
        "aspect_ratio": aspect_ratio,
        "resolution": resolution,
        "image_inputs": [_sanitize_url(url) for url in item.get("image_inputs", []) if url],
        "extra_params": item.get("extra_params", {})
    }
    
    # 处理 source 相关字段
    if item.get("source"):
        task_params["source"] = item["source"]
    if item.get("source_id"):
        task_params["source_id"] = item["source_id"]
    if item.get("project_id"):
        task_params["project_id"] = item["project_id"]
    
    model_name = item.get("model") or getattr(image_service, "model_id", "default_image")
    
    task = await task_manager.create_task(
        task_type="image",
        model=model_name,
        params=task_params,
        user_id=current_user.id,
        project_id=item.get("project_id")
    )
    
    return task.id


async def _create_video_task(item: Dict[str, Any], current_user: UserAuth) -> str:
    """创建视频生成任务"""
    video_service = resolve_service(item.get("model"), ModelType.VIDEO)
    
    if not video_service:
        raise ValueError("Video service not configured")
    
    # 检查 API Key
    provider = getattr(video_service, 'provider_id', None)
    if provider:
        check_user_api_key(current_user.id, provider)
    
    # 构建参数
    task_params = {
        "prompt": item.get("prompt", ""),
        "duration": item.get("duration", 5),
        "resolution": item.get("resolution", "720p"),
        "image_inputs": [_sanitize_url(url) for url in item.get("image_inputs", []) if url],
        "extra_params": item.get("extra_params", {})
    }
    
    if item.get("source"):
        task_params["source"] = item["source"]
    if item.get("source_id"):
        task_params["source_id"] = item["source_id"]
    if item.get("project_id"):
        task_params["project_id"] = item["project_id"]
    
    model_name = item.get("model") or getattr(video_service, "model_id", "default_video")
    
    task = await task_manager.create_task(
        task_type="video",
        model=model_name,
        params=task_params,
        user_id=current_user.id,
        project_id=item.get("project_id")
    )
    
    return task.id


async def _create_audio_task(item: Dict[str, Any], current_user: UserAuth) -> str:
    """创建音频生成任务"""
    audio_service = resolve_service(item.get("model"), ModelType.AUDIO)
    
    if not audio_service:
        raise ValueError("Audio service not configured")
    
    # 检查 API Key
    provider = getattr(audio_service, 'provider_id', None)
    if provider:
        check_user_api_key(current_user.id, provider)
    
    task_params = {
        "text": item.get("text", item.get("prompt", "")),
        "voice": item.get("voice", "default"),
        "speed": item.get("speed", 1.0),
        "extra_params": item.get("extra_params", {})
    }
    
    if item.get("source"):
        task_params["source"] = item["source"]
    if item.get("source_id"):
        task_params["source_id"] = item["source_id"]
    
    model_name = item.get("model") or getattr(audio_service, "model_id", "default_audio")
    
    task = await task_manager.create_task(
        task_type="audio",
        model=model_name,
        params=task_params,
        user_id=current_user.id,
        project_id=item.get("project_id")
    )
    
    return task.id


async def _create_music_task(item: Dict[str, Any], current_user: UserAuth) -> str:
    """创建音乐生成任务"""
    music_service = resolve_service(item.get("model"), ModelType.MUSIC)
    
    if not music_service:
        raise ValueError("Music service not configured")
    
    # 检查 API Key
    provider = getattr(music_service, 'provider_id', None)
    if provider:
        check_user_api_key(current_user.id, provider)
    
    task_params = {
        "prompt": item.get("prompt", ""),
        "duration": item.get("duration", 30),
        "extra_params": item.get("extra_params", {})
    }
    
    if item.get("source"):
        task_params["source"] = item["source"]
    if item.get("source_id"):
        task_params["source_id"] = item["source_id"]
    
    model_name = item.get("model") or getattr(music_service, "model_id", "default_music")
    
    task = await task_manager.create_task(
        task_type="music",
        model=model_name,
        params=task_params,
        user_id=current_user.id,
        project_id=item.get("project_id")
    )
    
    return task.id


def _resolve_image_size(image_service, aspect_ratio: str, resolution: str) -> str:
    """解析图片尺寸"""
    model_config = getattr(image_service, "config", {})
    resolutions_config = model_config.get("resolutions", {})
    
    if resolutions_config and resolution in resolutions_config:
        ratio_map = resolutions_config[resolution]
        if isinstance(ratio_map, dict) and aspect_ratio in ratio_map:
            return ratio_map[aspect_ratio]
    
    # 默认尺寸
    defaults = {
        "1K": {
            "16:9": "1280*720",
            "9:16": "720*1280",
            "1:1": "1024*1024"
        },
        "2K": {
            "16:9": "2560*1440",
            "9:16": "1440*2560",
            "1:1": "2048*2048"
        }
    }
    return defaults.get(resolution, {}).get(aspect_ratio, "1024*1024")


def _sanitize_url(url: str) -> str:
    """清理 URL"""
    if not url:
        return url
    return ensure_url(url)