database.py

"""
DetruAI — Production Database Layer v1.0
Uses SQLite — works perfectly on Hugging Face Spaces with no external dependencies.

Tables:
  users         — accounts, roles, password hashes
  subscriptions — Stripe subscription state per user
  scans         — rate-limit counter + full scan history
  api_keys      — Pro API keys
"""

import os
import sqlite3
import secrets
import hashlib
import time
from datetime import datetime, date
from contextlib import contextmanager

# ── DB file lives next to app.py ─────────────────────────────────────────────
DB_PATH = os.path.join(os.path.dirname(__file__), 'detruai.db')

# ── Tier limits ───────────────────────────────────────────────────────────────
TIER_LIMITS = {
    'free':       {'image_daily': 10,  'video': False, 'audio': False,
                   'batch': False, 'heatmap': False, 'history_days': 1,
                   'api': False, 'max_image_mb': 5,  'watermark': True},
    'pro':        {'image_daily': None, 'video': True,  'audio': True,
                   'batch': True,  'heatmap': True,  'history_days': 90,
                   'api': True,  'max_image_mb': 50, 'watermark': False},
    'enterprise': {'image_daily': None, 'video': True,  'audio': True,
                   'batch': True,  'heatmap': True,  'history_days': 9999,
                   'api': True,  'max_image_mb': 100,'watermark': False},
    # guest = same as free but no history at all
    'guest':      {'image_daily': 3,   'video': False, 'audio': False,
                   'batch': False, 'heatmap': False, 'history_days': 0,
                   'api': False, 'max_image_mb': 2,  'watermark': True},
}


# ══════════════════════════════════════════════════════════════════════════════
# CONNECTION
# ══════════════════════════════════════════════════════════════════════════════

