database.py

"""
Модуль для работы с базой данных SQLite
"""
import sqlite3
import os
from datetime import datetime
from typing import List, Dict, Optional, Tuple
from config import DATABASE_NAME


class Database:
    """Класс для работы с базой данных"""
    
    def __init__(self, db_name: str = None):
        # Используем значение из конфига, если не передано явно
        if db_name is None:
            db_name = DATABASE_NAME
        # Если и в конфиге None, используем значение по умолчанию
        if db_name is None:
            db_name = 'wealth_logger.db'
        self.db_name = db_name
        self.init_database()
    
    def get_connection(self):
        """Получить соединение с базой данных"""
        if self.db_name is None:
            raise ValueError("Имя базы данных не может быть None")
        return sqlite3.connect(self.db_name)
    
    def init_database(self):
        """Инициализация базы данных: создание таблиц"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Таблица пользователей
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                user_id INTEGER PRIMARY KEY,
                username TEXT,
                first_name TEXT,
                registered_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                subscription_type TEXT DEFAULT 'FREE',
                subscription_expires_at TEXT
            )
        ''')
        
        # Добавляем поля подписки к существующей таблице (миграция)
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN subscription_type TEXT DEFAULT "FREE"')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN subscription_expires_at TEXT')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        # Добавляем поля реферальной системы (миграция)
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN referral_code TEXT')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN referred_by INTEGER')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN referral_count INTEGER DEFAULT 0')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        try:
            cursor.execute('ALTER TABLE users ADD COLUMN referral_bonus_expires_at TEXT')
        except sqlite3.OperationalError:
            pass  # Колонка уже существует
        
        # Таблица транзакций (покупок)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS transactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                asset TEXT NOT NULL,
                amount_usd REAL NOT NULL,
                price REAL NOT NULL,
                quantity REAL NOT NULL,
                transaction_date DATE NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (user_id)
            )
        ''')
        
        # Таблица текущих цен (кэш)
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS prices (
                asset TEXT PRIMARY KEY,
                price REAL NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # Таблица настроек пользователей
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS user_settings (
                user_id INTEGER PRIMARY KEY,
                price_alerts_enabled INTEGER DEFAULT 1,
                price_change_threshold REAL DEFAULT 5.0,
                reserve_amount REAL DEFAULT 0.0,
                FOREIGN KEY (user_id) REFERENCES users (user_id)
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def register_user(self, user_id: int, username: str = None, first_name: str = None) -> bool:
        """Регистрация нового пользователя"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT OR IGNORE INTO users (user_id, username, first_name, subscription_type)
                VALUES (?, ?, ?, 'FREE')
            ''', (user_id, username, first_name))
            conn.commit()
            return cursor.rowcount > 0
        except Exception as e:
            print(f"Ошибка при регистрации пользователя: {e}")
            return False
        finally:
            conn.close()
    
    def add_transaction(self, user_id: int, asset: str, amount_usd: float, 
                       price: float, transaction_date: str) -> bool:
        """Добавление новой транзакции (покупки)"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Вычисляем количество купленного актива
        quantity = amount_usd / price if price > 0 else 0
        
        try:
            cursor.execute('''
                INSERT INTO transactions (user_id, asset, amount_usd, price, quantity, transaction_date)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (user_id, asset.lower(), amount_usd, price, quantity, transaction_date))
            conn.commit()
            return True
        except Exception as e:
            print(f"Ошибка при добавлении транзакции: {e}")
            return False
        finally:
            conn.close()
    
    def get_user_transactions(self, user_id: int, asset: str = None) -> List[Dict]:
        """Получить все транзакции пользователя, опционально фильтр по активу"""
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        if asset:
            cursor.execute('''
                SELECT * FROM transactions
                WHERE user_id = ? AND asset = ?
                ORDER BY transaction_date DESC, created_at DESC
            ''', (user_id, asset.lower()))
        else:
            cursor.execute('''
                SELECT * FROM transactions
                WHERE user_id = ?
                ORDER BY transaction_date DESC, created_at DESC
            ''', (user_id,))
        
        rows = cursor.fetchall()
        conn.close()
        
        return [dict(row) for row in rows]
    
    def get_portfolio_summary(self, user_id: int) -> Dict:
        """Получить сводку по портфелю пользователя (оптимизированная версия)"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # Получаем все транзакции пользователя и вычисляем взвешенную среднюю цену
        # Взвешенная средняя = сумма(amount_usd) / сумма(quantity)
        cursor.execute('''
            SELECT asset, 
                   SUM(amount_usd) as total_invested, 
                   SUM(quantity) as total_quantity,
                   CASE 
                       WHEN SUM(quantity) > 0 THEN SUM(amount_usd) / SUM(quantity)
                       ELSE 0
                   END as avg_price
            FROM transactions
            WHERE user_id = ?
            GROUP BY asset
        ''', (user_id,))
        
        rows = cursor.fetchall()
        conn.close()
        
        portfolio = {}
        for row in rows:
            asset, total_invested, total_quantity, avg_price = row
            portfolio[asset] = {
                'total_invested': total_invested or 0,
                'total_quantity': total_quantity or 0,
                'avg_price': avg_price or 0
            }
        
        return portfolio
    
    def get_asset_stats(self, user_id: int, asset: str) -> Optional[Dict]:
        """Получить статистику по конкретному активу"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT 
                COUNT(*) as transaction_count,
                SUM(amount_usd) as total_invested,
                SUM(quantity) as total_quantity,
                AVG(price) as avg_price,
                MIN(price) as min_price,
                MAX(price) as max_price,
                MIN(transaction_date) as first_purchase,
                MAX(transaction_date) as last_purchase
            FROM transactions
            WHERE user_id = ? AND asset = ?
        ''', (user_id, asset.lower()))
        
        row = cursor.fetchone()
        conn.close()
        
        if row and row[0] > 0:
            return {
                'transaction_count': row[0],
                'total_invested': row[1] or 0,
                'total_quantity': row[2] or 0,
                'avg_price': row[3] or 0,
                'min_price': row[4] or 0,
                'max_price': row[5] or 0,
                'first_purchase': row[6],
                'last_purchase': row[7]
            }
        return None
    
    def update_price(self, asset: str, price: float) -> bool:
        """Обновить текущую цену актива"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                INSERT OR REPLACE INTO prices (asset, price, updated_at)
                VALUES (?, ?, CURRENT_TIMESTAMP)
            ''', (asset.lower(), price))
            conn.commit()
            return True
        except Exception as e:
            print(f"Ошибка при обновлении цены: {e}")
            return False
        finally:
            conn.close()
    
    def get_price(self, asset: str) -> Optional[float]:
        """Получить текущую цену актива из кэша"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT price FROM prices
            WHERE asset = ?
        ''', (asset.lower(),))
        
        row = cursor.fetchone()
        conn.close()
        
        return row[0] if row else None
    
    def delete_transaction(self, user_id: int, transaction_id: int) -> bool:
        """Удалить транзакцию"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        try:
            cursor.execute('''
                DELETE FROM transactions
                WHERE id = ? AND user_id = ?
            ''', (transaction_id, user_id))
            conn.commit()
            deleted = cursor.rowcount > 0
            return deleted
        except Exception as e:
            print(f"Ошибка при удалении транзакции: {e}")
            return False
        finally:
            conn.close()
    
    def update_transaction(self, user_id: int, transaction_id: int, 
                          amount_usd: float = None, price: float = None,
                          transaction_date: str = None) -> bool:
        """Обновить транзакцию"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        try:
            updates = []
            params = []
            
            if amount_usd is not None:
                updates.append('amount_usd = ?')
                params.append(amount_usd)
            
            if price is not None:
                updates.append('price = ?')
                params.append(price)
            
            if transaction_date is not None:
                updates.append('transaction_date = ?')
                params.append(transaction_date)
            
            if not updates:
                return False
            
            # Пересчитываем quantity если изменились amount_usd или price
            if amount_usd is not None or price is not None:
                # Получаем текущие значения
                cursor.execute('SELECT amount_usd, price FROM transactions WHERE id = ? AND user_id = ?',
                             (transaction_id, user_id))
                row = cursor.fetchone()
                if not row:
                    return False
                
                current_amount = amount_usd if amount_usd is not None else row[0]
                current_price = price if price is not None else row[1]
                new_quantity = current_amount / current_price if current_price > 0 else 0
                
                updates.append('quantity = ?')
                params.append(new_quantity)
            
            params.extend([transaction_id, user_id])
            
            query = f'''
                UPDATE transactions
                SET {', '.join(updates)}
                WHERE id = ? AND user_id = ?
            '''
            
            cursor.execute(query, params)
            conn.commit()
            return cursor.rowcount > 0
        except Exception as e:
            print(f"Ошибка при обновлении транзакции: {e}")
            return False
        finally:
            conn.close()
    
    def get_transaction(self, user_id: int, transaction_id: int) -> Optional[Dict]:
        """Получить конкретную транзакцию"""
        conn = self.get_connection()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT * FROM transactions
            WHERE id = ? AND user_id = ?
        ''', (transaction_id, user_id))
        
        row = cursor.fetchone()
        conn.close()
        
        return dict(row) if row else None
    
    def get_user_settings(self, user_id: int) -> Dict:
        """Получить настройки пользователя"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT price_alerts_enabled, price_change_threshold, reserve_amount
            FROM user_settings
            WHERE user_id = ?
        ''', (user_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return {
                'price_alerts_enabled': bool(row[0]),
                'price_change_threshold': row[1] or 5.0,
                'reserve_amount': row[2] or 0.0
            }
        else:
            # Создаём настройки по умолчанию
            self.update_user_settings(user_id, True, 5.0, 0.0)
            return {
                'price_alerts_enabled': True,
                'price_change_threshold': 5.0,
                'reserve_amount': 0.0
            }
    
    def update_user_settings(self, user_id: int, price_alerts_enabled: bool = None,
                            price_change_threshold: float = None, reserve_amount: float = None) -> bool:
        """Обновить настройки пользователя"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        try:
            # Проверяем существование записи
            cursor.execute('SELECT user_id FROM user_settings WHERE user_id = ?', (user_id,))
            exists = cursor.fetchone()
            
            if exists:
                updates = []
                params = []
                
                if price_alerts_enabled is not None:
                    updates.append('price_alerts_enabled = ?')
                    params.append(1 if price_alerts_enabled else 0)
                
                if price_change_threshold is not None:
                    updates.append('price_change_threshold = ?')
                    params.append(price_change_threshold)
                
                if reserve_amount is not None:
                    updates.append('reserve_amount = ?')
                    params.append(reserve_amount)
                
                if updates:
                    params.append(user_id)
                    query = f'''
                        UPDATE user_settings
                        SET {', '.join(updates)}
                        WHERE user_id = ?
                    '''
                    cursor.execute(query, params)
            else:
                # Создаём новую запись
                cursor.execute('''
                    INSERT INTO user_settings (user_id, price_alerts_enabled, price_change_threshold, reserve_amount)
                    VALUES (?, ?, ?, ?)
                ''', (
                    user_id,
                    1 if price_alerts_enabled else 0 if price_alerts_enabled is not None else 1,
                    price_change_threshold if price_change_threshold is not None else 5.0,
                    reserve_amount if reserve_amount is not None else 0.0
                ))
            
            conn.commit()
            return True
        except Exception as e:
            print(f"Ошибка при обновлении настроек: {e}")
            return False
        finally:
            conn.close()