database.py

"""
database.py — All SQLite operations for the bot.
Uses aiosqlite for non-blocking async I/O compatible with python-telegram-bot v20.
"""

import aiosqlite
import logging
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

DB_PATH = "dropsvipbot.db"


# ══════════════════════════════════════════════════════════════════════════════
#  Schema
# ══════════════════════════════════════════════════════════════════════════════

SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
    user_id     INTEGER PRIMARY KEY,
    username    TEXT,
    first_name  TEXT,
    last_name   TEXT,
    created_at  TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    is_banned   INTEGER   DEFAULT 0
);

CREATE TABLE IF NOT EXISTS subscriptions (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id     INTEGER NOT NULL,
    sub_type    TEXT    NOT NULL,   -- 'paid' | 'free'
    start_date  TIMESTAMP NOT NULL,
    end_date    TIMESTAMP NOT NULL,
    is_active   INTEGER DEFAULT 1,
    FOREIGN KEY (user_id) REFERENCES users(user_id)
);

CREATE TABLE IF NOT EXISTS payments (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id     INTEGER NOT NULL,
    method      TEXT    NOT NULL,   -- 'stripe' | 'crypto' | 'stars' | 'paypal'
    amount      REAL,
    currency    TEXT,
    status      TEXT    DEFAULT 'pending',  -- 'pending' | 'approved' | 'rejected'
    proof       TEXT,               -- TXID, screenshot file_id, Stripe session id, etc.
    created_at  TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (user_id) REFERENCES users(user_id)
);

CREATE TABLE IF NOT EXISTS share_submissions (
    id           INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id      INTEGER NOT NULL,
    platform     TEXT    NOT NULL,  -- 'twitter' | 'discord' | 'telegram' | 'reddit' | 'instagram'
    proof        TEXT    NOT NULL,  -- URL or Telegram file_id
    proof_type   TEXT    NOT NULL,  -- 'url' | 'photo'
    status       TEXT    DEFAULT 'pending',  -- 'pending' | 'approved' | 'rejected'
    submitted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (user_id) REFERENCES users(user_id)
);
"""


# ══════════════════════════════════════════════════════════════════════════════
#  Init
# ══════════════════════════════════════════════════════════════════════════════

async def init_db() -> None:
    """Create all tables on first run."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.executescript(SCHEMA)
        await db.commit()
    logger.info("Database initialised at %s", DB_PATH)


# ══════════════════════════════════════════════════════════════════════════════
#  Users
# ══════════════════════════════════════════════════════════════════════════════

