database.py

"""
OmniParse AI - Database Layer

Dual-backend persistence:
  - **Supabase** (PostgreSQL) when SUPABASE_URL and SUPABASE_KEY are set.
  - **SQLite** as the fallback (file-based, suitable for HF Spaces).

All public functions return plain dicts / lists — never raw Row / tuple objects.
"""

import json
import os
import sqlite3
import time
import uuid

from config import USE_SUPABASE, supabase_client, SESSION_TTL_SECONDS

# ---------------------------------------------------------------------------
# SQLite setup
# ---------------------------------------------------------------------------

_DB_PATH = os.environ.get("SQLITE_DB_PATH", "omniparse.db")


def _sqlite_conn() -> sqlite3.Connection:
    """Return a new SQLite connection with row_factory enabled."""
    conn = sqlite3.connect(_DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def _row_to_dict(row) -> dict:
    """Convert a sqlite3.Row (or similar) to a plain dict."""
    if row is None:
        return None
    if isinstance(row, dict):
        return row
    return dict(row)


# ---------------------------------------------------------------------------
# Initialisation
# ---------------------------------------------------------------------------

def db_init() -> None:
    """Create tables and indexes if they do not exist yet."""
    if USE_SUPABASE and supabase_client is not None:
        # Supabase tables are expected to be created via migrations;
        # we do not attempt DDL here.
        return

    conn = _sqlite_conn()
    try:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id            TEXT PRIMARY KEY,
                email         TEXT    NOT NULL UNIQUE,
                name          TEXT    NOT NULL,
                password      TEXT    NOT NULL DEFAULT '',
                plan          TEXT    NOT NULL DEFAULT 'free',
                stripe_cid    TEXT    DEFAULT '',
                api_key       TEXT    DEFAULT '',
                created_at    TEXT    NOT NULL DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS sessions (
                token      TEXT PRIMARY KEY,
                user_id    TEXT    NOT NULL,
                expires_at REAL    NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS invoices (
                id            TEXT PRIMARY KEY,
                user_id       TEXT    NOT NULL,
                vendor        TEXT    DEFAULT '',
                inv_number    TEXT    DEFAULT '',
                inv_date      TEXT    DEFAULT '',
                due_date      TEXT    DEFAULT '',
                amount        REAL    DEFAULT 0,
                vat_amount    REAL    DEFAULT 0,
                total         REAL    DEFAULT 0,
                currency      TEXT    DEFAULT 'USD',
                status        TEXT    DEFAULT 'done',
                is_duplicate  INTEGER DEFAULT 0,
                confidence    REAL    DEFAULT 0,
                raw_json      TEXT    DEFAULT '',
                created_at    TEXT    NOT NULL DEFAULT (datetime('now')),
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            );

            CREATE INDEX IF NOT EXISTS idx_invoices_user_created
                ON invoices(user_id, created_at);
        """)
        conn.commit()
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# User CRUD
# ---------------------------------------------------------------------------

def create_user(email: str, name: str, password: str, plan: str = "free") -> dict | None:
    """Insert a new user. Returns the user dict on success, None on failure."""
    uid = str(uuid.uuid4())

    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("users")
                .insert({
                    "id": uid,
                    "email": email,
                    "name": name,
                    "password": password,
                    "plan": plan,
                    "stripe_cid": "",
                    "api_key": "",
                })
                .execute()
            )
            if row.data and len(row.data) > 0:
                return row.data[0]
        except Exception:
            return None
        return None

    # SQLite path
    conn = _sqlite_conn()
    try:
        conn.execute(
            "INSERT INTO users (id, email, name, password, plan) VALUES (?, ?, ?, ?, ?)",
            (uid, email, name, password, plan),
        )
        conn.commit()
        cur = conn.execute("SELECT * FROM users WHERE id = ?", (uid,))
        return _row_to_dict(cur.fetchone())
    except Exception:
        return None
    finally:
        conn.close()


def get_user_by_email(email: str) -> dict | None:
    """Fetch a user by email. Returns user dict or None."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("users")
                .select("*")
                .eq("email", email)
                .limit(1)
                .execute()
            )
            if row.data and len(row.data) > 0:
                return row.data[0]
        except Exception:
            pass
        return None

    conn = _sqlite_conn()
    try:
        cur = conn.execute("SELECT * FROM users WHERE email = ?", (email,))
        return _row_to_dict(cur.fetchone())
    finally:
        conn.close()


def get_user_by_id(uid: str) -> dict | None:
    """Fetch a user by primary-key ID. Returns user dict or None."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("users")
                .select("*")
                .eq("id", uid)
                .limit(1)
                .execute()
            )
            if row.data and len(row.data) > 0:
                return row.data[0]
        except Exception:
            pass
        return None

    conn = _sqlite_conn()
    try:
        cur = conn.execute("SELECT * FROM users WHERE id = ?", (uid,))
        return _row_to_dict(cur.fetchone())
    finally:
        conn.close()


def update_user(uid: str, fields: dict) -> dict | None:
    """Update arbitrary fields on a user. Returns updated user dict or None."""
    if not fields:
        return get_user_by_id(uid)

    allowed = {"name", "password", "plan", "stripe_cid", "api_key"}
    updates = {k: v for k, v in fields.items() if k in allowed}
    if not updates:
        return get_user_by_id(uid)

    if USE_SUPABASE and supabase_client is not None:
        try:
            (
                supabase_client.table("users")
                .update(updates)
                .eq("id", uid)
                .execute()
            )
            return get_user_by_id(uid)
        except Exception:
            return None

    # SQLite path
    set_clause = ", ".join(f"{k} = ?" for k in updates)
    values = list(updates.values()) + [uid]

    conn = _sqlite_conn()
    try:
        conn.execute(f"UPDATE users SET {set_clause} WHERE id = ?", values)
        conn.commit()
        return get_user_by_id(uid)
    except Exception:
        return None
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Sessions
# ---------------------------------------------------------------------------

def create_session(user_id: str, token: str) -> None:
    """Create a new session row."""
    expires_at = time.time() + SESSION_TTL_SECONDS

    if USE_SUPABASE and supabase_client is not None:
        try:
            supabase_client.table("sessions").insert({
                "token": token,
                "user_id": user_id,
                "expires_at": expires_at,
            }).execute()
        except Exception:
            pass
        return

    conn = _sqlite_conn()
    try:
        conn.execute(
            "INSERT INTO sessions (token, user_id, expires_at) VALUES (?, ?, ?)",
            (token, user_id, expires_at),
        )
        conn.commit()
    finally:
        conn.close()


def get_session_user(token: str) -> dict | None:
    """Return the user dict for a valid, non-expired session, or None."""
    now = time.time()

    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("sessions")
                .select("user_id, expires_at")
                .eq("token", token)
                .limit(1)
                .execute()
            )
            if not row.data or len(row.data) == 0:
                return None
            session = row.data[0]
            if session["expires_at"] < now:
                # Session expired — clean up
                supabase_client.table("sessions").delete().eq("token", token).execute()
                return None
            return get_user_by_id(session["user_id"])
        except Exception:
            return None

    conn = _sqlite_conn()
    try:
        cur = conn.execute(
            "SELECT user_id, expires_at FROM sessions WHERE token = ?",
            (token,),
        )
        session = cur.fetchone()
        if session is None:
            return None
        if session["expires_at"] < now:
            conn.execute("DELETE FROM sessions WHERE token = ?", (token,))
            conn.commit()
            return None
        return get_user_by_id(session["user_id"])
    finally:
        conn.close()


def delete_session(token: str) -> None:
    """Delete a single session (used on logout)."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            supabase_client.table("sessions").delete().eq("token", token).execute()
        except Exception:
            pass
        return

    conn = _sqlite_conn()
    try:
        conn.execute("DELETE FROM sessions WHERE token = ?", (token,))
        conn.commit()
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Invoices
# ---------------------------------------------------------------------------

def insert_invoice(user_id: str, data: dict) -> dict | None:
    """Insert a parsed invoice. *data* should contain the invoice fields.

    Returns the inserted invoice dict or None on failure.
    """
    inv_id = str(uuid.uuid4())
    raw_json = json.dumps(data.get("raw_data", data), default=str)

    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("invoices")
                .insert({
                    "id": inv_id,
                    "user_id": user_id,
                    "vendor": data.get("vendor", ""),
                    "inv_number": data.get("inv_number", ""),
                    "inv_date": data.get("inv_date", ""),
                    "due_date": data.get("due_date", ""),
                    "amount": data.get("amount", 0),
                    "vat_amount": data.get("vat_amount", 0),
                    "total": data.get("total", 0),
                    "currency": data.get("currency", "USD"),
                    "status": data.get("status", "done"),
                    "is_duplicate": 1 if data.get("is_duplicate") else 0,
                    "confidence": data.get("confidence", 0),
                    "raw_json": raw_json,
                })
                .execute()
            )
            if row.data and len(row.data) > 0:
                return row.data[0]
        except Exception:
            return None
        return None

    # SQLite path
    conn = _sqlite_conn()
    try:
        conn.execute(
            """INSERT INTO invoices
               (id, user_id, vendor, inv_number, inv_date, due_date,
                amount, vat_amount, total, currency, status,
                is_duplicate, confidence, raw_json)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                inv_id,
                user_id,
                data.get("vendor", ""),
                data.get("inv_number", ""),
                data.get("inv_date", ""),
                data.get("due_date", ""),
                data.get("amount", 0),
                data.get("vat_amount", 0),
                data.get("total", 0),
                data.get("currency", "USD"),
                data.get("status", "done"),
                1 if data.get("is_duplicate") else 0,
                data.get("confidence", 0),
                raw_json,
            ),
        )
        conn.commit()
        cur = conn.execute("SELECT * FROM invoices WHERE id = ?", (inv_id,))
        return _row_to_dict(cur.fetchone())
    except Exception:
        return None
    finally:
        conn.close()


def get_invoices(user_id: str, limit: int = 50, offset: int = 0) -> list[dict]:
    """Return a list of invoice dicts for *user_id*, newest first."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("invoices")
                .select("*")
                .eq("user_id", user_id)
                .order("created_at", desc=True)
                .range(offset, offset + limit - 1)
                .execute()
            )
            return row.data if row.data else []
        except Exception:
            return []

    conn = _sqlite_conn()
    try:
        cur = conn.execute(
            "SELECT * FROM invoices WHERE user_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
            (user_id, limit, offset),
        )
        return [_row_to_dict(r) for r in cur.fetchall()]
    finally:
        conn.close()


