video-gen/video-gen-api/app/services/generation/poll_schedule_service.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from app.config import settings
from app.enums.generation_task import GenerationType
from app.services.generation.pipeline.owner_service import GenerationOwner
from app.services.redis_registry_service import ensure_aware_utc


@dataclass(slots=True)
class PollScheduleDecision:
    delay_seconds: int
    next_poll_at: datetime
    poll_interval_seconds: int
    direct_countdown: bool
    final_poll: bool
    reason: str


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


def is_video_generation_task(task: GenerationOwner) -> bool:
    return str(getattr(task, "gen_type", "") or "").lower() == GenerationType.VIDEO.value


def video_final_deadline_from(now: datetime | None = None) -> datetime:
    current_time = now or utc_now()
    hours = max(1, int(settings.CHATAPI_ASYNC_VIDEO_FINAL_DEADLINE_HOURS or 24))
    return current_time + timedelta(hours=hours)


def ensure_video_poll_fields(task: GenerationOwner, *, now: datetime | None = None) -> None:
    """补齐视频轮询调度字段，兼容历史任务。"""
    if not is_video_generation_task(task):
        return

    current_time = now or utc_now()
    if ensure_aware_utc(getattr(task, "poll_started_at", None)) is None:
        task.poll_started_at = current_time
    if ensure_aware_utc(getattr(task, "deadline_at", None)) is None:
        resource_started_at = (
            ensure_aware_utc(getattr(task, "resource_generation_started_at", None))
            or ensure_aware_utc(getattr(task, "created_at", None))
            or current_time
        )
        task.deadline_at = video_final_deadline_from(resource_started_at)
    if getattr(task, "poll_interval_seconds", None) is None:
        task.poll_interval_seconds = 0


def is_final_poll_due(task: GenerationOwner, *, now: datetime | None = None) -> bool:
    deadline_at = ensure_aware_utc(getattr(task, "deadline_at", None))
    return bool(deadline_at and deadline_at <= (now or utc_now()))


def is_poll_not_due(task: GenerationOwner, *, now: datetime | None = None, tolerance_seconds: int = 1) -> bool:
    """判断当前 poll 任务是否早于 next_poll_at。只对视频降频轮询生效。"""
    if not is_video_generation_task(task):
        return False
    if is_final_poll_due(task, now=now):
        return False
    next_poll_at = ensure_aware_utc(getattr(task, "next_poll_at", None))
    if next_poll_at is None:
        return False
    return next_poll_at > (now or utc_now()) + timedelta(seconds=max(0, int(tolerance_seconds)))


def _clamp_positive_seconds(value: int | float | None, default: int) -> int:
    try:
        parsed = int(value if value is not None else default)
    except (TypeError, ValueError):
        parsed = default
    return max(1, parsed)


def build_video_pending_poll_schedule(
    task: GenerationOwner,
    *,
    now: datetime | None = None,
) -> PollScheduleDecision:
    """计算视频任务 pending/running 后的下一次轮询时间。"""
    current_time = now or utc_now()
    ensure_video_poll_fields(task, now=current_time)

    deadline_at = ensure_aware_utc(task.deadline_at)
    if deadline_at and deadline_at <= current_time:
        return PollScheduleDecision(
            delay_seconds=0,
            next_poll_at=current_time,
            poll_interval_seconds=int(task.poll_interval_seconds or 0),
            direct_countdown=True,
            final_poll=True,
            reason="video_final_poll_due",
        )

    poll_started_at = ensure_aware_utc(task.poll_started_at) or current_time
    elapsed_seconds = max(0, int((current_time - poll_started_at).total_seconds()))
    high_freq_seconds = max(0, int(settings.CHATAPI_ASYNC_VIDEO_HIGH_FREQ_MINUTES or 10)) * 60
    high_freq_poll_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_HIGH_FREQ_POLL_SECONDS, 30)
    initial_backoff_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_INITIAL_SECONDS, 60)
    multiplier = max(1, int(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_MULTIPLIER or 2))
    max_backoff_seconds = _clamp_positive_seconds(settings.CHATAPI_ASYNC_VIDEO_BACKOFF_MAX_SECONDS, 3600)

    if elapsed_seconds < high_freq_seconds:
        delay_seconds = high_freq_poll_seconds
        interval_seconds = int(task.poll_interval_seconds or 0)
        reason = "video_high_freq_poll"
    else:
        previous_interval = int(task.poll_interval_seconds or 0)
        if previous_interval < initial_backoff_seconds:
            interval_seconds = initial_backoff_seconds
        else:
            interval_seconds = min(previous_interval * multiplier, max_backoff_seconds)
        delay_seconds = interval_seconds
        reason = "video_backoff_poll"

    next_poll_at = current_time + timedelta(seconds=delay_seconds)
    final_poll = False
    if deadline_at and next_poll_at >= deadline_at:
        next_poll_at = deadline_at
        delay_seconds = max(0, int((deadline_at - current_time).total_seconds()))
        final_poll = delay_seconds <= 0
        reason = "video_schedule_to_final_deadline"

    direct_max = max(0, int(settings.CHATAPI_ASYNC_VIDEO_DIRECT_COUNTDOWN_MAX_SECONDS or 300))
    return PollScheduleDecision(
        delay_seconds=delay_seconds,
        next_poll_at=next_poll_at,
        poll_interval_seconds=interval_seconds,
        direct_countdown=delay_seconds <= direct_max,
        final_poll=final_poll,
        reason=reason,
    )


def build_default_poll_schedule(
    task: GenerationOwner,
    *,
    now: datetime | None = None,
    delay_seconds: int | None = None,
    reason: str = "default_poll",
) -> PollScheduleDecision:
    """图片和旧逻辑兼容用的固定间隔轮询计划。"""
    current_time = now or utc_now()
    delay = _clamp_positive_seconds(delay_seconds, int(settings.CHATAPI_ASYNC_POLL_INTERVAL_SECONDS or 30))
    next_poll_at = current_time + timedelta(seconds=delay)
    return PollScheduleDecision(
        delay_seconds=delay,
        next_poll_at=next_poll_at,
        poll_interval_seconds=int(getattr(task, "poll_interval_seconds", 0) or 0),
        direct_countdown=True,
        final_poll=False,
        reason=reason,
    )