zkxw/backend/app/services/cn_holiday_sync.py

"""中国节假日日历同步（ailcc batch）与进程内后台循环"""
from __future__ import annotations

import calendar
import json
import logging
import threading
import time
import uuid
from datetime import date, datetime, timedelta
from typing import Any, Callable, Dict, List, Optional, Tuple
from urllib.parse import urlencode
from urllib.request import Request, urlopen

from app.core.config import settings
from app.models.cn_holiday import SysCnHolidayDay

logger = logging.getLogger("sqxp.cn_holiday_sync")

_sync_lock = threading.Lock()
_stop_event: Optional[threading.Event] = None
_thread: Optional[threading.Thread] = None
_loop_lock = threading.Lock()


def _is_rest_from_type(day_type: int) -> int:
    return 1 if int(day_type) in (1, 2, 3) else 0


def map_ailcc_day_row(
    date_str: str,
    holiday: Optional[Dict[str, Any]],
    type_info: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
    """Normalize one day for upsert fields."""
    d = datetime.strptime(date_str, "%Y-%m-%d").date()
    if type_info and type_info.get("type") is not None:
        day_type = int(type_info["type"])
        day_name = type_info.get("name") or (holiday or {}).get("name")
        week = type_info.get("week")
        if week is None:
            week = d.isoweekday()
        else:
            week = int(week)
    else:
        # Infer from weekday: Sat/Sun → weekend(1), else workday(0)
        is_weekend = d.weekday() >= 5
        day_type = 1 if is_weekend else 0
        day_name = (holiday or {}).get("name") if holiday else None
        week = d.isoweekday()

    raw = {"date": date_str, "holiday": holiday, "type": type_info}
    return {
        "day_date": d,
        "day_type": day_type,
        "day_name": day_name,
        "week": week,
        "is_rest_day": _is_rest_from_type(day_type),
        "raw_json": json.dumps(raw, ensure_ascii=False),
    }


def iter_month_dates(year: int, month: int) -> List[str]:
    _, last = calendar.monthrange(year, month)
    return [date(year, month, day).isoformat() for day in range(1, last + 1)]


def _years_to_sync(today: Optional[date] = None) -> List[int]:
    today = today or date.today()
    return [today.year, today.year + 1]


def _months_for_year(year: int) -> List[Tuple[int, int]]:
    return [(year, m) for m in range(1, 13)]


def _fetch_batch(date_strs: List[str]) -> Dict[str, Any]:
    """GET batch API with retries. Raises on persistent failure."""
    base = (settings.CN_HOLIDAY_API_BASE or "").rstrip("/")
    params = urlencode({"d": ",".join(date_strs), "type": "Y"})
    url = f"{base}/api/holiday/batch?{params}"
    timeout = int(settings.CN_HOLIDAY_HTTP_TIMEOUT_SEC or 15)
    retries = max(0, int(settings.CN_HOLIDAY_HTTP_RETRIES or 0))
    last_err: Optional[BaseException] = None
    for attempt in range(retries + 1):
        try:
            req = Request(url, method="GET", headers={"Accept": "application/json"})
            with urlopen(req, timeout=timeout) as resp:
                body = resp.read().decode("utf-8")
            data = json.loads(body)
            if not isinstance(data, dict) or int(data.get("code", -1)) != 0:
                raise RuntimeError(f"ailcc batch unexpected response: {body[:500]}")
            return data
        except Exception as exc:  # noqa: BLE001 — retry then raise
            last_err = exc
            logger.warning(
                "cn_holiday batch fetch failed attempt=%s/%s: %s",
                attempt + 1,
                retries + 1,
                exc,
            )
            if attempt < retries:
                time.sleep(0.5 * (attempt + 1))
    assert last_err is not None
    raise last_err


def _days_in_year(year: int) -> int:
    return 366 if calendar.isleap(year) else 365


def is_year_calendar_complete(session: Any, year: int) -> bool:
    """True when the year already has one row per calendar day."""
    start = date(year, 1, 1)
    end = date(year, 12, 31)
    cnt = (
        session.query(SysCnHolidayDay)
        .filter(SysCnHolidayDay.day_date >= start, SysCnHolidayDay.day_date <= end)
        .count()
    )
    return int(cnt) >= _days_in_year(year)


def sync_cn_holidays(db_session_factory: Callable[[], Any]) -> int:
    """Fetch current + next year month-by-month; upsert by day_date. Return rows written.

    Skip any year that already has a full calendar (row count >= days in that year).
    """
    with _sync_lock:
        written = 0
        session = db_session_factory()
        try:
            now = datetime.now()
            for year in _years_to_sync():
                if is_year_calendar_complete(session, year):
                    logger.info("cn_holiday skip year=%s (calendar already complete)", year)
                    continue
                for y, month in _months_for_year(year):
                    date_strs = iter_month_dates(y, month)
                    payload = _fetch_batch(date_strs)
                    holiday_map = payload.get("holiday") or {}
                    type_map = payload.get("type") or {}
                    for ds in date_strs:
                        holiday = holiday_map.get(ds)
                        type_info = type_map.get(ds)
                        mapped = map_ailcc_day_row(ds, holiday, type_info)
                        existing = (
                            session.query(SysCnHolidayDay)
                            .filter(SysCnHolidayDay.day_date == mapped["day_date"])
                            .first()
                        )
                        if existing is None:
                            row = SysCnHolidayDay(
                                id=uuid.uuid4().hex,
                                day_date=mapped["day_date"],
                                day_type=mapped["day_type"],
                                day_name=mapped["day_name"],
                                week=mapped["week"],
                                is_rest_day=mapped["is_rest_day"],
                                raw_json=mapped["raw_json"],
                                sync_time=now,
                                create_time=now,
                                update_time=now,
                            )
                            session.add(row)
                        else:
                            existing.day_type = mapped["day_type"]
                            existing.day_name = mapped["day_name"]
                            existing.week = mapped["week"]
                            existing.is_rest_day = mapped["is_rest_day"]
                            existing.raw_json = mapped["raw_json"]
                            existing.sync_time = now
                            existing.update_time = now
                        written += 1
            session.commit()
            logger.info("cn_holiday sync ok rows=%s", written)
            return written
        except Exception:
            session.rollback()
            logger.exception("cn_holiday sync failed")
            raise
        finally:
            session.close()


def seconds_until_next_local_midnight(now: Optional[datetime] = None) -> float:
    now = now or datetime.now()
    # Always wait until the next calendar midnight (if already 00:00:00 → +1 day)
    nxt = now.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
    return max(0.0, (nxt - now).total_seconds())


def _loop_body(db_session_factory: Callable[[], Any], stop_event: threading.Event) -> None:
    try:
        sync_cn_holidays(db_session_factory)
    except Exception:
        logger.exception("cn_holiday initial sync failed; will retry at next midnight")
    while not stop_event.is_set():
        wait_sec = seconds_until_next_local_midnight()
        # Wait in chunks so stop responds quickly
        end_at = time.time() + wait_sec
        while not stop_event.is_set():
            remaining = end_at - time.time()
            if remaining <= 0:
                break
            stop_event.wait(timeout=min(remaining, 1.0))
        if stop_event.is_set():
            break
        try:
            sync_cn_holidays(db_session_factory)
        except Exception:
            logger.exception("cn_holiday midnight sync failed; keep old cache")


def start_cn_holiday_sync_loop(db_session_factory: Callable[[], Any]) -> None:
    """Start daemon thread if not running: sync immediately, then each local midnight."""
    global _stop_event, _thread
    with _loop_lock:
        if _thread is not None and _thread.is_alive():
            return
        _stop_event = threading.Event()
        t = threading.Thread(
            target=_loop_body,
            args=(db_session_factory, _stop_event),
            name="cn-holiday-sync",
            daemon=True,
        )
        _thread = t
        t.start()


def stop_cn_holiday_sync_loop(timeout: float = 2.0) -> None:
    """Set stop event and join the sync thread."""
    global _stop_event, _thread
    with _loop_lock:
        ev = _stop_event
        t = _thread
        if ev is not None:
            ev.set()
        if t is not None and t.is_alive():
            t.join(timeout=timeout)
        _thread = None
        _stop_event = None