from cachetools import TTLCache
import os

from utils.mongo import mongo

CACHE_SIZE = int(os.getenv("CACHE_SIZE", 1000))
CACHE_TTL = int(os.getenv("CACHE_TTL", 3600))

cache = TTLCache(CACHE_SIZE, CACHE_TTL)


def get_chat_info(chat_id: int) -> dict:
    cached_info = cache.get(chat_id)
    if cached_info is not None:
        return cached_info
    mongo_info = mongo.chats_collection.find_one({"chat_id": chat_id})
    if mongo_info is not None:
        cache[chat_id] = mongo_info
        return mongo_info
    chat_info = {"chat_id": chat_id, "state": "default", "probability": 100}
    mongo.chats_collection.insert_one(chat_info)
    cache[chat_id] = chat_info
    return chat_info


def set_values(chat_id: int, **values):
    cached_info = cache.get(chat_id)
    if cached_info is None:
        mongo_info = mongo.chats_collection.find_one({"chat_id": chat_id})
        if mongo_info is None:
            chat_info = {"chat_id": chat_id, "state": "default", "probability": 100}
            chat_info.update(values)
            mongo.chats_collection.insert_one(chat_info)
            cache[chat_id] = chat_info
        else:
            mongo.chats_collection.update_one({"chat_id": chat_id}, {"$set": values})
            mongo_info = dict(mongo_info)
            mongo_info.update(values)
            cache[chat_id] = mongo_info
    else:
        cached_info.update(values)
        mongo.chats_collection.update_one({"chat_id": chat_id}, {"$set": values})