db.py

# src/services/db.py
import aiosqlite
from datetime import datetime
import calendar
from pathlib import Path

class DB:
    def __init__(self, path: Path):
        self.path = Path(path)

    async def connect(self):
        self.conn = await aiosqlite.connect(self.path)
        # Performance tuning: WAL mode, larger cache, optimized sync
        await self.conn.execute('PRAGMA journal_mode=WAL')
        await self.conn.execute('PRAGMA synchronous=NORMAL')
        await self.conn.execute('PRAGMA cache_size=10000')
        await self.conn.execute('PRAGMA temp_store=MEMORY')
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS users(
            user_id INTEGER PRIMARY KEY,
            username TEXT,
            first_name TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        # Wallet balances (integer smallest currency unit; here we treat as VND)
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS wallets(
            user_id INTEGER PRIMARY KEY,
            balance INTEGER NOT NULL DEFAULT 0,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        # Incoming payments from PayOS
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS payments(
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            order_code TEXT UNIQUE,
            user_id INTEGER,
            amount INTEGER,
            status TEXT,
            raw JSON,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS orders(
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER, product TEXT, quantity INTEGER, filename TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        # Premium subscription info (optional)
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS premium_subs(
            user_id INTEGER PRIMARY KEY,
            plan TEXT,
            expires_at TIMESTAMP
        )''')
        # Referrals: each invited user can be linked to exactly one referrer
        await self.conn.execute('''CREATE TABLE IF NOT EXISTS referrals(
            invited_id INTEGER PRIMARY KEY,
            referrer_id INTEGER,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )''')
        await self.conn.commit()

    async def upsert_user(self, user_id:int, username:str, first_name:str):
        await self.conn.execute('INSERT INTO users(user_id, username, first_name) VALUES(?,?,?) '
                                'ON CONFLICT(user_id) DO UPDATE SET username=excluded.username, first_name=excluded.first_name',
                                (user_id, username, first_name))
        await self.conn.commit()

    # -------- Wallets --------
    async def get_balance(self, user_id:int) -> int:
        cur = await self.conn.execute('SELECT balance FROM wallets WHERE user_id=?', (user_id,))
        row = await cur.fetchone()
        await cur.close()
        return int(row[0]) if row else 0

    async def add_balance(self, user_id:int, amount:int) -> None:
        await self.conn.execute('INSERT OR IGNORE INTO wallets(user_id, balance) VALUES(?, 0)', (user_id,))
        await self.conn.execute('UPDATE wallets SET balance = balance + ?, updated_at=CURRENT_TIMESTAMP WHERE user_id=?', (amount, user_id))
        await self.conn.commit()

    async def deduct_balance(self, user_id:int, amount:int) -> bool:
        # Atomic check-and-deduct
        await self.conn.execute('INSERT OR IGNORE INTO wallets(user_id, balance) VALUES(?, 0)', (user_id,))
        cur = await self.conn.execute('UPDATE wallets SET balance = balance - ? WHERE user_id=? AND balance >= ?', (amount, user_id, amount))
        await self.conn.commit()
        return cur.rowcount > 0

    # -------- Payments --------
    async def create_or_update_payment_pending(self, order_code:str, user_id:int, amount:int) -> None:
        await self.conn.execute('INSERT INTO payments(order_code, user_id, amount, status) VALUES(?,?,?,?) '
                                'ON CONFLICT(order_code) DO UPDATE SET user_id=excluded.user_id, amount=excluded.amount, status=excluded.status, updated_at=CURRENT_TIMESTAMP',
                                (order_code, user_id, amount, 'PENDING'))
        await self.conn.commit()

    async def mark_payment_paid_and_credit(self, order_code:str, user_id:int, amount:int, raw_json:str) -> bool:
        # Atomic transaction: check status, update payment, credit wallet in one go
        async with self.conn.execute('BEGIN IMMEDIATE'):
            cur = await self.conn.execute('SELECT status FROM payments WHERE order_code=?', (order_code,))
            row = await cur.fetchone()
            await cur.close()
            if row and (row[0] or '').upper() == 'PAID':
                await self.conn.rollback()
                return False
            await self.conn.execute('INSERT OR IGNORE INTO payments(order_code, user_id, amount, status) VALUES(?,?,?,?)', (order_code, user_id, amount, 'PENDING'))
            await self.conn.execute('UPDATE payments SET status=?, amount=?, raw=?, updated_at=CURRENT_TIMESTAMP WHERE order_code=?', ('PAID', amount, raw_json, order_code))
            await self.conn.execute('INSERT OR IGNORE INTO wallets(user_id, balance) VALUES(?, 0)', (user_id,))
            await self.conn.execute('UPDATE wallets SET balance = balance + ?, updated_at=CURRENT_TIMESTAMP WHERE user_id=?', (amount, user_id))
            await self.conn.commit()
        return True

    async def add_order(self, user_id:int, product:str, quantity:int, filename:str):
        await self.conn.execute('INSERT INTO orders(user_id, product, quantity, filename) VALUES(?,?,?,?)',
                                (user_id, product, quantity, filename))
        await self.conn.commit()

    async def reset_all_balances(self) -> int:
        """Set all wallet balances to zero. Returns number of rows affected."""
        cur = await self.conn.execute('UPDATE wallets SET balance = 0, updated_at=CURRENT_TIMESTAMP WHERE balance <> 0')
        await self.conn.commit()
        return cur.rowcount or 0

    # -------- Profile/Referrals --------
    async def get_total_paid(self, user_id:int) -> int:
        cur = await self.conn.execute("SELECT COALESCE(SUM(amount),0) FROM payments WHERE user_id=? AND UPPER(status)='PAID'", (user_id,))
        row = await cur.fetchone()
        await cur.close()
        return int(row[0] or 0)

    async def set_referrer(self, invited_id:int, referrer_id:int) -> bool:
        if invited_id == referrer_id:
            return False
        try:
            await self.conn.execute('INSERT INTO referrals(invited_id, referrer_id) VALUES(?,?)', (invited_id, referrer_id))
            await self.conn.commit()
            return True
        except Exception:
            # Already referred or constraint violation
            return False

    async def get_ref_count(self, referrer_id:int) -> int:
        cur = await self.conn.execute('SELECT COUNT(1) FROM referrals WHERE referrer_id=?', (referrer_id,))
        row = await cur.fetchone()
        await cur.close()
        return int(row[0] or 0)

    async def get_total_paid_from_referrals(self, referrer_id:int) -> int:
        cur = await self.conn.execute(
            '''SELECT COALESCE(SUM(p.amount),0)
               FROM payments p
               WHERE UPPER(p.status)='PAID' AND p.user_id IN (
                    SELECT invited_id FROM referrals WHERE referrer_id=?
               )''', (referrer_id,)
        )
        row = await cur.fetchone()
        await cur.close()
        return int(row[0] or 0)

    async def get_premium(self, user_id:int):
        """Return (plan, expires_at) or (None, None) if not premium."""
        cur = await self.conn.execute('SELECT plan, expires_at FROM premium_subs WHERE user_id=?', (user_id,))
        row = await cur.fetchone()
        await cur.close()
        if not row:
            return None, None
        return row[0], row[1]

    async def grant_premium(self, user_id:int, plan:str, months:int) -> str:
        """Grant or extend premium for user and return new expires_at as ISO string."""
        # read current expiry
        cur = await self.conn.execute('SELECT expires_at FROM premium_subs WHERE user_id=?', (user_id,))
        row = await cur.fetchone()
        await cur.close()
        now = datetime.now()
        base = now
        if row and row[0]:
            try:
                current = datetime.fromisoformat(str(row[0]))
                if current > now:
                    base = current
            except Exception:
                base = now
        # add months
        year = base.year
        month = base.month + months
        year += (month - 1) // 12
        month = (month - 1) % 12 + 1
        day = min(base.day, calendar.monthrange(year, month)[1])
        new_exp = base.replace(year=year, month=month, day=day)
        exp_iso = new_exp.strftime('%Y-%m-%d %H:%M:%S')
        await self.conn.execute('INSERT INTO premium_subs(user_id, plan, expires_at) VALUES(?,?,?) '
                                'ON CONFLICT(user_id) DO UPDATE SET plan=excluded.plan, expires_at=excluded.expires_at',
                                (user_id, plan, exp_iso))
        await self.conn.commit()
        return exp_iso

    async def close(self):
        await self.conn.close()