app/db.py

"""Postgres persistence (Aiven). Stores users, API keys, sessions, and every message.

Connection string comes from DATABASE_URL (env). All access goes through a small
pool. If DATABASE_URL is unset, available() is False and the server runs without
persistence (in-memory only) so local dev still works.
"""
from __future__ import annotations

from contextlib import contextmanager
from datetime import datetime, timezone

import psycopg2
import psycopg2.extras
from psycopg2.pool import SimpleConnectionPool

from . import config

_pool: SimpleConnectionPool | None = None


def available() -> bool:
    return bool(config.DATABASE_URL)


def init_pool() -> None:
    global _pool
    if _pool is None and available():
        _pool = SimpleConnectionPool(1, 8, dsn=config.DATABASE_URL)


@contextmanager
def cursor(commit: bool = False):
    if _pool is None:
        init_pool()
    if _pool is None:
        raise RuntimeError("DATABASE_URL not configured")
    conn = _pool.getconn()
    try:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            yield cur
        if commit:
            conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        _pool.putconn(conn)


SCHEMA = """
CREATE TABLE IF NOT EXISTS app_users (
    id            SERIAL PRIMARY KEY,
    username      TEXT UNIQUE NOT NULL,
    password_hash TEXT NOT NULL,
    role          TEXT NOT NULL DEFAULT 'user',
    created_at    TIMESTAMPTZ NOT NULL DEFAULT now()
);

CREATE TABLE IF NOT EXISTS login_sessions (
    token      TEXT PRIMARY KEY,
    user_id    INTEGER REFERENCES app_users(id) ON DELETE CASCADE,
    role       TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
    expires_at TIMESTAMPTZ NOT NULL
);

CREATE TABLE IF NOT EXISTS api_keys (
    id            SERIAL PRIMARY KEY,
    key_hash      TEXT UNIQUE NOT NULL,
    key_prefix    TEXT NOT NULL,
    label         TEXT,
    role          TEXT NOT NULL DEFAULT 'user',
    scopes        TEXT[] NOT NULL DEFAULT ARRAY['chat','websearch']::TEXT[],
    created_by    INTEGER REFERENCES app_users(id) ON DELETE SET NULL,
    created_at    TIMESTAMPTZ NOT NULL DEFAULT now(),
    expires_at    TIMESTAMPTZ,                 -- NULL = lifetime
    revoked       BOOLEAN NOT NULL DEFAULT FALSE,
    request_count INTEGER NOT NULL DEFAULT 0,
    token_count   BIGINT  NOT NULL DEFAULT 0,
    last_used_at  TIMESTAMPTZ
);

CREATE TABLE IF NOT EXISTS chat_sessions (
    session_id      TEXT PRIMARY KEY,
    conversation_id TEXT,
    api_key_id      INTEGER REFERENCES api_keys(id) ON DELETE SET NULL,
    created_at      TIMESTAMPTZ NOT NULL DEFAULT now(),
    last_active     TIMESTAMPTZ NOT NULL DEFAULT now(),
    turns           INTEGER NOT NULL DEFAULT 0,
    context_tokens  INTEGER NOT NULL DEFAULT 0
);

CREATE TABLE IF NOT EXISTS messages (
    id          BIGSERIAL PRIMARY KEY,
    session_id  TEXT REFERENCES chat_sessions(session_id) ON DELETE CASCADE,
    api_key_id  INTEGER,
    role        TEXT NOT NULL,
    content     TEXT,
    tokens      INTEGER NOT NULL DEFAULT 0,
    has_image   BOOLEAN NOT NULL DEFAULT FALSE,
    created_at  TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);

CREATE TABLE IF NOT EXISTS app_settings (
    key        TEXT PRIMARY KEY,
    value      TEXT,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);

-- per-key session cap (NULL/0 = unlimited)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS max_sessions INTEGER;

-- per-user AGENTS.md context, applied only to this key's workspace (NULL = none)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS agents_md TEXT;
"""


