SupplyAgent/supply_infra/pipeline/dates.py

from __future__ import annotations

from datetime import datetime, timedelta
from zoneinfo import ZoneInfo

from supply_infra.config import InfraSettings, get_infra_settings

CHINA_TIMEZONE = ZoneInfo("Asia/Shanghai")


def validate_biz_dt(value: str) -> str:
    text = str(value).strip()
    parsed = datetime.strptime(text, "%Y%m%d")
    if parsed.strftime("%Y%m%d") != text:
        raise ValueError(f"Invalid biz_dt: {value!r}")
    return text


def resolve_biz_dt(
    biz_dt: str | None,
    *,
    settings: InfraSettings | None = None,
    now: datetime | None = None,
) -> str:
    if biz_dt:
        return validate_biz_dt(biz_dt)
    active = settings or get_infra_settings()
    local_now = now or datetime.now(ZoneInfo(active.scheduler_timezone))
    if local_now.tzinfo is None:
        local_now = local_now.replace(tzinfo=ZoneInfo(active.scheduler_timezone))
    return local_now.astimezone(ZoneInfo(active.scheduler_timezone)).strftime("%Y%m%d")


def build_date_snapshot(biz_dt: str) -> dict[str, str]:
    resolved = validate_biz_dt(biz_dt)
    day = datetime.strptime(resolved, "%Y%m%d")
    previous = (day - timedelta(days=1)).strftime("%Y%m%d")
    return {
        "biz_dt": resolved,
        "global_tree_partition": previous,
        "demand_pool_partition": resolved,
        "video_decode_partition": previous,
    }


def china_now() -> datetime:
    """Return a naive China Standard Time value for MySQL DATETIME columns."""
    return datetime.now(CHINA_TIMEZONE).replace(tzinfo=None)


def next_schedule_china(
    *,
    settings: InfraSettings | None = None,
    now: datetime | None = None,
) -> datetime:
    active = settings or get_infra_settings()
    zone = ZoneInfo(active.scheduler_timezone)
    local_now = now or datetime.now(zone)
    if local_now.tzinfo is None:
        local_now = local_now.replace(tzinfo=zone)
    else:
        local_now = local_now.astimezone(zone)
    candidate = local_now.replace(
        hour=active.scheduler_cron_hour,
        minute=active.scheduler_cron_minute,
        second=0,
        microsecond=0,
    )
    if candidate <= local_now:
        candidate += timedelta(days=1)
    return candidate.astimezone(CHINA_TIMEZONE).replace(tzinfo=None)


def next_daily_deadline_china(
    *,
    settings: InfraSettings | None = None,
    now: datetime | None = None,
) -> datetime:
    """Return the next local calendar day's scheduler start as China-time naive."""
    active = settings or get_infra_settings()
    zone = ZoneInfo(active.scheduler_timezone)
    local_now = now or datetime.now(zone)
    if local_now.tzinfo is None:
        local_now = local_now.replace(tzinfo=zone)
    else:
        local_now = local_now.astimezone(zone)
    candidate = (local_now + timedelta(days=1)).replace(
        hour=active.scheduler_cron_hour,
        minute=active.scheduler_cron_minute,
        second=0,
        microsecond=0,
    )
    return candidate.astimezone(CHINA_TIMEZONE).replace(tzinfo=None)


def deadline_for_trigger(
    trigger_type: str,
    *,
    settings: InfraSettings | None = None,
    now: datetime | None = None,
) -> datetime | None:
    """Scheduled daily runs have a next-day deadline; manual runs do not."""
    if trigger_type not in {"cron", "reconcile"}:
        return None
    return next_daily_deadline_china(settings=settings, now=now)