"""
Task Status Endpoints - 任务状态模块

包含任务管理相关的端点:
- GET /tasks - 列出任务
- GET /tasks/{task_id} - 获取任务状态
"""
import logging
import asyncio

from fastapi import APIRouter, Query, Request
from typing import Optional

from src.models.schemas import ResponseModel, PaginationParams
from src.utils.pagination import Paginator
from src.services.provider.registry import ModelRegistry, ModelType
from src.services.provider.base import TaskStatus
from src.services.task_manager import task_manager
from src.services.storage_service import storage_manager
from src.config.settings import OSS_BUCKET, OSS_REGION
from src.utils.errors import TaskNotFoundException, AppException, ErrorCode

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


@router.get("/tasks", response_model=ResponseModel)
async def list_tasks(
    request: Request,
    type: Optional[str] = Query(None, description="Filter by task type"),
    page: int = Query(1, ge=1, description="页码，从 1 开始"),
    page_size: int = Query(20, ge=1, le=100, description="每页数量，最大 100"),
):
    """列表 tasks with optional filtering.

    Args:
        type: 任务类型过滤
        page: 页码
        page_size: 每页数量

    Returns:
        分页的任务列表
    """
    # 计算偏移量
    offset = (page - 1) * page_size

    # 获取任务列表
    tasks = task_manager.list_tasks(type=type, limit=page_size, offset=offset)

    # 获取总数（简化处理，实际应该查询总数）
    total = len(tasks)

    # 创建分页器
    paginator = Paginator(
        items=tasks,
        total=total,
        page=page,
        page_size=page_size
    )

    return paginator.to_response(request)


@router.get("/tasks/{task_id}", response_model=ResponseModel)
async def get_task_status(task_id: str):
    """
    Check status of a task.
    """
    # 1. Check local TaskManager
    task = await task_manager.get_task(task_id)
    
    if task:
        # If task is active (pending/processing), check provider status
        active_statuses = ["pending", "processing", "queued", "running", "submitted"]
        if task.status.lower() in active_statuses and task.provider_task_id:
            service = ModelRegistry.get(task.model)
            # Fallback to default if model service not found by name
            if not service:
                if task.type == "image":
                    service = ModelRegistry.get_default(ModelType.IMAGE)
                elif task.type == "video":
                    service = ModelRegistry.get_default(ModelType.VIDEO)
                elif task.type == "audio":
                    service = ModelRegistry.get_default(ModelType.AUDIO)
                elif task.type == "music":
                    service = ModelRegistry.get_default(ModelType.MUSIC)
            
            if service and hasattr(service, 'check_status'):
                try:
                    result = await service.check_status(task.provider_task_id, user_id=task.user_id)
                    
                    # 更新 status
                    # 映射 TaskStatus enum to string
                    new_status = result.status.value.lower()
                    if new_status == "succeeded":
                        new_status = "success"
                    
                    # If succeeded, process results (OSS upload)
                    final_results = {}
                    if result.status == TaskStatus.SUCCEEDED and result.results:
                         processed_urls = []
                         for idx, item in enumerate(result.results):
                             if item.url:
                                 # Determine extension
                                 ext = "png"
                                 if task.type == "video":
                                     ext = "mp4"
                                 elif task.type in ("audio", "music"):
                                     ext = "mp3"
                                 
                                 key = f"generated/{task.type}s/{task.id}_{idx}.{ext}"
                                 
                                 # Save to storage (Local or OSS)
                                 oss_url = await asyncio.to_thread(storage_manager.save_from_url, item.url, key)
                                 processed_urls.append(oss_url or item.url)
                         
                         final_results = {"urls": processed_urls}
                    
                    # 更新 task in DB
                    update_data = {"status": new_status}
                    if final_results:
                        update_data["result"] = final_results
                    if result.error:
                        update_data["error"] = str(result.error)
                        
                    task = task_manager.update_task(task.id, **update_data)
                    
                except Exception as e:
                    logger.error(f"Failed to check provider status for task {task.id}: {e}")
                    # Don't fail the request, just return current state
        
        # Retry OSS upload if succeeded but urls are not OSS (resilience)
        elif (task.status.lower() == "succeeded" or task.status.lower() == "success") and task.result and "urls" in task.result:
            try:
                urls = task.result["urls"]
                updated_urls = []
                needs_update = False
                
                for idx, url in enumerate(urls):
                    # Check if it looks like an OSS url
                    is_oss = False
                    if OSS_BUCKET and OSS_REGION and f"{OSS_BUCKET}.{OSS_REGION}" in url:
                        is_oss = True
                    
                    if not is_oss:
                         # Attempt upload
                         ext = "png"
                         if task.type == "video":
                             ext = "mp4"
                         elif task.type in ("audio", "music"):
                             ext = "mp3"
                         key = f"generated/{task.type}s/{task.id}_{idx}.{ext}"
                         
                         new_url = await asyncio.to_thread(storage_manager.save_from_url, url, key)
                         if new_url and new_url != url:
                             updated_urls.append(new_url)
                             needs_update = True
                         else:
                             updated_urls.append(url)
                    else:
                        # It is OSS, maybe re-sign it
                        signed_url = storage_manager.sign_url(url)
                        updated_urls.append(signed_url if signed_url else url)
                
                if needs_update:
                    task = task_manager.update_task(task.id, result={"urls": updated_urls})
                    
            except Exception as e:
                logger.error(f"Failed to retry OSS upload for task {task.id}: {e}")

        return ResponseModel(data=task)

    raise AppException(
                message="Task not found",
                code=ErrorCode.NOT_FOUND,
                status_code=404
            )


@router.post("/tasks/{task_id}/cancel", response_model=ResponseModel)
async def cancel_task(task_id: str):
    """
    Cancel a pending/processing task.
    仅更新本地任务状态为 CANCELLED，不保证下游 provider 真正中止计算，
    但对于绝大多数异步生成场景已经足够（前端不再轮询，结果也不会再写入）。
    """
    try:
        cancelled = await task_manager.cancel_task(task_id)
    except TaskNotFoundException:
        raise AppException(
                message="Task not found",
                code=ErrorCode.NOT_FOUND,
                status_code=404
            )
    except Exception as e:
        logger.error(f"Failed to cancel task {task_id}: {e}")
        raise AppException(
                message="Failed to cancel task",
                code=ErrorCode.INTERNAL_ERROR,
                status_code=500
            )

    if not cancelled:
        # 已完成任务无法取消
        raise AppException(
                message="Task already finished and cannot be cancelled",
                code=ErrorCode.INVALID_PARAMETER,
                status_code=400
            )

    # 返回简单结果，前端如需最新任务详情可以再调用 GET /tasks/{task_id}
    return ResponseModel(data={"task_id": task_id, "cancelled": True})