def count_sessions(api_key_id: int) -> int:
    with cursor() as cur:
        cur.execute("SELECT count(*) AS n FROM chat_sessions WHERE api_key_id=%s", (api_key_id,))
        return int(cur.fetchone()["n"])


def session_exists(session_id: str) -> bool:
    with cursor() as cur:
        cur.execute("SELECT 1 FROM chat_sessions WHERE session_id=%s", (session_id,))
        return cur.fetchone() is not None


def get_setting(key: str) -> str | None:
    with cursor() as cur:
        cur.execute("SELECT value FROM app_settings WHERE key=%s", (key,))
        row = cur.fetchone()
        return row["value"] if row else None


def set_setting(key: str, value: str) -> None:
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO app_settings (key, value, updated_at) VALUES (%s,%s,now())
               ON CONFLICT (key) DO UPDATE SET value=EXCLUDED.value, updated_at=now()""",
            (key, value),
        )


def init_schema() -> None:
    with cursor(commit=True) as cur:
        cur.execute(SCHEMA)


# ---------- users / superadmin ----------
def upsert_user(username: str, password_hash: str, role: str) -> dict:
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO app_users (username, password_hash, role) VALUES (%s,%s,%s)
               ON CONFLICT (username) DO UPDATE SET password_hash=EXCLUDED.password_hash, role=EXCLUDED.role
               RETURNING id, username, role""",
            (username, password_hash, role),
        )
        return dict(cur.fetchone())


def get_user(username: str) -> dict | None:
    with cursor() as cur:
        cur.execute("SELECT * FROM app_users WHERE username=%s", (username,))
        row = cur.fetchone()
        return dict(row) if row else None


# ---------- login sessions ----------
def create_login(token: str, user_id: int, role: str, expires_at: datetime) -> None:
    with cursor(commit=True) as cur:
        cur.execute(
            "INSERT INTO login_sessions (token, user_id, role, expires_at) VALUES (%s,%s,%s,%s)",
            (token, user_id, role, expires_at),
        )


def get_login(token: str) -> dict | None:
    with cursor() as cur:
        cur.execute("SELECT * FROM login_sessions WHERE token=%s", (token,))
        row = cur.fetchone()
        if not row:
            return None
        if row["expires_at"] <= datetime.now(timezone.utc):
            return None  # expired
        return dict(row)


# ---------- api keys ----------
def create_api_key(key_hash: str, key_prefix: str, label: str, role: str,
                   scopes: list[str], created_by: int | None, expires_at: datetime | None,
                   max_sessions: int | None = None, agents_md: str | None = None) -> dict:
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO api_keys (key_hash, key_prefix, label, role, scopes, created_by, expires_at, max_sessions, agents_md)
               VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)
               RETURNING id, key_prefix, label, role, scopes, expires_at, created_at, max_sessions""",
            (key_hash, key_prefix, label, role, scopes, created_by, expires_at, max_sessions, agents_md),
        )
        return dict(cur.fetchone())


def get_api_key(key_id: int) -> dict | None:
    with cursor() as cur:
        cur.execute("SELECT * FROM api_keys WHERE id=%s", (key_id,))
        row = cur.fetchone()
        return dict(row) if row else None


def set_api_key_agents_md(key_id: int, agents_md: str | None) -> bool:
    with cursor(commit=True) as cur:
        cur.execute("UPDATE api_keys SET agents_md=%s WHERE id=%s", (agents_md, key_id))
        return cur.rowcount > 0


def get_api_key_by_hash(key_hash: str) -> dict | None:
    with cursor() as cur:
        cur.execute("SELECT * FROM api_keys WHERE key_hash=%s", (key_hash,))
        row = cur.fetchone()
        return dict(row) if row else None


def touch_api_key(key_id: int, tokens: int) -> None:
    with cursor(commit=True) as cur:
        cur.execute(
            """UPDATE api_keys SET request_count=request_count+1, token_count=token_count+%s,
               last_used_at=now() WHERE id=%s""",
            (tokens, key_id),
        )


def list_api_keys() -> list[dict]:
    with cursor() as cur:
        cur.execute(
            """SELECT id, key_prefix, label, role, scopes, created_at, expires_at, revoked,
                      request_count, token_count, last_used_at, max_sessions,
                      (agents_md IS NOT NULL AND agents_md <> '') AS has_agents,
                      (SELECT count(*) FROM chat_sessions s WHERE s.api_key_id = api_keys.id) AS sessions_used
               FROM api_keys ORDER BY created_at DESC"""
        )
        return [dict(r) for r in cur.fetchall()]


def revoke_api_key(key_id: int) -> bool:
    with cursor(commit=True) as cur:
        cur.execute("UPDATE api_keys SET revoked=TRUE WHERE id=%s", (key_id,))
        return cur.rowcount > 0


# ---------- chat sessions + messages ----------
def upsert_session(session_id: str, conversation_id: str | None, api_key_id: int | None,
                   turns: int, context_tokens: int) -> None:
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO chat_sessions (session_id, conversation_id, api_key_id, turns, context_tokens)
               VALUES (%s,%s,%s,%s,%s)
               ON CONFLICT (session_id) DO UPDATE
               SET conversation_id=EXCLUDED.conversation_id, last_active=now(),
                   turns=EXCLUDED.turns, context_tokens=EXCLUDED.context_tokens""",
            (session_id, conversation_id, api_key_id, turns, context_tokens),
        )


