simbarag/blueprints/beli/logic.py

import asyncio
import logging
import random
from datetime import datetime, timezone

from utils.beli_service import (
    RANKING_CATEGORIES,
    REQUEST_SPACING,
    authenticate,
    extract_business,
    fetch_bookmarks,
    fetch_user_rankings,
    is_token_expired,
    refresh_access_token,
)

from .models import BeliProfile, BeliRestaurant, BeliStatus

logger = logging.getLogger(__name__)

MAX_PROFILES_PER_USER = 3
NIGHTLY_BATCH_SIZE = 3
NIGHTLY_WINDOW_SECONDS = 4 * 3600  # spread refreshes across a 4-hour window


async def _get_access_token(profile: BeliProfile) -> str:
    """Return a live access token, refreshing or re-authenticating as needed."""
    if profile.beli_refresh_token and not is_token_expired(profile.beli_refresh_token):
        try:
            token = await refresh_access_token(profile.beli_refresh_token)
            return token
        except Exception:
            logger.warning(
                "Refresh token failed for profile %s, falling back to password",
                profile.id,
            )

    session = await authenticate(profile.beli_email, profile.beli_password)
    profile.beli_user_id = session["user_id"]
    profile.beli_refresh_token = session["refresh_token"]
    await profile.save()
    return session["access_token"]


async def refresh_profile(profile_id: str) -> dict:
    """Fetch the latest rankings and bookmarks for a profile and upsert them."""
    profile = await BeliProfile.get_or_none(id=profile_id)
    if not profile or not profile.enabled:
        return {"skipped": True}

    try:
        access_token = await _get_access_token(profile)

        upserted = 0

        # Fetch all ranking categories (restaurants the user has been to and ranked)
        for category in RANKING_CATEGORIES:
            await asyncio.sleep(REQUEST_SPACING)
            results = await fetch_user_rankings(
                access_token, profile.beli_user_id, category
            )
            for item in results:
                biz = extract_business(item)
                if not biz:
                    continue
                await BeliRestaurant.update_or_create(
                    profile=profile,
                    beli_business_id=biz["beli_business_id"],
                    status=BeliStatus.BEEN,
                    defaults={
                        "name": biz["name"],
                        "city": biz["city"],
                        "neighborhood": biz["neighborhood"],
                        "cuisine": biz["cuisine"],
                        "price_range": biz["price_range"],
                        "rank_index": biz["rank_index"],
                        "score": biz["score"],
                        "notes": biz["notes"],
                        "tags": biz["tags"],
                    },
                )
                upserted += 1

        # Fetch bookmarks (want-to-try list)
        await asyncio.sleep(REQUEST_SPACING)
        bookmarks = await fetch_bookmarks(access_token)
        for item in bookmarks:
            biz = extract_business(item)
            if not biz:
                continue
            await BeliRestaurant.update_or_create(
                profile=profile,
                beli_business_id=biz["beli_business_id"],
                status=BeliStatus.WANT_TO_TRY,
                defaults={
                    "name": biz["name"],
                    "city": biz["city"],
                    "neighborhood": biz["neighborhood"],
                    "cuisine": biz["cuisine"],
                    "price_range": biz["price_range"],
                    "rank_index": None,
                    "score": None,
                    "notes": biz["notes"],
                    "tags": biz["tags"],
                },
            )
            upserted += 1

        profile.last_refreshed_at = datetime.now(timezone.utc)
        profile.last_error = None
        await profile.save()

        logger.info(
            "Beli profile %s refreshed: %d restaurants upserted", profile_id, upserted
        )
        return {"ok": True, "upserted": upserted}

    except Exception as e:
        profile.last_error = str(e)
        await profile.save()
        logger.error("Failed to refresh Beli profile %s: %s", profile_id, e)
        return {"error": str(e)}


async def _refresh_after_delay(profile_id: str, delay: float) -> None:
    await asyncio.sleep(delay)
    await refresh_profile(profile_id)


async def schedule_nightly_refresh() -> dict:
    """Pick the NIGHTLY_BATCH_SIZE stalest profiles and schedule them with random delays."""
    profiles = (
        await BeliProfile.filter(enabled=True)
        .order_by("last_refreshed_at")
        .limit(NIGHTLY_BATCH_SIZE)
        .all()
    )

    if not profiles:
        logger.info("No enabled Beli profiles to refresh")
        return {"scheduled": 0}

    for profile in profiles:
        delay = random.uniform(0, NIGHTLY_WINDOW_SECONDS)
        asyncio.create_task(_refresh_after_delay(str(profile.id), delay))
        logger.info("Scheduled Beli refresh for profile %s in %.0fs", profile.id, delay)

    return {"scheduled": len(profiles)}