SmartUp/backend/app/services/scheduler.py

"""APScheduler background scheduler for upstream checks."""
from __future__ import annotations

import json
import logging
from datetime import datetime, timezone

from apscheduler.executors.pool import ThreadPoolExecutor
from apscheduler.schedulers.background import BackgroundScheduler
from sqlalchemy.orm import Session

from app.database import SessionLocal
from app.models.upstream import Upstream
from app.models.snapshot import UpstreamRateSnapshot
from app.services.upstream_client import UpstreamClient, UpstreamError, build_snapshot
from app.services.snapshot_service import diff_snapshots, prune_snapshots
from app.services import webhook_service
from app.services import website_sync
from app.config import get_settings

logger = logging.getLogger(__name__)

_scheduler = BackgroundScheduler(timezone="UTC", executors={"default": ThreadPoolExecutor(max_workers=1)})


def get_scheduler() -> BackgroundScheduler:
    return _scheduler


def _check_upstream(upstream_id: int) -> None:
    """Full upstream check executed by scheduler (runs in thread)."""
    settings = get_settings()
    db: Session = SessionLocal()
    try:
        upstream = db.query(Upstream).filter(Upstream.id == upstream_id).first()
        if not upstream or not upstream.enabled:
            _remove_job(upstream_id)
            return

        auth_config = json.loads(upstream.auth_config_json or "{}")
        client = UpstreamClient(
            base_url=upstream.base_url,
            api_prefix=upstream.api_prefix,
            auth_type=upstream.auth_type,
            auth_config=auth_config,
            timeout=float(upstream.timeout_seconds),
        )

        was_unhealthy = upstream.last_status == "unhealthy"

        try:
            client.login()
            groups = client.get_available_groups(upstream.groups_endpoint)
            raw_rates = client.get_group_rates(upstream.rate_endpoint)
            snapshot = build_snapshot(
                upstream.id, upstream.base_url, upstream.api_prefix, groups, raw_rates
            )
        except Exception as exc:
            # failure path
            upstream.consecutive_failures = (upstream.consecutive_failures or 0) + 1
            upstream.last_error = str(exc)
            upstream.last_checked_at = datetime.now(timezone.utc)
            threshold = settings.unhealthy_threshold
            if upstream.consecutive_failures >= threshold and upstream.last_status != "unhealthy":
                upstream.last_status = "unhealthy"
                db.commit()
                webhook_service.send_status_event(
                    db, upstream.id, upstream.name, upstream.base_url,
                    "upstream_unhealthy", str(exc)
                )
            else:
                db.commit()
            logger.warning("upstream %s check failed: %s", upstream.name, exc)
            return

        # success path
        prev_snapshot_row = (
            db.query(UpstreamRateSnapshot)
            .filter(UpstreamRateSnapshot.upstream_id == upstream_id)
            .order_by(UpstreamRateSnapshot.captured_at.desc())
            .first()
        )
        previous = json.loads(prev_snapshot_row.snapshot_json) if prev_snapshot_row else None
        changes = diff_snapshots(previous, snapshot)

        # save new snapshot
        new_row = UpstreamRateSnapshot(
            upstream_id=upstream_id,
            snapshot_json=json.dumps(snapshot, ensure_ascii=False),
            captured_at=datetime.now(timezone.utc),
        )
        db.add(new_row)

        # update upstream status
        upstream.last_status = "healthy"
        upstream.last_checked_at = datetime.now(timezone.utc)
        upstream.last_error = None
        upstream.consecutive_failures = 0
        prune_snapshots(db, upstream_id, settings.snapshot_retention_count)
        db.commit()

        if was_unhealthy:
            webhook_service.send_status_event(
                db, upstream.id, upstream.name, upstream.base_url, "upstream_recovered"
            )

        if changes:
            webhook_service.send_rate_changed(
                db, upstream.id, upstream.name, upstream.base_url, changes
            )
            website_sync.sync_affected_bindings(db, upstream.id, changes)
            logger.info("upstream %s: %d rate change(s)", upstream.name, len(changes))
        else:
            logger.debug("upstream %s: no changes", upstream.name)

    finally:
        db.close()


def _remove_job(upstream_id: int) -> None:
    job_id = f"upstream_{upstream_id}"
    if _scheduler.get_job(job_id):
        _scheduler.remove_job(job_id)


def refresh_upstream(upstream_id: int, interval_seconds: int = 0, enabled: bool = True) -> None:
    """Add/update/remove a scheduler job for the given upstream."""
    job_id = f"upstream_{upstream_id}"
    if not enabled or interval_seconds <= 0:
        _remove_job(upstream_id)
        return
    _scheduler.add_job(
        _check_upstream,
        "interval",
        seconds=interval_seconds,
        id=job_id,
        args=[upstream_id],
        replace_existing=True,
        coalesce=True,
        max_instances=1,
    )
    logger.info("scheduler job %s set to %ds interval", job_id, interval_seconds)


def start_scheduler() -> None:
    """Start scheduler and load all enabled upstreams."""
    _scheduler.start()
    db: Session = SessionLocal()
    try:
        upstreams = db.query(Upstream).filter(Upstream.enabled == True).all()
        for u in upstreams:
            refresh_upstream(u.id, u.check_interval_seconds, u.enabled)
        logger.info("scheduler started with %d upstream job(s)", len(upstreams))
    finally:
        db.close()


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