Crawl/backend/app/services/scheduler.py

from __future__ import annotations

import asyncio
import logging
import os
from datetime import datetime, timedelta, timezone

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.cron import CronTrigger

from app.db import get_pool
from app.services.scraper import ScraperService
from app.services.key_pool_pusher import push_all_groups_to_pool

logger = logging.getLogger(__name__)
_scheduler = AsyncIOScheduler()
_scraper = ScraperService()
JOB_ID = "auto_scrape"
JOB_ID_POOL_PUSH = "pool_push"


async def _run_auto_scrape() -> None:
    """爬取所有已注册模型。"""
    pool = get_pool()
    async with pool.acquire() as conn:
        cfg = await conn.fetchrow("SELECT * FROM scrape_schedule WHERE id = 1")

    if not cfg or not cfg["enabled"]:
        return

    async with pool.acquire() as conn:
        rows = await conn.fetch("SELECT url FROM models ORDER BY created_at")

    urls = [r["url"] for r in rows]
    if not urls:
        logger.info("[scheduler] 没有已注册的模型，跳过")
        return

    logger.info(f"[scheduler] 开始自动爬取 {len(urls)} 个模型")
    job_id_row = await pool.fetchrow(
        "INSERT INTO scrape_jobs (urls, status) VALUES ($1, 'pending') RETURNING id",
        urls,
    )
    asyncio.create_task(_scraper.run_job(str(job_id_row["id"]), urls, pool))


def _reschedule(interval_days: int, start_hour: int) -> None:
    if _scheduler.get_job(JOB_ID):
        _scheduler.remove_job(JOB_ID)
    # start_hour 是用户输入的本地时间（UTC+8），转换为 UTC
    tz_offset = int(os.environ.get("TZ_OFFSET_HOURS", "8"))
    start_hour_utc = (start_hour - tz_offset) % 24
    now = datetime.now(timezone.utc)
    next_run = now.replace(hour=start_hour_utc, minute=0, second=0, microsecond=0)
    if next_run <= now:
        next_run += timedelta(days=1)
    _scheduler.add_job(
        _run_auto_scrape,
        trigger=IntervalTrigger(days=interval_days, start_date=next_run),
        id=JOB_ID,
        replace_existing=True,
    )
    logger.info(f"[scheduler] 已设置每 {interval_days} 天本地 {start_hour}:00 (UTC {start_hour_utc}:00) 执行，下次: {next_run}")


async def start_scheduler() -> None:
    pool = get_pool()
    async with pool.acquire() as conn:
        cfg = await conn.fetchrow("SELECT * FROM scrape_schedule WHERE id = 1")
    interval = cfg["interval_days"] if cfg else 1
    start_hour = cfg["start_hour"] if cfg else 2
    _reschedule(interval, start_hour)

    # 添加号池推送定时任务（每天凌晨 3 点推送一次）
    tz_offset = int(os.environ.get("TZ_OFFSET_HOURS", "8"))
    start_hour_utc = (3 - tz_offset) % 24
    _scheduler.add_job(
        _run_pool_push,
        trigger=CronTrigger(hour=start_hour_utc, minute=0),
        id=JOB_ID_POOL_PUSH,
        replace_existing=True,
    )

    _scheduler.start()
    logger.info(f"[scheduler] 调度器已启动（号池推送每天凌晨 3:00 本地时间 / UTC {start_hour_utc}:00）")


async def _run_pool_push() -> None:
    """定时推送号池数据"""
    pool = get_pool()
    logger.info("[scheduler] 开始推送号池数据到company_base_pool")
    await push_all_groups_to_pool(pool)


async def stop_scheduler() -> None:
    if _scheduler.running:
        _scheduler.shutdown(wait=False)


def get_scheduler() -> AsyncIOScheduler:
    return _scheduler


def reschedule(interval_days: int, start_hour: int) -> None:
    _reschedule(interval_days, start_hour)