database.py

import os
import hashlib
import secrets
from datetime import datetime, timezone, timedelta
from pathlib import Path

_BACKEND_DIR = Path(__file__).resolve().parent
DB_PATH = _BACKEND_DIR / "pseudogen.db"

_USE_PG = bool(os.getenv("DATABASE_URL"))
P = "%s" if _USE_PG else "?"

GUEST_DAILY_LIMIT = 10
USER_DAILY_LIMIT = 10


def get_connection():
    if _USE_PG:
        import psycopg2
        url = os.getenv("DATABASE_URL", "")
        if url.startswith("postgres://"):
            url = url.replace("postgres://", "postgresql://", 1)
        return psycopg2.connect(url)
    else:
        import sqlite3
        conn = sqlite3.connect(DB_PATH)
        conn.row_factory = sqlite3.Row
        return conn


def _one(conn, sql, params=()):
    if _USE_PG:
        import psycopg2.extras
        cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cur.execute(sql, params)
        row = cur.fetchone()
        return dict(row) if row else None
    else:
        row = conn.execute(sql, params).fetchone()
        return dict(row) if row else None


def _exec(conn, sql, params=()):
    if _USE_PG:
        cur = conn.cursor()
        cur.execute(sql, params)
        return cur
    else:
        return conn.execute(sql, params)


def init_db():
    conn = get_connection()
    try:
        if _USE_PG:
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS users (
                    id BIGSERIAL PRIMARY KEY,
                    email TEXT UNIQUE NOT NULL,
                    hashed_password TEXT NOT NULL,
                    plan TEXT NOT NULL DEFAULT 'free',
                    created_at TEXT NOT NULL DEFAULT NOW()::TEXT
                )
            """)
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS daily_usage (
                    id BIGSERIAL PRIMARY KEY,
                    identifier TEXT NOT NULL,
                    date TEXT NOT NULL,
                    count INTEGER NOT NULL DEFAULT 0,
                    UNIQUE(identifier, date)
                )
            """)
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS refresh_tokens (
                    id BIGSERIAL PRIMARY KEY,
                    token_hash TEXT UNIQUE NOT NULL,
                    user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
                    expires_at TEXT NOT NULL,
                    created_at TEXT NOT NULL DEFAULT NOW()::TEXT
                )
            """)
        else:
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS users (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    email TEXT UNIQUE NOT NULL,
                    hashed_password TEXT NOT NULL,
                    plan TEXT NOT NULL DEFAULT 'free',
                    created_at TEXT NOT NULL DEFAULT (datetime('now'))
                )
            """)
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS daily_usage (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    identifier TEXT NOT NULL,
                    date TEXT NOT NULL,
                    count INTEGER NOT NULL DEFAULT 0,
                    UNIQUE(identifier, date)
                )
            """)
            _exec(conn, """
                CREATE TABLE IF NOT EXISTS refresh_tokens (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    token_hash TEXT UNIQUE NOT NULL,
                    user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
                    expires_at TEXT NOT NULL,
                    created_at TEXT NOT NULL DEFAULT (datetime('now'))
                )
            """)
        conn.commit()
    finally:
        conn.close()


def get_user_by_email(email: str) -> dict | None:
    conn = get_connection()
    try:
        return _one(
            conn,
            f"SELECT id, email, hashed_password, plan, created_at FROM users WHERE email = {P}",
            (email.strip().lower(),),
        )
    finally:
        conn.close()


def get_user_by_id(user_id: int) -> dict | None:
    conn = get_connection()
    try:
        return _one(
            conn,
            f"SELECT id, email, plan, created_at FROM users WHERE id = {P}",
            (user_id,),
        )
    finally:
        conn.close()


def create_user(email: str, hashed_password: str, plan: str = "free") -> dict:
    conn = get_connection()
    try:
        if _USE_PG:
            row = _one(
                conn,
                f"INSERT INTO users (email, hashed_password, plan) VALUES ({P}, {P}, {P}) RETURNING id, email, plan",
                (email.strip().lower(), hashed_password, plan),
            )
            conn.commit()
            return row
        else:
            cur = _exec(
                conn,
                f"INSERT INTO users (email, hashed_password, plan) VALUES ({P}, {P}, {P})",
                (email.strip().lower(), hashed_password, plan),
            )
            conn.commit()
            return {"id": cur.lastrowid, "email": email.strip().lower(), "plan": plan}
    finally:
        conn.close()


def _today_utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def get_usage_today(identifier: str) -> int:
    conn = get_connection()
    try:
        row = _one(
            conn,
            f"SELECT count FROM daily_usage WHERE identifier = {P} AND date = {P}",
            (identifier, _today_utc()),
        )
        return row["count"] if row else 0
    finally:
        conn.close()


def increment_usage_today(identifier: str) -> int:
    today = _today_utc()
    conn = get_connection()
    try:
        _exec(
            conn,
            f"""
            INSERT INTO daily_usage (identifier, date, count)
            VALUES ({P}, {P}, 1)
            ON CONFLICT(identifier, date) DO UPDATE SET count = daily_usage.count + 1
            """,
            (identifier, today),
        )
        conn.commit()
        row = _one(
            conn,
            f"SELECT count FROM daily_usage WHERE identifier = {P} AND date = {P}",
            (identifier, today),
        )
        return row["count"] if row else 1
    finally:
        conn.close()


# ── Refresh tokens ────────────────────────────────────────────────────────────

REFRESH_TOKEN_DAYS = 7


def _hash_token(value: str) -> str:
    return hashlib.sha256(value.encode()).hexdigest()


def create_refresh_token(user_id: int) -> str:
    """Create a refresh token, persist its hash, return the raw value."""
    value = secrets.token_urlsafe(32)
    token_hash = _hash_token(value)
    expires_at = (
        datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_DAYS)
    ).isoformat()
    conn = get_connection()
    try:
        _exec(
            conn,
            f"INSERT INTO refresh_tokens (token_hash, user_id, expires_at) VALUES ({P}, {P}, {P})",
            (token_hash, user_id, expires_at),
        )
        conn.commit()
    finally:
        conn.close()
    return value


def validate_refresh_token(value: str) -> int | None:
    """Return user_id if valid and not expired, else None."""
    token_hash = _hash_token(value)
    conn = get_connection()
    try:
        row = _one(
            conn,
            f"SELECT user_id, expires_at FROM refresh_tokens WHERE token_hash = {P}",
            (token_hash,),
        )
        if not row:
            return None
        expires = datetime.fromisoformat(row["expires_at"])
        if expires.tzinfo is None:
            expires = expires.replace(tzinfo=timezone.utc)
        if datetime.now(timezone.utc) > expires:
            _exec(conn, f"DELETE FROM refresh_tokens WHERE token_hash = {P}", (token_hash,))
            conn.commit()
            return None
        return row["user_id"]
    finally:
        conn.close()


def revoke_refresh_token(value: str) -> None:
    token_hash = _hash_token(value)
    conn = get_connection()
    try:
        _exec(conn, f"DELETE FROM refresh_tokens WHERE token_hash = {P}", (token_hash,))
        conn.commit()
    finally:
        conn.close()


def revoke_all_refresh_tokens(user_id: int) -> None:
    conn = get_connection()
    try:
        _exec(conn, f"DELETE FROM refresh_tokens WHERE user_id = {P}", (user_id,))
        conn.commit()
    finally:
        conn.close()