backend/storage.py

from __future__ import annotations

import json
import os
import sqlite3
import threading
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Sequence


DEFAULT_STORAGE_DIR = Path("/data/hackathon-space-recommender")


def _default_storage_dir() -> Path:
    custom = os.getenv("BSQF_STORAGE_DIR", "").strip()
    if custom:
        candidate = Path(custom).expanduser()
        if candidate.is_absolute() and str(candidate).startswith("/data/"):
            return candidate
        raise RuntimeError("BSQF_STORAGE_DIR must point inside the HF bucket mount under /data/")

    return DEFAULT_STORAGE_DIR


STORAGE_DIR = _default_storage_dir()
DB_PATH = STORAGE_DIR / "build_small_quest_finder.sqlite3"
DB_LOCK = threading.RLock()


def _ensure_parent(path: Path) -> None:
    if not Path("/data").exists():
        raise RuntimeError(
            "HF Storage Bucket mount not found at /data. "
            "This app is bucket-only and must run inside a Hugging Face Space with the bucket mounted."
        )
    path.parent.mkdir(parents=True, exist_ok=True)


@contextmanager
def get_connection():
    with DB_LOCK:
        _ensure_parent(DB_PATH)
        conn = sqlite3.connect(
            DB_PATH,
            timeout=30,
            check_same_thread=False,
            isolation_level=None,
        )
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA busy_timeout = 30000")
        conn.execute("PRAGMA foreign_keys = ON")
        conn.execute("PRAGMA journal_mode = DELETE")
        conn.execute("PRAGMA synchronous = NORMAL")
        try:
            yield conn
        finally:
            conn.close()


@contextmanager
def _transaction(conn: sqlite3.Connection):
    conn.execute("BEGIN IMMEDIATE")
    try:
        yield
        conn.commit()
    except Exception:
        conn.rollback()
        raise


def init_db() -> None:
    with get_connection() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS llm_cache (
                cache_key TEXT PRIMARY KEY,
                cache_value TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS quest_feedback (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                repo_id TEXT NOT NULL,
                space_title TEXT NOT NULL,
                track TEXT NOT NULL,
                question TEXT NOT NULL,
                answer TEXT NOT NULL,
                source_query TEXT,
                owner_session_id TEXT NOT NULL DEFAULT '',
                updated_at TEXT NOT NULL DEFAULT '',
                deleted_at TEXT,
                created_at TEXT NOT NULL
            )
            """
        )
        existing_columns = {row["name"] for row in conn.execute("PRAGMA table_info(quest_feedback)").fetchall()}
        if "owner_session_id" not in existing_columns:
            conn.execute("ALTER TABLE quest_feedback ADD COLUMN owner_session_id TEXT NOT NULL DEFAULT ''")
        if "updated_at" not in existing_columns:
            conn.execute("ALTER TABLE quest_feedback ADD COLUMN updated_at TEXT NOT NULL DEFAULT ''")
        if "deleted_at" not in existing_columns:
            conn.execute("ALTER TABLE quest_feedback ADD COLUMN deleted_at TEXT")
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS app_state (
                state_key TEXT PRIMARY KEY,
                state_value TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS space_metrics (
                repo_id TEXT PRIMARY KEY,
                likes INTEGER NOT NULL DEFAULT 0,
                synced_at TEXT NOT NULL DEFAULT '',
                updated_at TEXT NOT NULL
            )
            """
        )
        conn.commit()


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def cache_get(cache_key: str) -> str | None:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT cache_value FROM llm_cache WHERE cache_key = ?",
            (cache_key,),
        ).fetchone()
        return row["cache_value"] if row else None


def cache_set(cache_key: str, cache_value: str) -> None:
    with get_connection() as conn:
        with _transaction(conn):
            conn.execute(
                """
                INSERT INTO llm_cache (cache_key, cache_value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(cache_key) DO UPDATE SET
                    cache_value = excluded.cache_value,
                    updated_at = excluded.updated_at
                """,
                (cache_key, cache_value, now_iso()),
            )


def state_get(state_key: str) -> str | None:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT state_value FROM app_state WHERE state_key = ?",
            (state_key,),
        ).fetchone()
        return row["state_value"] if row else None


def state_set(state_key: str, state_value: str) -> None:
    with get_connection() as conn:
        with _transaction(conn):
            conn.execute(
                """
                INSERT INTO app_state (state_key, state_value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(state_key) DO UPDATE SET
                    state_value = excluded.state_value,
                    updated_at = excluded.updated_at
                """,
                (state_key, state_value, now_iso()),
            )