async def upsert_user(
    user_id: int,
    username: Optional[str],
    first_name: str,
    last_name: Optional[str],
) -> None:
    """Insert or update a user record."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            """
            INSERT INTO users (user_id, username, first_name, last_name)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(user_id) DO UPDATE SET
                username   = excluded.username,
                first_name = excluded.first_name,
                last_name  = excluded.last_name
            """,
            (user_id, username, first_name, last_name),
        )
        await db.commit()


async def get_user(user_id: int) -> Optional[Dict[str, Any]]:
    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 cur:
            row = await cur.fetchone()
            return dict(row) if row else None


async def get_all_active_subscriber_ids() -> List[int]:
    """Return user_ids with a currently active subscription."""
    now = datetime.utcnow().isoformat()
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            """
            SELECT DISTINCT user_id FROM subscriptions
            WHERE is_active = 1 AND end_date > ?
            """,
            (now,),
        ) as cur:
            rows = await cur.fetchall()
            return [r[0] for r in rows]


# ══════════════════════════════════════════════════════════════════════════════
#  Subscriptions
# ══════════════════════════════════════════════════════════════════════════════

async def add_subscription(
    user_id: int,
    sub_type: str,
    months: int,
) -> None:
    """Grant a subscription starting now for `months` months."""
    start = datetime.utcnow()
    # Add months by advancing the month field, rolling over year as needed
    month = start.month - 1 + months
    year = start.year + month // 12
    month = month % 12 + 1
    end = start.replace(year=year, month=month)

    async with aiosqlite.connect(DB_PATH) as db:
        # Deactivate any existing active subscriptions first
        await db.execute(
            "UPDATE subscriptions SET is_active = 0 WHERE user_id = ? AND is_active = 1",
            (user_id,),
        )
        await db.execute(
            """
            INSERT INTO subscriptions (user_id, sub_type, start_date, end_date)
            VALUES (?, ?, ?, ?)
            """,
            (user_id, sub_type, start.isoformat(), end.isoformat()),
        )
        await db.commit()
    logger.info("Subscription granted: user=%s type=%s until=%s", user_id, sub_type, end)


async def has_active_subscription(user_id: int) -> bool:
    now = datetime.utcnow().isoformat()
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            """
            SELECT id FROM subscriptions
            WHERE user_id = ? AND is_active = 1 AND end_date > ?
            LIMIT 1
            """,
            (user_id, now),
        ) as cur:
            return (await cur.fetchone()) is not None


async def get_subscription_expiry(user_id: int) -> Optional[str]:
    now = datetime.utcnow().isoformat()
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            """
            SELECT end_date FROM subscriptions
            WHERE user_id = ? AND is_active = 1 AND end_date > ?
            ORDER BY end_date DESC LIMIT 1
            """,
            (user_id, now),
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else None


# ══════════════════════════════════════════════════════════════════════════════
#  Payments
# ══════════════════════════════════════════════════════════════════════════════

async def create_payment(
    user_id: int,
    method: str,
    amount: float,
    currency: str,
    proof: Optional[str] = None,
) -> int:
    """Insert a pending payment and return its id."""
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            """
            INSERT INTO payments (user_id, method, amount, currency, proof)
            VALUES (?, ?, ?, ?, ?)
            """,
            (user_id, method, amount, currency, proof),
        )
        await db.commit()
        return cur.lastrowid


async def update_payment_proof(payment_id: int, proof: str) -> None:
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE payments SET proof = ? WHERE id = ?",
            (proof, payment_id),
        )
        await db.commit()


async def update_payment_status(payment_id: int, status: str) -> Optional[int]:
    """Update payment status and return user_id for notification."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE payments SET status = ? WHERE id = ?",
            (status, payment_id),
        )
        await db.commit()
        async with db.execute(
            "SELECT user_id FROM payments WHERE id = ?", (payment_id,)
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else None


async def get_pending_payment_by_user(user_id: int) -> Optional[Dict[str, Any]]:
    """Get the most recent pending payment for a user."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            """
            SELECT * FROM payments
            WHERE user_id = ? AND status = 'pending'
            ORDER BY created_at DESC LIMIT 1
            """,
            (user_id,),
        ) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None


async def approve_payment_by_user(user_id: int) -> bool:
    """Approve the most recent pending payment for a user. Returns True if found."""
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            """
            SELECT id FROM payments
            WHERE user_id = ? AND status = 'pending'
            ORDER BY created_at DESC LIMIT 1
            """,
            (user_id,),
        ) as cur:
            row = await cur.fetchone()
        if not row:
            return False
        await db.execute(
            "UPDATE payments SET status = 'approved' WHERE id = ?", (row[0],)
        )
        await db.commit()
        return True


# ══════════════════════════════════════════════════════════════════════════════
#  Share submissions
# ══════════════════════════════════════════════════════════════════════════════

async def create_share_submission(
    user_id: int,
    platform: str,
    proof: str,
    proof_type: str,
) -> int:
    """Insert a pending share submission and return its id."""
    async with aiosqlite.connect(DB_PATH) as db:
        cur = await db.execute(
            """
            INSERT INTO share_submissions (user_id, platform, proof, proof_type)
            VALUES (?, ?, ?, ?)
            """,
            (user_id, platform, proof, proof_type),
        )
        await db.commit()
        return cur.lastrowid


async def get_approved_share_count(user_id: int) -> int:
    """Count how many share submissions have been approved for this user."""
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            """
            SELECT COUNT(*) FROM share_submissions
            WHERE user_id = ? AND status = 'approved'
            """,
            (user_id,),
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else 0


async def get_pending_share_submission(submission_id: int) -> Optional[Dict[str, Any]]:
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT * FROM share_submissions WHERE id = ?", (submission_id,)
        ) as cur:
            row = await cur.fetchone()
            return dict(row) if row else None


async def update_share_status(submission_id: int, status: str) -> Optional[int]:
    """Update submission status and return user_id."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "UPDATE share_submissions SET status = ? WHERE id = ?",
            (status, submission_id),
        )
        await db.commit()
        async with db.execute(
            "SELECT user_id FROM share_submissions WHERE id = ?", (submission_id,)
        ) as cur:
            row = await cur.fetchone()
            return row[0] if row else None


async def get_pending_shares_for_user(user_id: int) -> List[Dict[str, Any]]:
    """Get all pending share submissions for a user."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            """
            SELECT * FROM share_submissions
            WHERE user_id = ? AND status = 'pending'
            ORDER BY submitted_at ASC
            """,
            (user_id,),
        ) as cur:
            rows = await cur.fetchall()
            return [dict(r) for r in rows]