"""
Image Generation Endpoints - 图片生成模块

包含图片生成相关的端点:
- POST /generations/image - 图片生成
- POST /image/upscale - 图片放大
"""
import logging

from fastapi import APIRouter, Depends

from src.models.schemas import (
    ResponseModel,
    ImageGenerationRequest,
    UpscaleImageRequest, UpscaleImageResponse
)
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 .helpers import resolve_service, ensure_url, check_user_api_key
from src.utils.errors import ErrorCode, AppException

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


@router.post("/generations/image", response_model=ResponseModel)
async def generate_image(
    request: ImageGenerationRequest,
    current_user: UserAuth = Depends(get_current_user)
):
    """ 泛型 Image Generation Endpoint.
    Supports generation for any context (Storyboard, Asset, etc.)
    """
    # Logging
    logger.info(f"Image generation request: model={request.model}, prompt={request.prompt[:50]}...")

    # 1. Resolve Model - 使用统一的路由逻辑
    image_service = resolve_service(request.model, ModelType.IMAGE)

    if not image_service:
        raise AppException(
                message="Image service not configured",
                code=ErrorCode.INTERNAL_ERROR,
                status_code=500
            )

    # 1.5 检查用户是否配置了 API Key
    # 从 image_service 获取 provider_id
    provider = getattr(image_service, 'provider_id', None)
    if provider:
        check_user_api_key(current_user.id, provider)
    
    # 2. Resolve size from aspect_ratio and resolution (BEFORE creating task)
    aspect_ratio = request.aspect_ratio
    size = None
    
    # 调试 logging
    logger.info(f"[ImageGeneration] Request received:")
    logger.info(f"  - Prompt: {request.prompt[:100]}...")
    logger.info(f"  - Model: {request.model}")
    logger.info(f"  - Aspect Ratio: {aspect_ratio}")
    logger.info(f"  - Resolution: {request.resolution}")
    logger.info(f"  - Image inputs count: {len(request.image_inputs or [])}")
    if request.image_inputs:
        for i, url in enumerate(request.image_inputs):
            logger.info(f"  - Image input {i}: {url[:100]}...")
    
    if aspect_ratio:
        # Look up resolution in model config
        model_config = getattr(image_service, "config", {})
        resolutions_config = model_config.get("resolutions", {})
        
        # Use provided resolution level (e.g. "1K", "2K", "4K") or default to "1K"
        res_level = request.resolution or "1K"
        
        if resolutions_config and res_level in resolutions_config and isinstance(resolutions_config[res_level], dict):
            ratio_map = resolutions_config[res_level]
            if aspect_ratio in ratio_map:
                size = ratio_map[aspect_ratio]
                logger.info(f"  - Resolved size from config: {size} (resolution={res_level}, ratio={aspect_ratio})")

        if not size:
            defaults = {
                "1K": {
                    "16:9": "1280*720",
                    "9:16": "720*1280",
                    "1:1": "1024*1024",
                    "4:3": "1280*960",
                    "3:4": "960*1280",
                    "2.35:1": "1280*544"
                },
                "2K": {
                    "16:9": "2560*1440",
                    "9:16": "1440*2560",
                    "1:1": "2048*2048",
                    "4:3": "2560*1920",
                    "3:4": "1920*2560"
                },
                "4K": {
                    "16:9": "3840*2160",
                    "9:16": "2160*3840",
                    "1:1": "4096*4096",
                    "4:3": "3840*2880",
                    "3:4": "2880*3840"
                }
            }
            size = defaults.get(res_level, {}).get(aspect_ratio)
            if size:
                logger.info(f"  - Resolved size from defaults: {size} (resolution={res_level}, ratio={aspect_ratio})")
    
    # 3. Create Task in DB with resolved params
    try:
        model_name = request.model or getattr(image_service, "model_id", "default_image")
        
        # 构建 params dict with resolved size
        task_params = request.model_dump()
        task_params["size"] = size  # Override with resolved size (provider expects 'size')
        
        # lora_strength 已经在 extra_params 中（前端传递），不需要额外处理
        # TaskManager 会从 extra_params 中提取并合并到一级参数
        
        # 清理 Image URLs (Convert relative paths, localhost URLs, Base64)
        logger.info(f"[ImageGeneration] Before URL sanitization - image_inputs: {task_params.get('image_inputs')}")
        
        for key in ["image_inputs"]:
            if key in task_params and isinstance(task_params[key], list):
                original_urls = task_params[key]
                task_params[key] = [ensure_url(url) for url in task_params[key] if url]
                logger.info(f"[ImageGeneration] Sanitized {key}: {original_urls} -> {task_params[key]}")
        
        # Use unified task manager
        # Get user_id from authenticated user
        user_id = current_user.id

        # Handle project_id from multiple sources
        # Priority 1: request.project_id (from request body, alias="projectId")
        # Priority 2: request.source_id if source=='project' (from extra_params, alias="sourceId")
        project_id = request.project_id
        if not project_id and request.source == 'project':
            project_id = request.source_id

        task = await task_manager.create_task(
            task_type="image",
            model=model_name,
            params=task_params,
            user_id=user_id,
            project_id=project_id
        )
    except Exception as e:
        logger.error(f"Failed to create task: {e}")
        raise AppException(
            message=f"Failed to create task: {e}",
            code=ErrorCode.INTERNAL_ERROR,
            status_code=500
        )
    
    # 日志 context if provided
    if request.source and request.source_id:
        logger.info(f"Image generation started for {request.source}:{request.source_id}, task_id: {task.id}")
    
    # Unified task manager handles execution automatically
    return ResponseModel(data={
        "task_id": task.id, 
        "status": task.status
    })


@router.post("/image/upscale", response_model=ResponseModel)
async def upscale_image(request: UpscaleImageRequest):
    """Upscale image resolution"""
    upscale_service = ModelRegistry.get_default(ModelType.UPSCALE)
    if not upscale_service:
        raise AppException(
                message="Default Upscale service not configured",
                code=ErrorCode.INTERNAL_ERROR,
                status_code=500
            )
        
    try:
        # Prepare kwargs
        upscale_kwargs = {
            "image_url": request.image_url,
            "rate": request.rate
        }
        if request.extra_params:
            upscale_kwargs.update(request.extra_params)
            
        upscaled_url = await upscale_service.upscale_image(**upscale_kwargs)
        
        if not upscaled_url:
             raise AppException(
                message="Upscaling failed or returned no URL",
                code=ErrorCode.INTERNAL_ERROR,
                status_code=500
            )
             
        return ResponseModel(data=UpscaleImageResponse(
            original_url=request.image_url,
            upscaled_url=upscaled_url
        ))
    except Exception as e:
        raise AppException(
                message=str(e),
                code=ErrorCode.INTERNAL_ERROR,
                status_code=500
            )
