mdb.py

import os
import time
from pymongo import MongoClient

# -------- CONFIG --------
MONGO_URI = os.getenv("MONGO_URI")
DB_NAME = "tgstoragebot"   # tum change kar sakte ho
COLLECTION_NAME = "UserStorageData"
# ------------------------

if not MONGO_URI:
    raise ValueError("MONGO_URI not set in environment variables")

client = MongoClient(MONGO_URI)
db = client[DB_NAME]
col = db[COLLECTION_NAME]

# -------- CACHE --------
_cache = {}
CACHE_TTL = 600  # 10 min


# -------- HELPERS --------

def _normalize_type(dtype):
    return dtype.lower() if dtype else "string"


def _cache_key(key, dtype):
    return f"{key}:{dtype}"


def _set_cache(key, dtype, value):
    _cache[_cache_key(key, dtype)] = {
        "value": value,
        "time": time.time()
    }


def _get_cache(key, dtype):
    ckey = _cache_key(key, dtype)
    data = _cache.get(ckey)

    if not data:
        return None

    if time.time() - data["time"] > CACHE_TTL:
        del _cache[ckey]
        return None

    # 🔥 refresh TTL
    data["time"] = time.time()

    return data["value"]


def _delete_cache(key):
    for k in list(_cache.keys()):
        if k.startswith(f"{key}:"):
            del _cache[k]


# -------- CORE --------

def set(key, dtype, value):
    dtype = _normalize_type(dtype)

    col.update_one(
        {"key": key, "type": dtype},
        {"$set": {"key": key, "type": dtype, "value": value}},
        upsert=True
    )

    _set_cache(key, dtype, value)
    return True


def get(key, dtype="string"):
    dtype = _normalize_type(dtype)

    cached = _get_cache(key, dtype)
    if cached is not None:
        return cached

    doc = col.find_one({"key": key, "type": dtype})
    if not doc:
        return None

    value = doc["value"]
    _set_cache(key, dtype, value)

    return value


def update(key, dtype, value):
    dtype = _normalize_type(dtype)

    result = col.update_one(
        {"key": key, "type": dtype},
        {"$set": {"value": value}}
    )

    if result.matched_count == 0:
        return False

    _set_cache(key, dtype, value)
    return True


def delete(key):
    col.delete_many({"key": key})
    _delete_cache(key)
    return True