simbarag/scripts/manage_beli.py

"""
Manage Beli profiles and restaurant data.

Usage:
  uv run python scripts/manage_beli.py add <username> <beli_email> <beli_password>
  uv run python scripts/manage_beli.py list
  uv run python scripts/manage_beli.py refresh <profile_id>
  uv run python scripts/manage_beli.py refresh-all
  uv run python scripts/manage_beli.py remove <profile_id>

Examples:
  uv run python scripts/manage_beli.py add ryan me@gmail.com mybelipass
  uv run python scripts/manage_beli.py list
  uv run python scripts/manage_beli.py refresh 3f2a1b...
  uv run python scripts/manage_beli.py refresh-all
"""

import asyncio
import os
import sys

from dotenv import load_dotenv
from tortoise import Tortoise

load_dotenv()

DATABASE_URL = os.getenv(
    "DATABASE_URL", "postgres://raggr:changeme@localhost:5432/raggr"
)

TORTOISE_CONFIG = {
    "connections": {"default": DATABASE_URL},
    "apps": {
        "models": {
            "models": [
                "blueprints.users.models",
                "blueprints.beli.models",
                "blueprints.email.models",
                "blueprints.conversation.models",
                "blueprints.scheduled_messages.models",
                "blueprints.cron.models",
            ],
            "default_connection": "default",
        }
    },
}


async def _init():
    await Tortoise.init(config=TORTOISE_CONFIG)


async def _close():
    await Tortoise.close_connections()


async def cmd_add(username: str, beli_email: str, beli_password: str):
    from blueprints.beli.logic import MAX_PROFILES_PER_USER
    from blueprints.beli.models import BeliProfile
    from blueprints.users.models import User
    from utils.beli_service import authenticate

    user = await User.get_or_none(username=username)
    if not user:
        print(f"Error: no user with username '{username}'")
        return False

    count = await BeliProfile.filter(user=user).count()
    if count >= MAX_PROFILES_PER_USER:
        print(
            f"Error: user already has {MAX_PROFILES_PER_USER} Beli profiles (maximum)"
        )
        return False

    if await BeliProfile.exists(user=user, beli_email=beli_email):
        print(f"Error: profile for '{beli_email}' already exists for this user")
        return False

    print(f"Authenticating with Beli as {beli_email}...")
    try:
        session = await authenticate(beli_email, beli_password)
    except Exception as e:
        print(f"Error: Beli authentication failed — {e}")
        return False

    profile = await BeliProfile.create(
        user=user,
        beli_email=beli_email,
        beli_password=beli_password,
        beli_user_id=session["user_id"],
        beli_refresh_token=session["refresh_token"],
        enabled=True,
    )
    print(f"Profile created: {profile.id}")
    print(f"  Beli user ID : {profile.beli_user_id}")
    print(f"  Beli email   : {profile.beli_email}")
    print(f"  Owner        : {username}")
    return True


async def cmd_list():
    from blueprints.beli.models import BeliProfile, BeliRestaurant

    profiles = await BeliProfile.all().prefetch_related("user").order_by("created_at")
    if not profiles:
        print("No Beli profiles found.")
        return

    for p in profiles:
        count = await BeliRestaurant.filter(profile=p).count()
        refreshed = (
            p.last_refreshed_at.strftime("%Y-%m-%d %H:%M UTC")
            if p.last_refreshed_at
            else "never"
        )
        status = "enabled" if p.enabled else "DISABLED"
        print(f"{p.id}  [{status}]")
        print(f"  Owner    : {p.user.email}")
        print(f"  Beli     : {p.beli_email}  (user_id={p.beli_user_id})")
        print(f"  Restaurants: {count}  |  Last refresh: {refreshed}")
        if p.last_error:
            print(f"  Last error : {p.last_error}")
        print()


async def cmd_refresh(profile_id: str):
    from blueprints.beli.logic import refresh_profile
    from blueprints.beli.models import BeliProfile

    profile = await BeliProfile.get_or_none(id=profile_id)
    if not profile:
        print(f"Error: no profile with id '{profile_id}'")
        return False

    print(f"Refreshing profile {profile_id} ({profile.beli_email})...")
    result = await refresh_profile(profile_id)
    if result.get("ok"):
        print(f"Done — {result['upserted']} restaurants upserted")
    elif result.get("skipped"):
        print("Skipped (profile disabled)")
    else:
        print(f"Failed: {result.get('error')}")
        return False
    return True


async def cmd_refresh_all():
    from blueprints.beli.models import BeliProfile
    from blueprints.beli.logic import refresh_profile

    profiles = (
        await BeliProfile.filter(enabled=True).order_by("last_refreshed_at").all()
    )
    if not profiles:
        print("No enabled Beli profiles found.")
        return

    print(f"Refreshing {len(profiles)} profile(s)...")
    for p in profiles:
        print(f"\n  {p.beli_email} ({p.id})")
        result = await refresh_profile(str(p.id))
        if result.get("ok"):
            print(f"  → {result['upserted']} restaurants upserted")
        else:
            print(f"  → failed: {result.get('error')}")


async def cmd_remove(profile_id: str):
    from blueprints.beli.models import BeliProfile, BeliRestaurant

    profile = await BeliProfile.get_or_none(id=profile_id)
    if not profile:
        print(f"Error: no profile with id '{profile_id}'")
        return False

    deleted = await BeliRestaurant.filter(profile=profile).count()
    await BeliRestaurant.filter(profile=profile).delete()
    await profile.delete()
    print(
        f"Removed profile {profile_id} ({profile.beli_email}) and {deleted} restaurants"
    )
    return True


def usage():
    print(__doc__)


async def main():
    if len(sys.argv) < 2:
        usage()
        sys.exit(1)

    cmd = sys.argv[1]

    await _init()
    try:
        if cmd == "add":
            if len(sys.argv) != 5:
                print(
                    "Usage: uv run python scripts/manage_beli.py add <username> <beli_email> <beli_password>"
                )
                sys.exit(1)
            ok = await cmd_add(sys.argv[2], sys.argv[3], sys.argv[4])
            sys.exit(0 if ok else 1)

        elif cmd == "list":
            await cmd_list()

        elif cmd == "refresh":
            if len(sys.argv) != 3:
                print(
                    "Usage: uv run python scripts/manage_beli.py refresh <profile_id>"
                )
                sys.exit(1)
            ok = await cmd_refresh(sys.argv[2])
            sys.exit(0 if ok else 1)

        elif cmd == "refresh-all":
            await cmd_refresh_all()

        elif cmd == "remove":
            if len(sys.argv) != 3:
                print("Usage: uv run python scripts/manage_beli.py remove <profile_id>")
                sys.exit(1)
            ok = await cmd_remove(sys.argv[2])
            sys.exit(0 if ok else 1)

        else:
            print(f"Unknown command: {cmd}")
            usage()
            sys.exit(1)
    finally:
        await _close()


if __name__ == "__main__":
    asyncio.run(main())