"""API 对外开放接口的 Celery 任务。

处理视频和图片的异步生成流程：
- api_create_generation_task: 创建供应商任务（调用 Volcano Ark SDK）
- api_poll_generation_task: 轮询视频任务状态
- api_download_generation_result_task: 下载生成结果
- api_upscale_finalize_task: 超分完成后更新 API 任务

Worker 启动命令示例：
    celery -A app.tasks.celery_app worker \\
        --include=app.tasks.api_generation_tasks \\
        --queue=gen_api_create,gen_api_poll,gen_api_download \\
        --concurrency=4
"""

from __future__ import annotations

import json
import logging
import os
import time
from datetime import datetime, timedelta, timezone
from typing import Any

from app.config import settings
from app.enums.celery_queue import CeleryQueue
from app.models.api.api_generation_task import ApiGenerationTask
from app.models.api.api_key import ApiKey
from app.models.base import async_session
from sqlalchemy import select
from app.models.image_engine import ImageEngine
from app.models.video_engine import VideoEngine
from app.services.api_v3 import upscale_service, usage_log_service
from app.services.api_v3.logging_service import log_model_response, log_upscale_poll, log_upscale_poll_start, log_upscale_poll_end, log_error
from app.services.api_v3.quota_service import get_queued_video_tasks, can_start_video_task
from app.services.redis_registry_service import redis_acquire_lock
from app.services.generation.poll_schedule_service import (
    build_video_pending_poll_schedule,
    ensure_video_poll_fields,
    is_poll_not_due,
)
from app.services.video_gen import poll_task_status, submit_video_task
from app.tasks.api_celery_app import celery_app, run_async

logger = logging.getLogger("videogen")


# === 队列名称常量 ===
QUEUE_CREATE = "gen_api_create"
QUEUE_POLL = "gen_api_poll"
QUEUE_DOWNLOAD = "gen_api_download"


def _now() -> datetime:
    return datetime.now(timezone.utc)


async def _get_quota_info(db, api_key_id: str) -> tuple[float | None, float | None]:
    """获取当前配额信息。

    Returns:
        (quota_before, quota_after) - 当前余额作为 before，after 需要计算
    """
    from app.models.api.api_key import ApiKey
    key = await db.get(ApiKey, api_key_id)
    if key:
        return key.quota_used, key.quota_used
    return None, None


async def _refund_quota(db, task: ApiGenerationTask):
    """退回预扣配额。"""
    from app.models.api.api_key import ApiKey

    pre_deducted = task.credits_cost or 0.0
    if pre_deducted <= 0:
        return

    key_result = await db.execute(
        select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
    )
    key = key_result.scalar_one_or_none()
    if key:
        key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)
        task.credits_cost = 0  # 标记已退回


async def _start_next_queued_task(db, api_key_id: str):
    """检查并启动下一个排队的视频任务。

    当一个任务完成/失败时调用，检查是否有排队的任务可以启动。
    """
    from app.models.api.api_key import ApiKey

    # 加载 API Key
    key = await db.get(ApiKey, api_key_id)
    if not key:
        return

    # 检查是否可以启动新任务
    if not await can_start_video_task(key, db):
        return

    # 获取最早的排队任务
    queued_tasks = await get_queued_video_tasks(key, db, limit=1)
    if not queued_tasks:
        return

    next_task = queued_tasks[0]

    # 更新状态并启动
    next_task.status = "pending"
    next_task.pipeline_stage = "queued"
    await db.commit()

    # 入队 Celery 创建任务
    api_create_generation_task.apply_async(
        args=[next_task.id],
        queue=QUEUE_CREATE,
    )

    logger.info("Started queued API task: %s (key=%s)", next_task.id, api_key_id)


