database.py

import aiosqlite
import asyncio
from config import DB_PATH
from datetime import datetime


async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS users (
                user_id INTEGER PRIMARY KEY,
                username TEXT,
                first_name TEXT,
                language TEXT DEFAULT 'ru',
                referrer_id INTEGER,
                referral_code TEXT UNIQUE,
                balance_ton REAL DEFAULT 0.0,
                total_earned_ton REAL DEFAULT 0.0,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS orders (
                id TEXT PRIMARY KEY,
                user_id INTEGER,
                order_type TEXT,  -- 'stars' or 'premium'
                target_username TEXT,
                stars_count INTEGER,
                premium_months INTEGER,
                payment_method TEXT,  -- 'ton', 'sbp', 'cryptobot'
                amount_ton REAL,
                amount_usdt REAL,
                amount_rub REAL,
                memo TEXT,
                status TEXT DEFAULT 'pending',  -- pending, paid, completed, cancelled
                promo_code TEXT,
                discount_percent INTEGER DEFAULT 0,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                paid_at TEXT
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS referral_earnings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                referrer_id INTEGER,
                from_user_id INTEGER,
                order_id TEXT,
                amount_ton REAL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS withdrawals (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                amount_ton REAL,
                wallet_address TEXT,
                status TEXT DEFAULT 'pending',
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS promo_codes (
                code TEXT PRIMARY KEY,
                discount_percent INTEGER,
                max_uses INTEGER DEFAULT -1,
                used_count INTEGER DEFAULT 0,
                is_active INTEGER DEFAULT 1,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        await db.commit()


async def get_or_create_user(user_id: int, username: str = None, first_name: str = None, referrer_id: int = None):
    import secrets
    import string
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)) as cursor:
            user = await cursor.fetchone()
        if not user:
            ref_code = ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(12))
            await db.execute(
                "INSERT INTO users (user_id, username, first_name, referral_code, referrer_id) VALUES (?, ?, ?, ?, ?)",
                (user_id, username, first_name, ref_code, referrer_id)
            )
            await db.commit()
            async with db.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)) as cursor:
                user = await cursor.fetchone()
        return user


async def get_user(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)) as cursor:
            return await cursor.fetchone()


async def get_user_by_ref_code(ref_code: str):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM users WHERE referral_code = ?", (ref_code,)) as cursor:
            return await cursor.fetchone()


async def create_order(order_id: str, user_id: int, order_type: str, payment_method: str,
                       stars_count: int = None, premium_months: int = None,
                       target_username: str = None, amount_ton: float = None,
                       amount_usdt: float = None, amount_rub: float = None, memo: str = None):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            INSERT INTO orders (id, user_id, order_type, target_username, stars_count, premium_months,
                                payment_method, amount_ton, amount_usdt, amount_rub, memo)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (order_id, user_id, order_type, target_username, stars_count, premium_months,
              payment_method, amount_ton, amount_usdt, amount_rub, memo))
        await db.commit()


async def get_order(order_id: str):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("SELECT * FROM orders WHERE id = ?", (order_id,)) as cursor:
            return await cursor.fetchone()


async def update_order_status(order_id: str, status: str):
    async with aiosqlite.connect(DB_PATH) as db:
        paid_at = datetime.now().isoformat() if status == 'paid' else None
        if paid_at:
            await db.execute("UPDATE orders SET status = ?, paid_at = ? WHERE id = ?", (status, paid_at, order_id))
        else:
            await db.execute("UPDATE orders SET status = ? WHERE id = ?", (status, order_id))
        await db.commit()


async def apply_promo_to_order(order_id: str, promo_code: str, discount: int):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE orders SET promo_code = ?, discount_percent = ? WHERE id = ?",
            (promo_code, discount, order_id)
        )
        await db.commit()


async def get_referral_stats(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT COUNT(*) FROM users WHERE referrer_id = ?", (user_id,)) as cursor:
            total_refs = (await cursor.fetchone())[0]
        async with db.execute(
            "SELECT COALESCE(SUM(amount_ton), 0) FROM referral_earnings WHERE referrer_id = ?", (user_id,)
        ) as cursor:
            total_earned = (await cursor.fetchone())[0]
        async with db.execute("SELECT balance_ton FROM users WHERE user_id = ?", (user_id,)) as cursor:
            row = await cursor.fetchone()
            balance = row[0] if row else 0.0
        return {"total_refs": total_refs, "total_earned": total_earned, "balance": balance}


async def get_referral_earnings(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM referral_earnings WHERE referrer_id = ? ORDER BY created_at DESC",
            (user_id,)
        ) as cursor:
            return await cursor.fetchall()


async def get_withdrawal_history(user_id: int):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM withdrawals WHERE user_id = ? ORDER BY created_at DESC",
            (user_id,)
        ) as cursor:
            return await cursor.fetchall()


async def create_withdrawal(user_id: int, amount: float, wallet: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO withdrawals (user_id, amount_ton, wallet_address) VALUES (?, ?, ?)",
            (user_id, amount, wallet)
        )
        await db.execute(
            "UPDATE users SET balance_ton = balance_ton - ? WHERE user_id = ?",
            (amount, user_id)
        )
        await db.commit()


async def get_promo(code: str):
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM promo_codes WHERE code = ? AND is_active = 1",
            (code.upper(),)
        ) as cursor:
            promo = await cursor.fetchone()
        if not promo:
            return None
        if promo['max_uses'] != -1 and promo['used_count'] >= promo['max_uses']:
            return None
        return promo


async def use_promo(code: str):
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE promo_codes SET used_count = used_count + 1 WHERE code = ?",
            (code.upper(),)
        )
        await db.commit()