def save_feedback(
    repo_id: str,
    space_title: str,
    track: str,
    question: str,
    answer: str,
    source_query: str = "",
    owner_session_id: str = "",
) -> dict:
    created_at = now_iso()
    record = {
        "repo_id": repo_id,
        "space_title": space_title,
        "track": track,
        "question": question,
        "answer": answer,
        "source_query": source_query,
        "owner_session_id": owner_session_id,
        "created_at": created_at,
        "updated_at": created_at,
        "deleted_at": None,
    }
    with get_connection() as conn:
        with _transaction(conn):
            cursor = conn.execute(
                """
                INSERT INTO quest_feedback (repo_id, space_title, track, question, answer, source_query, owner_session_id, updated_at, deleted_at, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    repo_id,
                    space_title,
                    track,
                    question,
                    answer,
                    source_query,
                    owner_session_id,
                    record["updated_at"],
                    None,
                    record["created_at"],
                ),
            )
        record["id"] = int(cursor.lastrowid or 0)
    return record


def get_recent_feedback(limit: int = 10, repo_id: str | None = None) -> list[dict]:
    query = """
        SELECT id, repo_id, space_title, track, question, answer, source_query,
               owner_session_id, updated_at, deleted_at, created_at
        FROM quest_feedback
        WHERE COALESCE(deleted_at, '') = ''
    """
    params: list[object] = []
    if repo_id:
        query += " AND repo_id = ?"
        params.append(repo_id)
    query += " ORDER BY id DESC LIMIT ?"
    params.append(limit)
    try:
        with get_connection() as conn:
            db_rows = conn.execute(query, tuple(params)).fetchall()
    except Exception:
        return []
    return list(reversed([dict(row) for row in db_rows]))


def get_space_metrics(repo_ids: Sequence[str] | None = None) -> dict[str, dict]:
    query = "SELECT repo_id, likes, synced_at, updated_at FROM space_metrics"
    params: tuple[object, ...] = ()
    if repo_ids is not None:
        cleaned = [str(repo_id).strip() for repo_id in repo_ids if str(repo_id).strip()]
        if not cleaned:
            return {}
        placeholders = ",".join("?" for _ in cleaned)
        query += f" WHERE repo_id IN ({placeholders})"
        params = tuple(cleaned)

    try:
        with get_connection() as conn:
            rows = conn.execute(query, params).fetchall()
    except Exception:
        return {}
    return {str(row["repo_id"]): dict(row) for row in rows}


def upsert_space_metrics(rows: Sequence[dict]) -> None:
    cleaned_rows: list[tuple[str, int, str, str]] = []
    for row in rows:
        repo_id = str(row.get("repo_id", "")).strip()
        if not repo_id:
            continue
        likes = int(row.get("likes", 0) or 0)
        synced_at = str(row.get("synced_at", "") or now_iso())
        updated_at = str(row.get("updated_at", "") or synced_at)
        cleaned_rows.append((repo_id, likes, synced_at, updated_at))

    if not cleaned_rows:
        return

    with get_connection() as conn:
        with _transaction(conn):
            conn.executemany(
                """
                INSERT INTO space_metrics (repo_id, likes, synced_at, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(repo_id) DO UPDATE SET
                    likes = excluded.likes,
                    synced_at = excluded.synced_at,
                    updated_at = excluded.updated_at
                """,
                cleaned_rows,
            )


def update_feedback_comment(comment_id: int, owner_session_id: str, answer: str) -> bool:
    answer = (answer or "").strip()
    owner_session_id = (owner_session_id or "").strip()
    if not comment_id or not answer or not owner_session_id:
        return False

    updated_at = now_iso()
    with get_connection() as conn:
        with _transaction(conn):
            row = conn.execute(
                """
                SELECT id FROM quest_feedback
                WHERE id = ? AND owner_session_id = ? AND COALESCE(deleted_at, '') = ''
                """,
                (comment_id, owner_session_id),
            ).fetchone()
            if not row:
                return False
            conn.execute(
                """
                UPDATE quest_feedback
                SET answer = ?, updated_at = ?
                WHERE id = ? AND owner_session_id = ? AND COALESCE(deleted_at, '') = ''
                """,
                (answer, updated_at, comment_id, owner_session_id),
            )
    return True


def delete_feedback_comment(comment_id: int, owner_session_id: str) -> bool:
    owner_session_id = (owner_session_id or "").strip()
    if not comment_id or not owner_session_id:
        return False

    deleted_at = now_iso()
    with get_connection() as conn:
        with _transaction(conn):
            row = conn.execute(
                """
                SELECT id FROM quest_feedback
                WHERE id = ? AND owner_session_id = ? AND COALESCE(deleted_at, '') = ''
                """,
                (comment_id, owner_session_id),
            ).fetchone()
            if not row:
                return False
            conn.execute(
                """
                UPDATE quest_feedback
                SET deleted_at = ?, updated_at = ?
                WHERE id = ? AND owner_session_id = ? AND COALESCE(deleted_at, '') = ''
                """,
                (deleted_at, deleted_at, comment_id, owner_session_id),
            )
    return True


def dump_json(table_name: str) -> str:
    if table_name == "quest_feedback":
        rows = get_recent_feedback(limit=100000)
    else:
        rows = []
    return json.dumps(rows, ensure_ascii=False, indent=2)