def _build_optimized_prompt(task: ApiGenerationTask) -> str:
    """构建优化后的提示词（追加参数信息）。"""
    base = (task.original_prompt or "").strip().rstrip("，,。；; \n\t")
    if not base:
        return task.original_prompt or ""

    parts = []
    if task.gen_type == "video":
        parts = [
            f"时长：{task.duration or 4}秒",
            f"画面比例：{task.aspect_ratio or '16:9'}",
            f"分辨率：{task.provider_generation_resolution or task.resolution or '480p'}",
        ]
    suffix = "，".join(parts)
    return f"{base}，{suffix}" if base and suffix else base


# === 任务 1: 创建供应商任务 ===

@celery_app.task(
    name="api.create_generation_task",
    bind=True,
    max_retries=3,
    soft_time_limit=300,
    time_limit=600,
    acks_late=True,
)
def api_create_generation_task(self, task_id: str) -> dict[str, Any]:
    """创建视频生成任务并提交到 Volcano Ark SDK。"""
    return run_async(_create_generation_task(self, task_id))


async def _create_generation_task(self, task_id: str) -> dict[str, Any]:
    lock_key = f"vg:lock:api_generation:create:{task_id}:attempt:{1}"

    token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=120)
    if not token:
        logger.warning("API create task lock not acquired: %s", task_id)
        return {"status": "lock_not_acquired", "task_id": task_id}

    try:
        async with async_session() as db:
            # 加载任务（带行锁）
            task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
            if not task or task.deleted_at:
                return {"status": "not_found", "task_id": task_id}

            if task.status not in ("pending", "generating"):
                return {"status": "skipped", "task_id": task_id, "current_status": task.status}

            # 解析引擎配置
            try:
                engine_snapshot = json.loads(task.engine_snapshot_json) if task.engine_snapshot_json else {}
            except (json.JSONDecodeError, TypeError):
                engine_snapshot = {}

            engine_id = task.engine_id or engine_snapshot.get("id")
            if not engine_id:
                # 失败：退回预扣配额
                await _refund_quota(db, task)
                task.status = "failed"
                task.error_message = "无法解析引擎配置"
                await db.commit()
                return {"status": "failed", "task_id": task_id, "error": "no_engine"}

            # 加载引擎
            engine = await db.get(VideoEngine, engine_id) or await db.get(ImageEngine, engine_id)
            if not engine:
                # 失败：退回预扣配额
                await _refund_quota(db, task)
                task.status = "failed"
                task.error_message = f"引擎 {engine_id} 不存在"
                await db.commit()
                return {"status": "failed", "task_id": task_id, "error": "engine_not_found"}

            # 设置优化提示词
            task.optimized_prompt = _build_optimized_prompt(task)
            task.pipeline_stage = "creating_provider_task"
            await db.flush()

            try:
                # 调用 Volcano Ark SDK
                provider_task_id = await submit_video_task(
                    db=db,
                    engine=engine,
                    record=task,
                    include_media_references=True,
                )

                # 更新任务状态
                task.provider_task_id = provider_task_id
                task.provider_response_json = json.dumps({"task_id": provider_task_id}, ensure_ascii=False)
                task.pipeline_stage = "waiting_remote"
                task.status = "generating"
                task.resource_generation_started_at = _now()

                # 设置轮询字段
                ensure_video_poll_fields(task)
                task.poll_started_at = _now()
                task.poll_interval_seconds = 30
                task.next_poll_at = _now() + timedelta(seconds=30)

                await db.commit()

                logger.info("API video submitted: task_id=%s provider_task_id=%s", task_id, provider_task_id)

                # 记录模型调用成功
                log_model_response(
                    engine_id=engine_id,
                    model_name=task.model_name,
                    task_id=task_id,
                    success=True,
                    result={"provider_task_id": provider_task_id},
                )

                # 入队轮询任务
                api_poll_generation_task.apply_async(
                    args=[task_id],
                    countdown=30,
                    queue=QUEUE_POLL,
                )

                return {"status": "submitted", "task_id": task_id, "provider_task_id": provider_task_id}

            except Exception as exc:
                logger.exception("API video submit failed: task_id=%s", task_id)

                # 记录模型调用失败
                log_model_response(
                    engine_id=engine_id or "unknown",
                    model_name=task.model_name,
                    task_id=task_id,
                    success=False,
                    error=str(exc)[:500],
                )

                # 失败：退回预扣配额
                pre_deducted = task.credits_cost or 0.0
                await _refund_quota(db, task)
                task.status = "failed"
                task.error_message = f"提交失败: {str(exc)[:500]}"
                task.pipeline_stage = "failed"
                await db.commit()

                # 记录失败日志
                try:
                    quota_before, _ = await _get_quota_info(db, task.api_key_id)
                    await usage_log_service.record_usage(
                        db=db,
                        api_key_id=task.api_key_id,
                        request_type="video_create",
                        model_name=task.model_name,
                        gen_type="video",
                        status="failed",
                        task_id=task.id,
                        credits_cost=pre_deducted,
                        refund_amount=pre_deducted,
                        price_action="refund",
                        error_message=str(exc)[:500],
                        error_code="submit_failed",
                        quota_before=quota_before,
                        quota_after=quota_before + pre_deducted if quota_before else None,
                    )
                    await db.commit()
                except Exception as log_exc:
                    logger.error("Failed to record usage log: %s", log_exc)

                # 失败释放并发槽位，检查排队任务
                await _start_next_queued_task(db, task.api_key_id)

                return {"status": "failed", "task_id": task_id, "error": str(exc)}

    finally:
        # 释放锁
        from app.services.redis_registry_service import redis_release_lock
        await redis_release_lock(lock_key=lock_key, token=token)


