CatismImage/models.py

from pymongo import MongoClient
from config import Config
import random
import string
import datetime
from bson.int64 import Int64
from bson.decimal128 import Decimal128

client = MongoClient(Config.MONGO_URI)
db = client[Config.MONGO_DB]

class User:
    collection = db['users']
    
    @staticmethod
    def generate_key(length=12):
        """生成由小写字母和数字组成的随机key"""
        chars = string.ascii_lowercase + string.digits
        return ''.join(random.choice(chars) for _ in range(length))
    
    @classmethod
    def create_user(cls, nickname):
        """创建新用户并返回生成的key"""
        key = cls.generate_key()
        user_data = {
            'key': key,
            'nickname': nickname,
            'register_time': datetime.datetime.utcnow(),
            'upload_count': Int64(0),
            'storage_limit': Int64(100 * 1024 * 1024),  # 默认100MB
            'used_storage': Int64(0),
            'traffic_limit': Int64(1000 * 1024 * 1024),  # 默认1000MB流量限制
            'monthly_traffic': {
                f"{datetime.datetime.utcnow().year}-{datetime.datetime.utcnow().month}": Int64(0)
            },
            'meowcoin': 0.0,
            'bills': [],
            'custom': {
                'theme_preference': 'light',
                'theme_color': '#3296fa'
            }
        }
        cls.collection.insert_one(user_data)
        return key
        
    @classmethod
    def increment_upload_count(cls, key):
        """增加用户上传图片计数"""
        return cls.collection.update_one(
            {'key': key},
            {'$inc': {'upload_count': 1}}
        )
        
    @classmethod
    def decrement_upload_count(cls, key):
        """减少用户上传图片计数"""
        return cls.collection.update_one(
            {'key': key},
            {'$inc': {'upload_count': -1}}
        )
    
    @classmethod
    def get_user(cls, key):
        """根据key获取用户信息"""
        user = cls.collection.find_one({'key': key})
        if user:
            user['register_time'] = user.get('register_time', datetime.datetime.utcnow())
            user['upload_count'] = Int64(user.get('upload_count', 0))
            user['storage_limit'] = Int64(user.get('storage_limit', 100 * 1024 * 1024))
            user['used_storage'] = Int64(user.get('used_storage', 0))
            user['traffic_limit'] = Int64(user.get('traffic_limit', 1000 * 1024 * 1024))
            # 获取当月流量
            now = datetime.datetime.utcnow()
            month_key = f"{now.year}-{now.month}"
            user['monthly_traffic'] = user.get('monthly_traffic', {})
            user['current_month_traffic'] = Int64(user['monthly_traffic'].get(month_key, 0))
            user['meowcoin'] = float(user.get('meowcoin', 0.0))
            user['bills'] = user.get('bills', [])
            user['custom'] = user.get('custom', {
                'theme_preference': 'light',
                'theme_color': '#3296fa'
            })
        return user
    
    @classmethod
    def update_nickname(cls, key, new_nickname):
        """更新用户昵称"""
        return cls.collection.update_one(
            {'key': key},
            {'$set': {'nickname': new_nickname}}
        )
        
    @classmethod
    def update_storage(cls, key, size):
        """更新用户存储使用情况"""
        return cls.collection.update_one(
            {'key': key},
            {'$inc': {'used_storage': size}}
        )
        
    @classmethod
    def update_traffic(cls, key, size):
        """更新用户流量使用情况"""
        now = datetime.datetime.utcnow()
        month_key = f"{now.year}-{now.month}"
        return cls.collection.update_one(
            {'key': key},
            {'$inc': {
                f'monthly_traffic.{month_key}': size
            }}
        )
        
    @classmethod
    def update_meowcoin(cls, key, amount):
        """更新用户meowcoin余额"""
        return cls.collection.update_one(
            {'key': key},
            {'$inc': {'meowcoin': float(amount)}}
        )
        
    @classmethod
    def add_bill(cls, key, trade_id, user_id, plan_id, show_amount, discount, total_amount):
        """新增账单记录"""
        bill = {
            'trade_id': trade_id,
            'user_id': user_id,
            'plan_id': plan_id,
            'amount': show_amount,
            'discount': discount,
            'actual_amount': total_amount,
            'time': datetime.datetime.utcnow()
        }
        return cls.collection.update_one(
            {'key': key},
            {'$push': {'bills': bill}}
        )