@contextmanager
def _db():
    """Thread-safe SQLite connection with WAL mode for concurrent reads."""
    conn = sqlite3.connect(DB_PATH, timeout=10, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


# ══════════════════════════════════════════════════════════════════════════════
# SCHEMA
# ══════════════════════════════════════════════════════════════════════════════

def init_db():
    """Create all tables if they don't exist. Safe to call on every startup."""
    with _db() as conn:
        conn.executescript("""
        CREATE TABLE IF NOT EXISTS users (
            username        TEXT PRIMARY KEY,
            email           TEXT UNIQUE NOT NULL,
            password_hash   TEXT NOT NULL,
            role            TEXT NOT NULL DEFAULT 'free',
            created_at      TEXT NOT NULL DEFAULT (datetime('now')),
            reset_token     TEXT,
            reset_expires   TEXT,
            last_login      TEXT
        );

        CREATE TABLE IF NOT EXISTS subscriptions (
            username            TEXT PRIMARY KEY REFERENCES users(username),
            stripe_customer_id  TEXT,
            stripe_sub_id       TEXT,
            plan                TEXT NOT NULL DEFAULT 'free',
            status              TEXT NOT NULL DEFAULT 'inactive',
            current_period_end  TEXT,
            cancel_at_period_end INTEGER DEFAULT 0,
            updated_at          TEXT NOT NULL DEFAULT (datetime('now'))
        );

        CREATE TABLE IF NOT EXISTS daily_usage (
            username    TEXT NOT NULL,
            use_date    TEXT NOT NULL,
            scan_count  INTEGER NOT NULL DEFAULT 0,
            PRIMARY KEY (username, use_date),
            FOREIGN KEY (username) REFERENCES users(username)
        );

        CREATE TABLE IF NOT EXISTS scans (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            username    TEXT NOT NULL,
            scan_type   TEXT NOT NULL,
            filename    TEXT,
            label       TEXT,
            confidence  REAL,
            result_json TEXT,
            created_at  TEXT NOT NULL DEFAULT (datetime('now')),
            FOREIGN KEY (username) REFERENCES users(username)
        );

        CREATE TABLE IF NOT EXISTS api_keys (
            key_hash    TEXT PRIMARY KEY,
            username    TEXT NOT NULL,
            label       TEXT,
            created_at  TEXT NOT NULL DEFAULT (datetime('now')),
            last_used   TEXT,
            is_active   INTEGER NOT NULL DEFAULT 1,
            FOREIGN KEY (username) REFERENCES users(username)
        );

        CREATE INDEX IF NOT EXISTS idx_scans_user   ON scans(username, created_at);
        CREATE INDEX IF NOT EXISTS idx_usage_user   ON daily_usage(username, use_date);
        CREATE INDEX IF NOT EXISTS idx_users_email  ON users(email);
        """)
    print("[DB] Initialised — SQLite WAL mode active.")


# ══════════════════════════════════════════════════════════════════════════════
# USER MANAGEMENT
# ══════════════════════════════════════════════════════════════════════════════

def create_user(username: str, email: str, password_hash: str) -> dict:
    with _db() as conn:
        conn.execute(
            "INSERT INTO users (username, email, password_hash, role) VALUES (?,?,?,'free')",
            (username, email, password_hash)
        )
        conn.execute(
            "INSERT INTO subscriptions (username, plan, status) VALUES (?, 'free', 'active')",
            (username,)
        )
    return get_user(username)


def get_user(username: str) -> dict | None:
    with _db() as conn:
        row = conn.execute(
            """SELECT u.*, COALESCE(s.plan,'free') AS plan, s.status AS sub_status,
                      s.stripe_customer_id, s.stripe_sub_id, s.current_period_end,
                      s.cancel_at_period_end
               FROM users u
               LEFT JOIN subscriptions s ON s.username = u.username
               WHERE u.username = ?""",
            (username,)
        ).fetchone()
    return dict(row) if row else None


def get_user_by_email(email: str) -> dict | None:
    with _db() as conn:
        row = conn.execute(
            """SELECT u.*, COALESCE(s.plan,'free') AS plan
               FROM users u
               LEFT JOIN subscriptions s ON s.username = u.username
               WHERE u.email = ?""",
            (email,)
        ).fetchone()
    return dict(row) if row else None


def update_last_login(username: str):
    with _db() as conn:
        conn.execute("UPDATE users SET last_login=datetime('now') WHERE username=?", (username,))


def set_reset_token(username: str, token: str, expires: str):
    with _db() as conn:
        conn.execute(
            "UPDATE users SET reset_token=?, reset_expires=? WHERE username=?",
            (token, expires, username)
        )


def consume_reset_token(token: str) -> dict | None:
    """Find user by token if not expired. Clears token on success."""
    with _db() as conn:
        row = conn.execute(
            "SELECT * FROM users WHERE reset_token=? AND reset_expires > datetime('now')",
            (token,)
        ).fetchone()
        if not row:
            return None
        conn.execute(
            "UPDATE users SET reset_token=NULL, reset_expires=NULL WHERE username=?",
            (row['username'],)
        )
    return dict(row)


def update_password(username: str, password_hash: str):
    with _db() as conn:
        conn.execute(
            "UPDATE users SET password_hash=? WHERE username=?",
            (password_hash, username)
        )


def migrate_from_json(users_json_path: str):
    """
    One-time migration: import existing users.json accounts into SQLite.
    Safe to call on startup — skips users that already exist.
    """
    import json
    if not os.path.exists(users_json_path):
        return
    try:
        with open(users_json_path) as f:
            data = json.load(f)
        migrated = 0
        with _db() as conn:
            for email_key, u in data.items():
                uname = u.get('username', email_key)
                email = u.get('email', email_key)
                ph    = u.get('password_hash', '')
                role  = u.get('role', 'free')
                # Map old 'user' role to 'free'
                if role == 'user':
                    role = 'free'
                try:
                    conn.execute(
                        "INSERT OR IGNORE INTO users (username,email,password_hash,role,created_at) VALUES (?,?,?,?,?)",
                        (uname, email, ph, role, u.get('created','2024-01-01'))
                    )
                    conn.execute(
                        "INSERT OR IGNORE INTO subscriptions (username,plan,status) VALUES (?,?,?)",
                        (uname, role, 'active')
                    )
                    migrated += 1
                except Exception as e:
                    print(f"  [Migrate] skip {uname}: {e}")
        print(f"[DB] Migrated {migrated} user(s) from users.json → SQLite")
    except Exception as e:
        print(f"[DB] Migration error: {e}")


# ══════════════════════════════════════════════════════════════════════════════
# RATE LIMITING
# ══════════════════════════════════════════════════════════════════════════════

def get_plan(username: str) -> str:
    """Return the user's current active plan ('free', 'pro', 'enterprise', 'guest')."""
    with _db() as conn:
        row = conn.execute(
            """SELECT s.plan, s.status, s.current_period_end
               FROM subscriptions s WHERE s.username=?""",
            (username,)
        ).fetchone()
    if not row:
        return 'free'
    plan   = row['plan']
    status = row['status']
    # Check if subscription has expired
    if plan in ('pro', 'enterprise') and status == 'active':
        end = row['current_period_end']
        if end and datetime.fromisoformat(end) < datetime.utcnow():
            # Expired — downgrade
            _downgrade_expired(username)
            return 'free'
        return plan
    return 'free'


def _downgrade_expired(username: str):
    with _db() as conn:
        conn.execute(
            "UPDATE subscriptions SET plan='free', status='expired' WHERE username=?",
            (username,)
        )


def get_tier_limits(plan: str) -> dict:
    return TIER_LIMITS.get(plan, TIER_LIMITS['free'])


def check_and_increment_usage(username: str, scan_type: str = 'image') -> dict:
    """
    Check if user can perform a scan. If allowed, increment counter.
    Returns:
      {'allowed': True/False, 'used': N, 'limit': N or None, 'remaining': N or None}
    """
    plan   = get_plan(username)
    limits = get_tier_limits(plan)
    today  = date.today().isoformat()

    # Video/Audio require Pro
    if scan_type in ('video', 'audio') and not limits.get('video'):
        return {'allowed': False, 'reason': 'pro_required',
                'used': 0, 'limit': 0, 'remaining': 0,
                'plan': plan}

    # Batch requires Pro
    if scan_type == 'batch' and not limits.get('batch'):
        return {'allowed': False, 'reason': 'pro_required',
                'used': 0, 'limit': 0, 'remaining': 0,
                'plan': plan}

    daily_limit = limits.get('image_daily')

    # Unlimited plan
    if daily_limit is None:
        _record_usage(username, today)
        return {'allowed': True, 'used': None, 'limit': None,
                'remaining': None, 'plan': plan}

    # Check and increment atomically
    with _db() as conn:
        row = conn.execute(
            "SELECT scan_count FROM daily_usage WHERE username=? AND use_date=?",
            (username, today)
        ).fetchone()
        used = row['scan_count'] if row else 0

        if used >= daily_limit:
            return {'allowed': False, 'reason': 'daily_limit',
                    'used': used, 'limit': daily_limit,
                    'remaining': 0, 'plan': plan,
                    'resets': 'midnight UTC'}

        if row:
            conn.execute(
                "UPDATE daily_usage SET scan_count=scan_count+1 WHERE username=? AND use_date=?",
                (username, today)
            )
        else:
            # Ensure user row exists before inserting into daily_usage.
            # Guards against guest sessions (guest_xxxx) that have no DB row,
            # which would cause a FOREIGN KEY constraint failure.
            conn.execute(
                """INSERT OR IGNORE INTO users (username, email, password_hash, role)
                   VALUES (?, ?, '', 'guest')""",
                (username, username)
            )
            conn.execute(
                """INSERT OR IGNORE INTO subscriptions (username, plan, status)
                   VALUES (?, 'guest', 'active')""",
                (username,)
            )
            conn.execute(
                "INSERT INTO daily_usage (username, use_date, scan_count) VALUES (?,?,1)",
                (username, today)
            )
        used += 1

    return {'allowed': True, 'used': used, 'limit': daily_limit,
            'remaining': daily_limit - used, 'plan': plan}


def _record_usage(username: str, today: str):
    """Record usage for unlimited-plan users (for analytics)."""
    with _db() as conn:
        # Ensure user exists (same FK guard as check_and_increment_usage)
        conn.execute(
            """INSERT OR IGNORE INTO users (username, email, password_hash, role)
               VALUES (?, ?, '', 'guest')""",
            (username, username)
        )
        conn.execute(
            """INSERT OR IGNORE INTO subscriptions (username, plan, status)
               VALUES (?, 'guest', 'active')""",
            (username,)
        )
        conn.execute("""
            INSERT INTO daily_usage (username, use_date, scan_count) VALUES (?,?,1)
            ON CONFLICT(username, use_date) DO UPDATE SET scan_count=scan_count+1
        """, (username, today))


def get_usage_today(username: str) -> dict:
    plan   = get_plan(username)
    limits = get_tier_limits(plan)
    today  = date.today().isoformat()
    with _db() as conn:
        row = conn.execute(
            "SELECT scan_count FROM daily_usage WHERE username=? AND use_date=?",
            (username, today)
        ).fetchone()
    used        = row['scan_count'] if row else 0
    daily_limit = limits.get('image_daily')
    return {
        'plan':      plan,
        'used':      used,
        'limit':     daily_limit,
        'remaining': (daily_limit - used) if daily_limit is not None else None,
        'features':  {k: v for k, v in limits.items() if k not in ('image_daily',)}
    }


# ══════════════════════════════════════════════════════════════════════════════
# SCAN HISTORY
# ══════════════════════════════════════════════════════════════════════════════

def save_scan(username: str, scan_type: str, filename: str,
              label: str, confidence: float, result: dict):
    """Save a scan result. Enforces history retention by plan."""
    plan   = get_plan(username)
    limits = get_tier_limits(plan)
    days   = limits.get('history_days', 1)
    if days == 0:
        return   # guest — no history

    import json
    with _db() as conn:
        conn.execute(
            """INSERT INTO scans (username, scan_type, filename, label, confidence, result_json)
               VALUES (?,?,?,?,?,?)""",
            (username, scan_type, filename, label, confidence,
             json.dumps(result, default=str))
        )
        # Prune old records beyond retention window
        if days < 9999:
            conn.execute(
                """DELETE FROM scans WHERE username=?
                   AND created_at < datetime('now', ?)""",
                (username, f'-{days} days')
            )
        # Cap to 500 most recent per user
        conn.execute(
            """DELETE FROM scans WHERE username=? AND id NOT IN (
                SELECT id FROM scans WHERE username=? ORDER BY id DESC LIMIT 500
            )""",
            (username, username)
        )


def get_history(username: str, limit: int = 100) -> list:
    plan   = get_plan(username)
    limits = get_tier_limits(plan)
    if limits.get('history_days', 0) == 0:
        return []
    import json
    with _db() as conn:
        rows = conn.execute(
            """SELECT id, scan_type, filename, label, confidence, result_json, created_at
               FROM scans WHERE username=? ORDER BY id DESC LIMIT ?""",
            (username, limit)
        ).fetchall()
    result = []
    for r in rows:
        try:
            data = json.loads(r['result_json']) if r['result_json'] else {}
        except Exception:
            data = {}
        result.append({
            'id':         r['id'],
            'scan_type':  r['scan_type'],
            'filename':   r['filename'],
            'label':      r['label'],
            'confidence': r['confidence'],
            'created_at': r['created_at'],
            **data,
        })
    return result


def clear_history(username: str):
    with _db() as conn:
        conn.execute("DELETE FROM scans WHERE username=?", (username,))


# ══════════════════════════════════════════════════════════════════════════════
# SUBSCRIPTIONS (Stripe)
# ══════════════════════════════════════════════════════════════════════════════

def upsert_subscription(username: str, stripe_customer_id: str, stripe_sub_id: str,
                        plan: str, status: str, period_end: str,
                        cancel_at_period_end: bool = False):
    with _db() as conn:
        conn.execute("""
            INSERT INTO subscriptions
                (username, stripe_customer_id, stripe_sub_id, plan, status,
                 current_period_end, cancel_at_period_end, updated_at)
            VALUES (?,?,?,?,?,?,?,datetime('now'))
            ON CONFLICT(username) DO UPDATE SET
                stripe_customer_id  = excluded.stripe_customer_id,
                stripe_sub_id       = excluded.stripe_sub_id,
                plan                = excluded.plan,
                status              = excluded.status,
                current_period_end  = excluded.current_period_end,
                cancel_at_period_end= excluded.cancel_at_period_end,
                updated_at          = datetime('now')
        """, (username, stripe_customer_id, stripe_sub_id, plan, status,
              period_end, int(cancel_at_period_end)))


def get_subscription(username: str) -> dict | None:
    with _db() as conn:
        row = conn.execute(
            "SELECT * FROM subscriptions WHERE username=?", (username,)
        ).fetchone()
    return dict(row) if row else None


def cancel_subscription(username: str):
    """Downgrade to free immediately (for webhook-triggered cancellations)."""
    with _db() as conn:
        conn.execute(
            "UPDATE subscriptions SET plan='free', status='cancelled', updated_at=datetime('now') WHERE username=?",
            (username,)
        )


# ══════════════════════════════════════════════════════════════════════════════
# API KEYS
# ══════════════════════════════════════════════════════════════════════════════

def create_api_key(username: str, label: str = 'Default') -> str:
    """Generate a new API key for a Pro user. Returns the raw key (shown once)."""
    raw_key  = 'dkai_' + secrets.token_urlsafe(32)
    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
    with _db() as conn:
        conn.execute(
            "INSERT INTO api_keys (key_hash, username, label) VALUES (?,?,?)",
            (key_hash, username, label)
        )
    return raw_key


def validate_api_key(raw_key: str) -> dict | None:
    """Returns user dict if key is valid + active Pro/Enterprise user."""
    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
    with _db() as conn:
        row = conn.execute(
            """SELECT ak.username, ak.label, u.role, COALESCE(s.plan,'free') as plan
               FROM api_keys ak
               JOIN users u ON u.username = ak.username
               LEFT JOIN subscriptions s ON s.username = ak.username
               WHERE ak.key_hash=? AND ak.is_active=1""",
            (key_hash,)
        ).fetchone()
        if not row:
            return None
        plan = row['plan']
        if plan not in ('pro', 'enterprise'):
            return None
        conn.execute(
            "UPDATE api_keys SET last_used=datetime('now') WHERE key_hash=?",
            (key_hash,)
        )
    return dict(row)


def list_api_keys(username: str) -> list:
    with _db() as conn:
        rows = conn.execute(
            "SELECT label, created_at, last_used, is_active FROM api_keys WHERE username=? ORDER BY created_at DESC",
            (username,)
        ).fetchall()
    return [dict(r) for r in rows]


def revoke_api_key(username: str, label: str):
    with _db() as conn:
        conn.execute(
            "UPDATE api_keys SET is_active=0 WHERE username=? AND label=?",
            (username, label)
        )


# ══════════════════════════════════════════════════════════════════════════════
# ADMIN ANALYTICS
# ══════════════════════════════════════════════════════════════════════════════

def admin_stats() -> dict:
    with _db() as conn:
        total_users  = conn.execute("SELECT COUNT(*) FROM users").fetchone()[0]
        pro_users    = conn.execute(
            "SELECT COUNT(*) FROM subscriptions WHERE plan='pro' AND status='active'"
        ).fetchone()[0]
        scans_today  = conn.execute(
            "SELECT COALESCE(SUM(scan_count),0) FROM daily_usage WHERE use_date=date('now')"
        ).fetchone()[0]
        total_scans  = conn.execute("SELECT COUNT(*) FROM scans").fetchone()[0]
        fake_count   = conn.execute("SELECT COUNT(*) FROM scans WHERE label='FAKE'").fetchone()[0]
    return {
        'total_users':  total_users,
        'pro_users':    pro_users,
        'free_users':   total_users - pro_users,
        'scans_today':  scans_today,
        'total_scans':  total_scans,
        'fake_rate':    round(fake_count / max(total_scans, 1) * 100, 1),
    }