# === 任务 2: 轮询任务状态 ===

@celery_app.task(
    name="api.poll_generation_task",
    bind=True,
    max_retries=3,
    soft_time_limit=120,
    time_limit=300,
    acks_late=True,
)
def api_poll_generation_task(self, task_id: str) -> dict[str, Any]:
    """轮询视频任务状态。"""
    return run_async(_poll_generation_task(self, task_id))


async def _poll_generation_task(self, task_id: str) -> dict[str, Any]:
    lock_key = f"vg:lock:api_generation:poll:{task_id}"

    token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=60)
    if not token:
        return {"status": "lock_not_acquired", "task_id": task_id}

    try:
        async with async_session() as db:
            task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
            if not task or task.deleted_at:
                return {"status": "not_found", "task_id": task_id}

            if task.status != "generating" or not task.provider_task_id:
                return {"status": "skipped", "task_id": task_id}

            # 检查是否到轮询时间
            if is_poll_not_due(task):
                # 重新调度
                schedule = build_video_pending_poll_schedule(task)
                task.next_poll_at = schedule.next_poll_at
                task.poll_interval_seconds = schedule.poll_interval_seconds
                await db.commit()

                api_poll_generation_task.apply_async(
                    args=[task_id],
                    countdown=schedule.delay_seconds,
                    queue=QUEUE_POLL,
                )
                return {"status": "rescheduled", "task_id": task_id, "delay": schedule.delay_seconds}

            # 检查截止时间
            if task.deadline_at and task.deadline_at <= _now():
                # 超时：退回预扣配额
                await _refund_quota(db, task)
                task.status = "failed"
                task.error_message = "任务超时（24小时）"
                task.pipeline_stage = "timeout"
                await db.commit()
                return {"status": "timeout", "task_id": task_id}

            # 解析引擎
            try:
                engine_snapshot = json.loads(task.engine_snapshot_json) if task.engine_snapshot_json else {}
            except (json.JSONDecodeError, TypeError):
                engine_snapshot = {}

            engine_id = task.engine_id or engine_snapshot.get("id")
            engine = await db.get(VideoEngine, engine_id) if engine_id else None
            if not engine:
                task.status = "failed"
                task.error_message = f"引擎 {engine_id} 不存在"
                await db.commit()
                return {"status": "failed", "task_id": task_id, "error": "engine_not_found"}

            # 轮询状态
            task.last_poll_at = _now()
            task.poll_count = (task.poll_count or 0) + 1
            await db.flush()

            try:
                poll_result = await poll_task_status(engine, task.provider_task_id)
            except Exception as exc:
                logger.warning("API poll failed: task_id=%s error=%s", task_id, exc)
                task.poll_error_count = (task.poll_error_count or 0) + 1
                await db.commit()

                # 重新调度
                schedule = build_video_pending_poll_schedule(task)
                task.next_poll_at = schedule.next_poll_at
                task.poll_interval_seconds = schedule.poll_interval_seconds
                await db.commit()

                api_poll_generation_task.apply_async(
                    args=[task_id],
                    countdown=schedule.delay_seconds,
                    queue=QUEUE_POLL,
                )
                return {"status": "poll_error", "task_id": task_id}

            status = poll_result.get("status")

            if status == "succeeded":
                # 成功：入队下载
                task.remote_result_url = poll_result.get("video_url")
                task.provider_response_json = poll_result.get("response_data", "")
                task.pipeline_stage = "result_ready"
                task.video_tokens_used = poll_result.get("video_tokens", 0)
                await db.commit()

                api_download_generation_result_task.apply_async(
                    args=[task_id],
                    queue=QUEUE_DOWNLOAD,
                )
                return {"status": "succeeded", "task_id": task_id}

            elif status == "failed":
                # 失败：退回预扣配额
                await _refund_quota(db, task)
                task.status = "failed"
                task.error_message = poll_result.get("error", "视频生成失败")
                task.pipeline_stage = "failed"
                task.provider_response_json = poll_result.get("response_data", "")
                await db.commit()

                # 获取预扣金额（退回前）
                pre_deducted = task.credits_cost or 0.0
                await usage_log_service.record_usage(
                    db=db,
                    api_key_id=task.api_key_id,
                    request_type="video_create",
                    model_name=task.model_name,
                    gen_type="video",
                    status="failed",
                    task_id=task.id,
                    credits_cost=pre_deducted,
                    refund_amount=pre_deducted,
                    price_action="refund",
                    error_message=task.error_message,
                    error_code="generation_failed",
                )
                await db.commit()

                # 失败释放并发槽位，检查排队任务
                await _start_next_queued_task(db, task.api_key_id)

                return {"status": "failed", "task_id": task_id}

            else:
                # 仍在处理中：重新调度
                schedule = build_video_pending_poll_schedule(task)
                task.next_poll_at = schedule.next_poll_at
                task.poll_interval_seconds = schedule.poll_interval_seconds
                await db.commit()

                api_poll_generation_task.apply_async(
                    args=[task_id],
                    countdown=schedule.delay_seconds,
                    queue=QUEUE_POLL,
                )
                return {"status": "pending", "task_id": task_id, "delay": schedule.delay_seconds}

    finally:
        from app.services.redis_registry_service import redis_release_lock
        await redis_release_lock(lock_key=lock_key, token=token)


