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_following,
    fetch_user_rankings,
    is_token_expired,
    refresh_access_token,
)

from tortoise.functions import Max

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
FRIEND_SAMPLE_SIZE = 5  # friends to randomly refresh per profile per night


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_friend_sample(profile: BeliProfile, access_token: str) -> int:
    """Fetch rankings for a random sample of followed users and upsert them."""
    try:
        following = await fetch_following(access_token, profile.beli_user_id)
    except Exception as e:
        logger.debug("Beli following fetch skipped for profile %s: %s", profile.id, e)
        return 0

    if not following:
        return 0

    # Pick the FRIEND_SAMPLE_SIZE stalest friends — never-fetched first, then oldest
    fetched_rows = (
        await BeliRestaurant.filter(profile=profile)
        .exclude(friend_beli_user_id="")
        .group_by("friend_beli_user_id")
        .annotate(last_fetched=Max("fetched_at"))
        .values("friend_beli_user_id", "last_fetched")
    )
    last_fetched: dict[str, datetime] = {
        r["friend_beli_user_id"]: r["last_fetched"] for r in fetched_rows
    }
    _epoch = datetime.fromtimestamp(0, tz=timezone.utc)
    sample = sorted(
        following,
        key=lambda f: last_fetched.get(str(f.get("id") or ""), _epoch),
    )[:FRIEND_SAMPLE_SIZE]
    upserted = 0

    for friend in sample:
        friend_id = str(friend.get("id") or "")
        friend_name = friend.get("full_name") or friend.get("username") or friend_id
        if not friend_id:
            continue

        for category in RANKING_CATEGORIES:
            await asyncio.sleep(REQUEST_SPACING)
            try:
                results = await fetch_user_rankings(access_token, friend_id, category)
            except Exception as e:
                # 405 = private profile; log quietly and skip
                logger.debug(
                    "Beli friend ranking skipped (%s/%s): %s", friend_id, category, e
                )
                continue

            for idx, item in enumerate(results):
                biz = extract_business(item, rank_index=idx + 1)
                if not biz:
                    continue
                await BeliRestaurant.update_or_create(
                    profile=profile,
                    beli_business_id=biz["beli_business_id"],
                    status=BeliStatus.BEEN,
                    friend_beli_user_id=friend_id,
                    beli_category=category,
                    defaults={
                        "friend_name": friend_name,
                        "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": None,
                        "tags": biz["tags"],
                    },
                )
                upserted += 1

    if upserted:
        logger.info(
            "Beli profile %s: %d friend restaurants upserted from %d friends",
            profile.id,
            upserted,
            len(sample),
        )
    return upserted


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 idx, item in enumerate(results):
                biz = extract_business(item, rank_index=idx + 1)
                if not biz:
                    continue
                await BeliRestaurant.update_or_create(
                    profile=profile,
                    beli_business_id=biz["beli_business_id"],
                    status=BeliStatus.BEEN,
                    friend_beli_user_id="",
                    beli_category=category,
                    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) — non-fatal if unavailable
        await asyncio.sleep(REQUEST_SPACING)
        try:
            bookmarks = await fetch_bookmarks(access_token)
        except Exception as e:
            logger.warning(
                "Beli bookmarks fetch failed for profile %s: %s", profile_id, e
            )
            bookmarks = []
        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,
                friend_beli_user_id="",
                beli_category="",
                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)
    profile = await BeliProfile.get_or_none(id=profile_id)
    if not profile:
        return
    await refresh_profile(profile_id)
    # Friend sampling only runs during nightly scheduled refreshes, not manual ones
    try:
        access_token = await _get_access_token(profile)
        await _refresh_friend_sample(profile, access_token)
    except Exception as e:
        logger.warning("Friend sample failed for profile %s: %s", profile_id, e)


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)}