def insert_message(session_id: str | None, api_key_id: int | None, role: str,
                   content: str, tokens: int, has_image: bool = False) -> None:
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO messages (session_id, api_key_id, role, content, tokens, has_image)
               VALUES (%s,%s,%s,%s,%s,%s)""",
            (session_id, api_key_id, role, content, tokens, has_image),
        )


def reserve_session(session_id: str, api_key_id: int | None) -> bool:
    """Create an empty session row (no conversation yet). No-op if it already exists."""
    with cursor(commit=True) as cur:
        cur.execute(
            """INSERT INTO chat_sessions (session_id, api_key_id) VALUES (%s,%s)
               ON CONFLICT (session_id) DO NOTHING""",
            (session_id, api_key_id),
        )
        return cur.rowcount > 0


def get_session(session_id: str) -> dict | None:
    with cursor() as cur:
        cur.execute("SELECT * FROM chat_sessions WHERE session_id=%s", (session_id,))
        row = cur.fetchone()
        return dict(row) if row else None


def list_sessions_for_key(api_key_id: int) -> list[dict]:
    with cursor() as cur:
        cur.execute(
            """SELECT session_id, conversation_id, created_at, last_active, turns, context_tokens
               FROM chat_sessions WHERE api_key_id=%s ORDER BY last_active DESC""",
            (api_key_id,),
        )
        return [dict(r) for r in cur.fetchall()]


def delete_session_for_key(session_id: str, api_key_id: int) -> bool:
    """Delete a session ONLY if it belongs to this key (ownership enforced)."""
    with cursor(commit=True) as cur:
        cur.execute("DELETE FROM chat_sessions WHERE session_id=%s AND api_key_id=%s",
                    (session_id, api_key_id))
        return cur.rowcount > 0


def list_all_sessions(api_key_id: int | None = None) -> list[dict]:
    """Admin view: every session (optionally filtered by key), with the owner label."""
    with cursor() as cur:
        base = (
            """SELECT s.session_id, s.conversation_id, s.api_key_id, s.turns, s.context_tokens,
                      s.created_at, s.last_active, k.label AS owner_label, k.key_prefix
               FROM chat_sessions s LEFT JOIN api_keys k ON k.id = s.api_key_id """
        )
        if api_key_id is not None:
            cur.execute(base + "WHERE s.api_key_id=%s ORDER BY s.last_active DESC", (api_key_id,))
        else:
            cur.execute(base + "ORDER BY s.last_active DESC")
        return [dict(r) for r in cur.fetchall()]


def delete_session(session_id: str) -> bool:
    """Admin: delete any session by id (messages cascade)."""
    with cursor(commit=True) as cur:
        cur.execute("DELETE FROM chat_sessions WHERE session_id=%s", (session_id,))
        return cur.rowcount > 0