# === 任务 3: 下载生成结果 ===

@celery_app.task(
    name="api.download_generation_result_task",
    bind=True,
    max_retries=3,
    soft_time_limit=600,
    time_limit=900,
    acks_late=True,
)
def api_download_generation_result_task(self, task_id: str) -> dict[str, Any]:
    """下载视频结果并触发超分（如启用）。"""
    return run_async(_download_generation_result(self, task_id))


async def _download_generation_result(self, task_id: str) -> dict[str, Any]:
    lock_key = f"vg:lock:api_generation:download:{task_id}"

    token = await redis_acquire_lock(lock_key=lock_key, ttl_seconds=300)
    if not token:
        return {"status": "lock_not_acquired", "task_id": task_id}

    try:

        async with async_session() as db:
            task = await db.get(ApiGenerationTask, task_id, with_for_update=True)
            if not task or task.deleted_at:
                return {"status": "not_found", "task_id": task_id}

            if not task.remote_result_url:
                task.status = "failed"
                task.error_message = "无远程结果URL"
                await db.commit()
                return {"status": "failed", "task_id": task_id}

            # 检查是否需要超分
            use_upscale = bool(
                task.gen_type == "video"
                and task.video_upscale_enabled_snapshot
                and task.video_upscale_snapshot_json
            )

            if use_upscale:
                # 下载视频到 upscaled 目录（作为超分源）
                import os
                from app.services.video_gen import download_video

                date_dir = datetime.now().strftime("%Y%m%d")
                # 源文件存储路径
                dest_path = f"./storage/generate/api/upscaled/{date_dir}/{task.id}_source.mp4"
                abs_dest_path = os.path.abspath(dest_path)
                os.makedirs(os.path.dirname(abs_dest_path), exist_ok=True)

                logger.info("Downloading source video to: %s", abs_dest_path)
                try:
                    await download_video(task.remote_result_url, abs_dest_path)
                    # 注意：超分时不设置 video_url，等超分完成后再设置
                    task.local_path = dest_path  # 源文件路径
                    task.download_storage_date_dir = date_dir
                    logger.info("Source video downloaded successfully: %s", dest_path)

                    # 获取视频信息（尺寸、时长、文件大小）
                    try:
                        abs_path = os.path.abspath(dest_path)
                        if os.path.exists(abs_path):
                            task.source_file_size_bytes = os.path.getsize(abs_path)
                            # 使用 probe_video 探测实际视频尺寸和时长
                            from app.services.video_upscale.media_service import probe_video
                            source_info = await probe_video(abs_path)
                            if source_info:
                                task.source_width = source_info.width
                                task.source_height = source_info.height
                                task.source_duration_seconds = round(source_info.duration_seconds, 2)
                    except Exception as exc:
                        logger.warning("Failed to probe video info: %s", exc)

                    await db.flush()

                    # 创建超分任务（与状态更新在同一事务中）
                    try:
                        upscale_task = await upscale_service.prepare_api_upscale_task(
                            db=db,
                            api_task=task,
                            source_local_path=dest_path,
                        )

                        # 如果已存在超分任务（重复调用），检查超分状态
                        if upscale_task is None:
                            logger.info("Upscale task already exists for %s, checking status", task_id)
                            # 重新查询超分任务状态
                            from app.models.video_upscale_task import VideoUpscaleTask
                            from sqlalchemy import select
                            upscale_result = await db.execute(
                                select(VideoUpscaleTask).where(
                                    VideoUpscaleTask.api_generation_task_id == task.id
                                ).limit(1)
                            )
                            existing_upscale = upscale_result.scalar_one_or_none()

                            if existing_upscale and existing_upscale.status == "completed":
                                # 超分已完成
                                task.video_url = existing_upscale.final_resource_url or task.remote_result_url
                                task.status = "completed"
                                task.pipeline_stage = "done"
                                task.generated_at = _now()
                            else:
                                # 超分仍在进行中，保持 generating 状态
                                task.status = "generating"
                                task.pipeline_stage = "upscale_processing"
                            await db.commit()
                            return {"status": task.status, "task_id": task_id, "note": "upscale_already_exists"}

                        # 记录超分开始
                        log_upscale_poll_start(task_id=upscale_task.id, api_task_id=task_id)

                        # 入队超分任务（使用简化版 API v3 专用任务）
                        from app.tasks.api_upscale_tasks import api_upscale_execute_local_simple, api_upscale_submit_remote_simple

                        processor_key = upscale_task.processor_key

                        if processor_key in ("local_ffmpeg_crop_v1",):
                            api_upscale_execute_local_simple.apply_async(
                                args=[upscale_task.id],
                                queue="gen_api_upscale",
                            )
                        else:
                            # 远程超分（火山 MediaKit）
                            api_upscale_submit_remote_simple.apply_async(
                                args=[upscale_task.id],
                                queue="gen_api_upscale",
                            )
                        await db.commit()

                        return {"status": "upscale_queued", "task_id": task_id, "upscale_task_id": upscale_task.id}

                    except Exception as upscale_exc:
                        # 超分创建失败：记录错误，但视频已下载成功
                        # 将任务标记为 completed（有视频但无超分）
                        log_error(
                            "UPSCALE_CREATE_ERROR",
                            f"超分任务创建失败: {str(upscale_exc)[:500]}",
                            {"task_id": task_id, "video_url": task.remote_result_url}
                        )
                        task.video_url = task.remote_result_url
                        task.status = "completed"
                        task.pipeline_stage = "done"
                        task.generated_at = _now()
                        task.error_message = f"超分创建失败，返回原始视频: {str(upscale_exc)[:200]}"
                        await db.commit()

                        # 检查并启动下一个排队任务
                        await _start_next_queued_task(db, task.api_key_id)

                        return {"status": "completed_without_upscale", "task_id": task_id, "error": str(upscale_exc)[:500]}

                except Exception as exc:
                    # 下载失败：退回预扣配额
                    from app.models.api.api_key import ApiKey
                    pre_deducted = task.credits_cost or 0.0
                    if pre_deducted > 0:
                        key_result = await db.execute(
                            select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
                        )
                        key = key_result.scalar_one_or_none()
                        if key:
                            key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)

                    log_error(
                        "DOWNLOAD_ERROR",
                        f"视频下载失败: {str(exc)[:500]}",
                        {"task_id": task_id}
                    )
                    logger.exception("API video download failed: task_id=%s", task_id)
                    task.status = "failed"
                    task.error_message = f"下载失败: {str(exc)[:500]}"
                    task.credits_cost = 0
                    await db.commit()

                    # 记录失败日志
                    await usage_log_service.record_usage(
                        db=db,
                        api_key_id=task.api_key_id,
                        request_type="video_create",
                        model_name=task.model_name,
                        gen_type="video",
                        status="failed",
                        task_id=task.id,
                        credits_cost=pre_deducted,
                        refund_amount=pre_deducted,
                        price_action="refund",
                        resolution=task.resolution,
                        duration=task.duration,
                        error_message=str(exc)[:500],
                        error_code="download_failed",
                    )

                    # 失败释放并发槽位，检查排队任务
                    await _start_next_queued_task(db, task.api_key_id)

                    return {"status": "failed", "task_id": task_id, "error": str(exc)}

            else:
                # 直接下载最终结果
                import os
                from app.services.video_gen import download_video

                date_dir = datetime.now().strftime("%Y%m%d")
                # 统一路径格式
                dest_path = f"./storage/generate/api/videos/{date_dir}/{task.id}.mp4"
                abs_dest_path = os.path.abspath(dest_path)
                os.makedirs(os.path.dirname(abs_dest_path), exist_ok=True)

                logger.info("Downloading final video to: %s", abs_dest_path)
                try:
                    await download_video(task.remote_result_url, abs_dest_path)
                    logger.info("Final video downloaded successfully: %s", dest_path)

                    # 配额已在创建时预扣，此处不再重复扣减
                    task.video_url = dest_path  # 使用相对路径
                    task.local_path = dest_path
                    task.download_storage_date_dir = date_dir
                    task.status = "completed"
                    task.pipeline_stage = "done"
                    task.generated_at = _now()
                    await db.commit()

                    # 记录成功日志（配额已在创建时预扣）
                    try:
                        await usage_log_service.record_usage(
                            db=db,
                            api_key_id=task.api_key_id,
                            request_type="video_create",
                            model_name=task.model_name,
                            gen_type="video",
                            status="success",
                            task_id=task.id,
                            credits_cost=task.credits_cost,
                            tokens_used=task.video_tokens_used,
                            price_action="deduct",
                            resolution=task.resolution,
                            duration=task.duration,
                        )
                        await db.commit()
                    except Exception as log_exc:
                        logger.error("Failed to record usage log: %s", log_exc)
                        # 日志记录失败不应影响任务完成

                    # 检查并启动下一个排队任务
                    await _start_next_queued_task(db, task.api_key_id)

                    return {"status": "completed", "task_id": task_id}

                except Exception as exc:
                    # 下载失败：退回预扣配额
                    from app.models.api.api_key import ApiKey
                    pre_deducted = task.credits_cost or 0.0
                    if pre_deducted > 0:
                        key_result = await db.execute(
                            select(ApiKey).where(ApiKey.id == task.api_key_id).with_for_update()
                        )
                        key = key_result.scalar_one_or_none()
                        if key:
                            key.quota_used = round(max(0, (key.quota_used or 0.0) - pre_deducted), 2)

                    logger.exception("API video download failed: task_id=%s", task_id)
                    task.status = "failed"
                    task.error_message = f"下载失败: {str(exc)[:500]}"
                    task.credits_cost = 0
                    await db.commit()

                    # 失败释放并发槽位，检查排队任务
                    await _start_next_queued_task(db, task.api_key_id)

                    return {"status": "failed", "task_id": task_id, "error": str(exc)}

    finally:
        from app.services.redis_registry_service import redis_release_lock
        await redis_release_lock(lock_key=lock_key, token=token)


