subscriptions.py

"""
Модуль для работы с подписками и проверки прав доступа
"""
from enum import Enum
from datetime import datetime
from typing import Optional, Tuple
from database import Database
from config import OWNER_USER_IDS

# Инициализация базы данных для проверок
db = Database()


class SubscriptionType(Enum):
    """Типы подписок"""
    FREE = "FREE"
    PRO = "PRO"
    PRO_LIFETIME = "PRO_LIFETIME"


# Лимиты для FREE пользователей
FREE_MAX_PURCHASES = 10
FREE_ALLOWED_ASSETS = ['btc']  # Только BTC для FREE


def get_user_subscription(user_id: int) -> Tuple[SubscriptionType, Optional[datetime]]:
    """
    Получить тип подписки пользователя и дату истечения
    
    Returns:
        tuple: (SubscriptionType, expires_at или None для lifetime)
    """
    conn = db.get_connection()
    cursor = conn.cursor()
    
    try:
        cursor.execute('''
            SELECT subscription_type, subscription_expires_at
            FROM users
            WHERE user_id = ?
        ''', (user_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            sub_type_str, expires_at_str = row
            if sub_type_str:
                try:
                    sub_type = SubscriptionType(sub_type_str)
                except ValueError:
                    sub_type = SubscriptionType.FREE
            else:
                sub_type = SubscriptionType.FREE
            
            expires_at = None
            if expires_at_str:
                try:
                    expires_at = datetime.fromisoformat(expires_at_str)
                except (ValueError, TypeError):
                    expires_at = None
            
            # Проверяем, не истекла ли подписка PRO
            if sub_type == SubscriptionType.PRO and expires_at:
                if expires_at < datetime.now():
                    # Подписка истекла, возвращаем FREE
                    update_user_subscription(user_id, SubscriptionType.FREE, None)
                    return SubscriptionType.FREE, None
            
            return sub_type, expires_at
        else:
            # Пользователь не найден, возвращаем FREE по умолчанию
            return SubscriptionType.FREE, None
            
    except Exception as e:
        print(f"Ошибка при получении подписки пользователя: {e}")
        conn.close()
        return SubscriptionType.FREE, None


def update_user_subscription(user_id: int, subscription_type: SubscriptionType, 
                            expires_at: Optional[datetime] = None) -> bool:
    """
    Обновить подписку пользователя
    
    Args:
        user_id: ID пользователя
        subscription_type: Тип подписки
        expires_at: Дата истечения (None для lifetime)
    
    Returns:
        bool: Успешность операции
    """
    conn = db.get_connection()
    cursor = conn.cursor()
    
    try:
        expires_at_str = expires_at.isoformat() if expires_at else None
        
        cursor.execute('''
            UPDATE users
            SET subscription_type = ?, subscription_expires_at = ?
            WHERE user_id = ?
        ''', (subscription_type.value, expires_at_str, user_id))
        
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"Ошибка при обновлении подписки: {e}")
        conn.close()
        return False


def has_pro_access(user_id: int) -> bool:
    """
    Проверить, есть ли у пользователя PRO доступ
    
    Returns:
        bool: True если PRO или PRO_LIFETIME или владелец бота
    """
    # Владельцы бота всегда имеют полный доступ
    if user_id in OWNER_USER_IDS:
        return True
    
    sub_type, _ = get_user_subscription(user_id)
    return sub_type in [SubscriptionType.PRO, SubscriptionType.PRO_LIFETIME]


def can_add_asset(user_id: int, asset: str) -> Tuple[bool, str]:
    """
    Проверить, может ли пользователь добавить актив
    
    Args:
        user_id: ID пользователя
        asset: Название актива (btc, gold)
    
    Returns:
        tuple: (разрешено, сообщение об ошибке)
    """
    # Владельцы бота могут добавлять любые активы
    if user_id in OWNER_USER_IDS:
        return True, ""
    
    sub_type, _ = get_user_subscription(user_id)
    
    if sub_type in [SubscriptionType.PRO, SubscriptionType.PRO_LIFETIME]:
        return True, ""
    
    # FREE пользователь может добавлять только BTC
    if asset.lower() not in FREE_ALLOWED_ASSETS:
        return False, "gold_not_allowed"
    
    return True, ""


def can_add_purchase(user_id: int) -> Tuple[bool, str, int]:
    """
    Проверить, может ли пользователь добавить покупку
    
    Args:
        user_id: ID пользователя
    
    Returns:
        tuple: (разрешено, сообщение, текущее количество покупок)
    """
    # Владельцы бота не имеют ограничений
    if user_id in OWNER_USER_IDS:
        return True, "", 0
    
    sub_type, _ = get_user_subscription(user_id)
    
    if sub_type in [SubscriptionType.PRO, SubscriptionType.PRO_LIFETIME]:
        return True, "", 0
    
    # Для FREE пользователя проверяем лимит покупок
    transactions = db.get_user_transactions(user_id)
    purchase_count = len(transactions)
    
    if purchase_count >= FREE_MAX_PURCHASES:
        return False, "max_purchases_reached", purchase_count
    
    return True, "", purchase_count


def can_view_statistics(user_id: int) -> bool:
    """
    Проверить, может ли пользователь просматривать детальную статистику и графики
    
    Returns:
        bool: True если PRO или PRO_LIFETIME
    """
    return has_pro_access(user_id)


def can_view_charts(user_id: int) -> bool:
    """
    Проверить, может ли пользователь просматривать графики
    
    Returns:
        bool: True если PRO или PRO_LIFETIME
    """
    return has_pro_access(user_id)


def can_use_extended_reserve(user_id: int) -> bool:
    """
    Проверить, может ли пользователь использовать расширенный резерв
    
    Returns:
        bool: True если PRO или PRO_LIFETIME
    """
    return has_pro_access(user_id)


def get_subscription_info(user_id: int) -> dict:
    """
    Получить информацию о подписке пользователя для отображения
    
    Returns:
        dict: Информация о подписке
    """
    # Владельцы бота всегда имеют PRO_LIFETIME
    if user_id in OWNER_USER_IDS:
        return {
            'type': SubscriptionType.PRO_LIFETIME,
            'type_name': 'PRO_LIFETIME',
            'expires_at': None,
            'is_pro': True,
            'days_left': None,  # Бессрочно
            'is_owner': True
        }
    
    sub_type, expires_at = get_user_subscription(user_id)
    
    info = {
        'type': sub_type,
        'type_name': sub_type.value,
        'expires_at': expires_at,
        'is_pro': has_pro_access(user_id),
        'is_owner': False
    }
    
    if sub_type == SubscriptionType.PRO and expires_at:
        info['days_left'] = (expires_at - datetime.now()).days
    elif sub_type == SubscriptionType.PRO_LIFETIME:
        info['days_left'] = None  # Бессрочно
    
    return info