WORDLE

database/crud/wordle.py

@@ -1,32 +1,47 @@
 from typing import Optional
 
 from database.enums import TempStorageKey
-from database.mongo_types import MongoCollection
-from database.schemas.mongo import WordleGame
+from database.mongo_types import MongoDatabase
+from database.schemas.mongo import TemporaryStorage, WordleGame
 from database.utils.datetime import today_only_date
 
-__all__ = ["get_active_wordle_game", "make_wordle_guess", "start_new_wordle_game"]
+__all__ = [
+    "get_active_wordle_game",
+    "make_wordle_guess",
+    "start_new_wordle_game",
+    "set_daily_word",
+    "reset_wordle_games",
+]
 
 
-async def get_active_wordle_game(collection: MongoCollection, user_id: int) -> Optional[WordleGame]:
+async def get_active_wordle_game(database: MongoDatabase, user_id: int) -> Optional[WordleGame]:
     """Find a player's active game"""
-    return await collection.find_one({"user_id": user_id})
+    collection = database[WordleGame.collection()]
+    result = await collection.find_one({"user_id": user_id})
+    if result is None:
+        return None
+
+    return WordleGame(**result)
 
 
-async def start_new_wordle_game(collection: MongoCollection, user_id: int) -> WordleGame:
+async def start_new_wordle_game(database: MongoDatabase, user_id: int) -> WordleGame:
     """Start a new game"""
+    collection = database[WordleGame.collection()]
     game = WordleGame(user_id=user_id)
     await collection.insert_one(game.dict(by_alias=True))
     return game
 
 
-async def make_wordle_guess(collection: MongoCollection, user_id: int, guess: str):
+async def make_wordle_guess(database: MongoDatabase, user_id: int, guess: str):
     """Make a guess in your current game"""
+    collection = database[WordleGame.collection()]
     await collection.update_one({"user_id": user_id}, {"$push": {"guesses": guess}})
 
 
-async def get_daily_word(collection: MongoCollection) -> Optional[str]:
+async def get_daily_word(database: MongoDatabase) -> Optional[str]:
     """Get the word of today"""
+    collection = database[TemporaryStorage.collection()]
+
     result = await collection.find_one({"key": TempStorageKey.WORDLE_WORD, "day": today_only_date()})
     if result is None:
         return None
@@ -34,16 +49,33 @@ async def get_daily_word(collection: MongoCollection) -> Optional[str]:
     return result["word"]
 
 
-async def set_daily_word(collection: MongoCollection, word: str):
+async def set_daily_word(database: MongoDatabase, word: str, *, forced: bool = False) -> str:
     """Set the word of today
 
     This does NOT overwrite the existing word if there is one, so that it can safely run
-    on startup every time
+    on startup every time.
+
+    In order to always overwrite the current word, set the "forced"-kwarg to True.
+
+    Returns the word that was chosen. If one already existed, return that instead.
     """
-    current_word = await get_daily_word(collection)
+    collection = database[TemporaryStorage.collection()]
+
+    current_word = None if forced else await get_daily_word(collection)
     if current_word is not None:
-        return
+        return current_word
 
     await collection.update_one(
-        {"key": TempStorageKey.WORDLE_WORD}, {"day": today_only_date(), "word": word}, upsert=True
+        {"key": TempStorageKey.WORDLE_WORD}, {"$set": {"day": today_only_date(), "word": word}}, upsert=True
     )
+
+    # Remove all active games
+    await reset_wordle_games(database)
+
+    return word
+
+
+async def reset_wordle_games(database: MongoDatabase):
+    """Reset all active games"""
+    collection = database[WordleGame.collection()]
+    await collection.drop()