# === 任务 4: 超分完成后更新 API 任务 ===

@celery_app.task(
    name="api.upscale_finalize_task",
    bind=True,
    max_retries=3,
    soft_time_limit=120,
    time_limit=300,
    acks_late=True,
)
def api_upscale_finalize_task(self, api_task_id: str, upscale_task_id: str) -> dict[str, Any]:
    """超分完成后更新 API 任务状态。"""
    return run_async(_upscale_finalize(self, api_task_id, upscale_task_id))


async def _upscale_finalize(self, api_task_id: str, upscale_task_id: str) -> dict[str, Any]:
    async with async_session() as db:
        from app.models.video_upscale_task import VideoUpscaleTask

        task = await db.get(ApiGenerationTask, api_task_id)
        upscale_task = await db.get(VideoUpscaleTask, upscale_task_id)

        if not task or task.deleted_at:
            return {"status": "not_found", "api_task_id": api_task_id}

        if upscale_task and upscale_task.status == "completed":
            # 超分成功：更新视频URL（配额已在创建时预扣）
            task.video_url = upscale_task.provider_output_url or task.remote_result_url
            task.status = "completed"
            task.pipeline_stage = "done"
            task.generated_at = _now()
            await db.commit()

            # 记录超分完成日志
            log_upscale_poll_end(
                task_id=upscale_task_id,
                api_task_id=api_task_id,
                success=True,
                final_status="completed",
                total_attempts=upscale_task.attempt_count or 1,
            )

            await usage_log_service.record_usage(
                db=db,
                api_key_id=task.api_key_id,
                request_type="video_create",
                model_name=task.model_name,
                gen_type="video",
                status="success",
                task_id=task.id,
                credits_cost=task.credits_cost,
                tokens_used=task.video_tokens_used,
            )
            await db.commit()

            # 检查并启动下一个排队任务
            await _start_next_queued_task(db, task.api_key_id)

            return {"status": "completed", "api_task_id": api_task_id}

        elif upscale_task and upscale_task.status == "failed":
            # 超分失败：回退到原始视频
            task.video_url = task.remote_result_url
            task.status = "completed"
            task.pipeline_stage = "done"
            task.generated_at = _now()
            task.error_message = "超分失败，返回原始视频"
            await db.commit()

            # 记录超分失败日志
            log_upscale_poll_end(
                task_id=upscale_task_id,
                api_task_id=api_task_id,
                success=False,
                final_status="failed",
                total_attempts=upscale_task.attempt_count or 1,
            )

            # 检查并启动下一个排队任务
            await _start_next_queued_task(db, task.api_key_id)

            return {"status": "completed_with_fallback", "api_task_id": api_task_id}

        else:
            # 超分仍在处理中：重新调度
            log_upscale_poll(
                task_id=upscale_task_id,
                api_task_id=api_task_id,
                status=upscale_task.status or "unknown",
                attempt=upscale_task.attempt_count or 0,
            )
            api_upscale_finalize_task.apply_async(
                args=[api_task_id, upscale_task_id],
                countdown=60,
                queue=QUEUE_DOWNLOAD,
            )
            return {"status": "waiting_upscale", "api_task_id": api_task_id}