def count_invoices_this_month(user_id: str) -> int:
    """Count how many invoices *user_id* created in the current calendar month.

    Uses an SQL WHERE clause for efficiency (no Python-side filtering).
    """
    import datetime

    now = datetime.datetime.utcnow()
    month_start = now.strftime("%Y-%m-01 00:00:00")

    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("invoices")
                .select("id", count="exact")
                .eq("user_id", user_id)
                .gte("created_at", month_start)
                .execute()
            )
            return row.count if hasattr(row, "count") else len(row.data)
        except Exception:
            return 0

    conn = _sqlite_conn()
    try:
        cur = conn.execute(
            "SELECT COUNT(*) AS cnt FROM invoices WHERE user_id = ? AND created_at >= ?",
            (user_id, month_start),
        )
        row = cur.fetchone()
        return row["cnt"] if row else 0
    finally:
        conn.close()


def check_duplicate(user_id: str, vendor: str, total: float) -> bool:
    """Return True if an invoice with the same vendor and total already exists."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("invoices")
                .select("id")
                .eq("user_id", user_id)
                .eq("vendor", vendor)
                .eq("total", total)
                .limit(1)
                .execute()
            )
            return bool(row.data and len(row.data) > 0)
        except Exception:
            return False

    conn = _sqlite_conn()
    try:
        cur = conn.execute(
            "SELECT id FROM invoices WHERE user_id = ? AND vendor = ? AND total = ? LIMIT 1",
            (user_id, vendor, total),
        )
        return cur.fetchone() is not None
    finally:
        conn.close()


def delete_invoice(user_id: str, invoice_id: str) -> bool:
    """Delete an invoice belonging to *user_id*. Returns True if a row was removed."""
    if USE_SUPABASE and supabase_client is not None:
        try:
            row = (
                supabase_client.table("invoices")
                .delete()
                .eq("id", invoice_id)
                .eq("user_id", user_id)
                .execute()
            )
            return bool(row.data and len(row.data) > 0)
        except Exception:
            return False

    conn = _sqlite_conn()
    try:
        cur = conn.execute(
            "DELETE FROM invoices WHERE id = ? AND user_id = ?",
            (invoice_id, user_id),
        )
        conn.commit()
        return cur.rowcount > 0
    finally:
        conn.close()