import asyncio
import argparse

from app.core.database import AsyncSessionLocal, init_db
from app.services.bazi_backfill_service import backfill_bazi_wuxing_balance
from app.services.cleanup_service import cleanup_expired_images


async def cleanup_command() -> None:
    await init_db()
    async with AsyncSessionLocal() as session:
        count = await cleanup_expired_images(session)
        await session.commit()
    print(f"cleaned {count} expired images")


async def backfill_bazi_wuxing_command() -> None:
    await init_db()
    async with AsyncSessionLocal() as session:
        count = await backfill_bazi_wuxing_balance(session)
        await session.commit()
    print(f"backfilled {count} bazi wuxing balances")


async def main() -> None:
    parser = argparse.ArgumentParser(description="Cyber Mister maintenance commands")
    parser.add_argument("command", nargs="?", default="cleanup-expired-images", choices=["cleanup-expired-images", "backfill-bazi-wuxing"])
    args = parser.parse_args()

    if args.command == "backfill-bazi-wuxing":
        await backfill_bazi_wuxing_command()
    else:
        await cleanup_command()


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