app.py

import os
import time
import sqlite3
import tempfile
import textwrap
import threading
import zipfile
from collections import Counter
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional

import faiss
import gradio as gr
import numpy as np
import plotly.graph_objects as go
from PIL import Image, ImageDraw, ImageFont

try:
    import spaces  # HF ZeroGPU: required to declare GPU-using funcs
    ZEROGPU = True
except ImportError:
    ZEROGPU = False

    class _NoGPU:
        def GPU(self, *a, **kw):
            def deco(fn):
                return fn
            return deco

    spaces = _NoGPU()

from sentence_transformers import SentenceTransformer

try:
    import pytesseract
    OCR_AVAILABLE = True
except ImportError:
    OCR_AVAILABLE = False

try:
    from groq import Groq
    GROQ_SDK = True
except ImportError:
    GROQ_SDK = False

try:
    from google import genai as google_genai
    GEMINI_SDK = True
except ImportError:
    GEMINI_SDK = False

# Space-owner fallback keys (optional). Per-user keys take precedence.
OWNER_GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "")
OWNER_GROQ_KEY = os.environ.get("GROQ_API_KEY", "")


# ---------- config ----------
MODEL_NAME = "clip-ViT-B-32"
EMBED_DIM = 512
MAX_IMAGE_SIZE = 1024
THUMBNAIL_SIZE = 256
TOP_K_DEFAULT = 8
MAX_VISION_IMAGES = 4  # cap images sent to Gemini per turn
IS_SPACE = "SPACE_ID" in os.environ


def _resolve_base_dir() -> Path:
    """Prefer the persistent bucket at /data on HF Spaces; fall back to a tmp dir locally."""
    persistent = Path("/data")
    if persistent.exists() and os.access(persistent, os.W_OK):
        return persistent
    fallback = Path(tempfile.gettempdir()) / "memory_search"
    fallback.mkdir(parents=True, exist_ok=True)
    return fallback


BASE_DIR = _resolve_base_dir()
print(f"Data root: {BASE_DIR}  (persistent={BASE_DIR == Path('/data')})")

print("Loading CLIP…")
model = SentenceTransformer(MODEL_NAME)
print("Model loaded.")


# ---------- per-user store ----------
class Store:
    """One user's SQLite index + FAISS index + thumbnails, isolated on disk under BASE_DIR/<user_id>/."""

    def __init__(self, root: Path):
        self.root = root
        self.thumbs = root / "thumbs"
        self.db_path = root / "memory.db"
        self.index_path = root / "faiss.index"
        self.export_dir = root / "exports"
        self.root.mkdir(parents=True, exist_ok=True)
        self.thumbs.mkdir(exist_ok=True)
        self.export_dir.mkdir(exist_ok=True)
        self._init_db()
        self.index = self._load_index()
        # Guards concurrent FAISS mutations from the same user across events.
        self.lock = threading.Lock()

    def db(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        with self.db() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS items (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    type TEXT NOT NULL,
                    source TEXT,
                    text TEXT,
                    thumb_path TEXT,
                    timestamp TEXT NOT NULL,
                    embed_idx INTEGER NOT NULL UNIQUE
                )
                """
            )

    def _load_index(self) -> faiss.IndexFlatIP:
        if self.index_path.exists():
            return faiss.read_index(str(self.index_path))
        return faiss.IndexFlatIP(EMBED_DIM)

    def save_index(self) -> None:
        faiss.write_index(self.index, str(self.index_path))

    def reset(self) -> None:
        with self.db() as conn:
            conn.execute("DELETE FROM items")
        for f in self.thumbs.glob("*.jpg"):
            f.unlink()
        self.index = faiss.IndexFlatIP(EMBED_DIM)
        self.save_index()


_stores: Dict[str, Store] = {}
_stores_lock = threading.Lock()


def _sanitize_user_id(raw: str) -> str:
    # HF usernames are already URL-safe, but strip anything unexpected just in case.
    safe = "".join(c for c in raw if c.isalnum() or c in "-_")
    return safe or "anon"


def get_store(profile: Optional[gr.OAuthProfile]) -> Optional[Store]:
    """Return the caller's Store, or None if they must sign in first."""
    if profile is None:
        if IS_SPACE:
            return None  # anonymous visitors on the public Space cannot persist
        user_id = "_local"  # local dev: single implicit user
    else:
        user_id = _sanitize_user_id(profile.username)

    with _stores_lock:
        store = _stores.get(user_id)
        if store is None:
            store = Store(BASE_DIR / user_id)
            _stores[user_id] = store
        return store


LOGIN_REQUIRED_MD = (
    "🔒 **Sign in with your Hugging Face account** (button above) to save and search "
    "your personal memories. Each user gets their own private store."
)


# ---------- embeddings (GPU on ZeroGPU) ----------
@spaces.GPU
def embed_image(img: Image.Image) -> np.ndarray:
    return model.encode(img, convert_to_numpy=True, normalize_embeddings=True).astype("float32")


@spaces.GPU
def embed_text(text: str) -> np.ndarray:
    return model.encode(text, convert_to_numpy=True, normalize_embeddings=True).astype("float32")


# ---------- thumbnails ----------
def prep_image(img: Image.Image) -> Image.Image:
    img = img.convert("RGB")
    img.thumbnail((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE))
    return img


def save_thumb_from_image(img: Image.Image, key: str, thumbs_dir: Path) -> str:
    t = img.copy()
    t.thumbnail((THUMBNAIL_SIZE, THUMBNAIL_SIZE))
    p = thumbs_dir / f"{key}.jpg"
    t.save(p, "JPEG", quality=85)
    return str(p)


def render_text_thumb(text: str, key: str, thumbs_dir: Path) -> str:
    img = Image.new("RGB", (THUMBNAIL_SIZE, THUMBNAIL_SIZE), color=(245, 240, 235))
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 14)
    except Exception:
        font = ImageFont.load_default()
    snippet = text.strip().replace("\n", " ")
    if len(snippet) > 240:
        snippet = snippet[:240] + "…"
    y = 16
    for line in textwrap.wrap(snippet, width=24)[:11]:
        draw.text((14, y), line, fill=(50, 45, 40), font=font)
        y += 20
    p = thumbs_dir / f"{key}.jpg"
    img.save(p, "JPEG", quality=85)
    return str(p)


def ocr(img: Image.Image) -> str:
    if not OCR_AVAILABLE:
        return ""
    try:
        return pytesseract.image_to_string(img).strip()
    except Exception as e:
        print(f"OCR failed: {e}")
        return ""


# ---------- ingest ----------
def ingest_images(
    files,
    profile: Optional[gr.OAuthProfile] = None,
    progress=gr.Progress(),
) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    if not files:
        return "No files selected."
    added = 0
    with store.db() as conn:
        for fp in progress.tqdm(files, desc="Indexing"):
            try:
                img = prep_image(Image.open(fp))
            except Exception as e:
                print(f"Skip {fp}: {e}")
                continue
            vec = embed_image(img)
            with store.lock:
                embed_idx = store.index.ntotal
                store.index.add(vec.reshape(1, -1))
            key = f"img_{embed_idx}_{int(time.time())}"
            thumb_path = save_thumb_from_image(img, key, store.thumbs)
            text = ocr(img)
            conn.execute(
                "INSERT INTO items(type,source,text,thumb_path,timestamp,embed_idx) VALUES (?,?,?,?,?,?)",
                ("image", Path(fp).name, text, thumb_path, datetime.now().isoformat(), embed_idx),
            )
            added += 1
    with store.lock:
        store.save_index()
    return f"✅ Indexed {added} image(s). Total: {store.index.ntotal}."


def ingest_text(text: str, profile: Optional[gr.OAuthProfile] = None) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    text = (text or "").strip()
    if not text:
        return "Nothing to save."
    vec = embed_text(text)
    with store.lock:
        embed_idx = store.index.ntotal
        store.index.add(vec.reshape(1, -1))
    key = f"txt_{embed_idx}_{int(time.time())}"
    thumb_path = render_text_thumb(text, key, store.thumbs)
    with store.db() as conn:
        conn.execute(
            "INSERT INTO items(type,source,text,thumb_path,timestamp,embed_idx) VALUES (?,?,?,?,?,?)",
            ("text", "note", text, thumb_path, datetime.now().isoformat(), embed_idx),
        )
    with store.lock:
        store.save_index()
    return f"✅ Saved note. Total: {store.index.ntotal}."


# ---------- search ----------
def _search(store: Store, query: str, top_k: int) -> List[Dict]:
    if not query.strip() or store.index.ntotal == 0:
        return []
    qvec = embed_text(query).reshape(1, -1)
    top_k = min(top_k, store.index.ntotal)
    scores, ids = store.index.search(qvec, top_k)
    out: List[Dict] = []
    with store.db() as conn:
        for score, embed_idx in zip(scores[0], ids[0]):
            if embed_idx < 0:
                continue
            row = conn.execute("SELECT * FROM items WHERE embed_idx=?", (int(embed_idx),)).fetchone()
            if row is None:
                continue
            out.append({**dict(row), "score": float(score)})
    return out


def search_ui(query: str, top_k, profile: Optional[gr.OAuthProfile] = None):
    store = get_store(profile)
    if store is None:
        return [], LOGIN_REQUIRED_MD
    hits = _search(store, query, int(top_k))
    if not hits:
        return [], "*No matches. Add memories in the ➕ tab first.*"
    gallery, lines = [], []
    for h in hits:
        label = f"{h['score']:.2f} · {h['timestamp'][:16]}"
        if h["thumb_path"] and Path(h["thumb_path"]).exists():
            gallery.append((h["thumb_path"], label))
        preview = (h["text"] or "").strip().replace("\n", " ")
        if len(preview) > 220:
            preview = preview[:220] + "…"
        lines.append(
            f"- **[{h['score']:.2f}]** `{h['timestamp'][:16]}` · *{h['type']}* · "
            f"{h['source'] or ''} — {preview}"
        )
    return gallery, "\n".join(lines)


# ---------- timeline heatmap ----------
def _empty_plot(title: str) -> go.Figure:
    fig = go.Figure()
    fig.update_layout(
        title=title,
        height=230,
        margin=dict(l=40, r=10, t=50, b=10),
        plot_bgcolor="white",
        xaxis=dict(visible=False),
        yaxis=dict(visible=False),
    )
    return fig


def build_timeline_plot(profile: Optional[gr.OAuthProfile] = None):
    store = get_store(profile)
    if store is None:
        return _empty_plot("Sign in to see your memory timeline")
    with store.db() as conn:
        rows = conn.execute("SELECT timestamp, type FROM items").fetchall()

    end = datetime.now().date()
    start = end - timedelta(days=364)
    counts: Counter = Counter()
    for r in rows:
        d = datetime.fromisoformat(r[0]).date()
        if start <= d <= end:
            counts[d] += 1

    first_mon = start - timedelta(days=start.weekday())
    weeks = []
    cur = first_mon
    while cur <= end:
        week = []
        for i in range(7):
            d = cur + timedelta(days=i)
            week.append((d, counts.get(d, 0) if start <= d <= end else None))
        weeks.append(week)
        cur += timedelta(days=7)

    z = [[weeks[w][d][1] for w in range(len(weeks))] for d in range(7)]
    hover = [
        [f"{weeks[w][d][0].isoformat()} · {weeks[w][d][1] or 0} memories" for w in range(len(weeks))]
        for d in range(7)
    ]

    tickvals, ticktext = [], []
    seen_months = set()
    for w_i, week in enumerate(weeks):
        for d, _ in week:
            key = (d.year, d.month)
            if key not in seen_months and d >= start:
                seen_months.add(key)
                tickvals.append(w_i)
                ticktext.append(d.strftime("%b"))
                break

    fig = go.Figure(
        data=go.Heatmap(
            z=z,
            text=hover,
            hoverinfo="text",
            colorscale=[
                [0.0, "#ebedf0"],
                [0.25, "#9be9a8"],
                [0.5, "#40c463"],
                [0.75, "#30a14e"],
                [1.0, "#216e39"],
            ],
            showscale=False,
            xgap=3,
            ygap=3,
            zmin=0,
        )
    )
    fig.update_layout(
        title=f"Memory density · last 365 days · {sum(counts.values())} in window",
        yaxis=dict(
            tickmode="array",
            tickvals=list(range(7)),
            ticktext=["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"],
            autorange="reversed",
            showgrid=False,
        ),
        xaxis=dict(
            tickmode="array",
            tickvals=tickvals,
            ticktext=ticktext,
            side="top",
            showgrid=False,
        ),
        height=230,
        margin=dict(l=40, r=10, t=50, b=10),
        plot_bgcolor="white",
    )
    return fig


def timeline_summary(profile: Optional[gr.OAuthProfile] = None) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    with store.db() as conn:
        total = conn.execute("SELECT COUNT(*) FROM items").fetchone()[0]
        by_type = conn.execute("SELECT type, COUNT(*) FROM items GROUP BY type").fetchall()
        recent = conn.execute(
            "SELECT timestamp, type, source, text FROM items ORDER BY id DESC LIMIT 8"
        ).fetchall()
    if total == 0:
        return "Nothing yet. Add memories in the ➕ tab."
    lines = [f"### {total} total · " + " · ".join(f"{t}: {c}" for t, c in by_type), "", "**Recent:**"]
    for r in recent:
        preview = (r["text"] or "").strip().replace("\n", " ")[:80]
        lines.append(f"- `{r['timestamp'][:16]}` · *{r['type']}* · {r['source'] or ''} — {preview}")
    return "\n".join(lines)


# ---------- chat ----------
SYSTEM_PROMPT = (
    "You are the user's memory assistant. Answer using ONLY the retrieved memories below. "
    "Cite memories by their [#N] number. If nothing is relevant, say so honestly. Be concise."
)


def _chat_gemini(question: str, hits: List[Dict], api_key: str) -> str:
    client = google_genai.Client(api_key=api_key)
    parts: List = [SYSTEM_PROMPT, "\nMemories:"]
    imgs_added = 0
    for i, h in enumerate(hits, 1):
        meta = f"[#{i}] ({h['timestamp'][:16]}, {h['type']}, {h['source'] or 'note'})"
        if h["type"] == "image" and h["thumb_path"] and imgs_added < MAX_VISION_IMAGES:
            try:
                parts.append(meta + ":")
                parts.append(Image.open(h["thumb_path"]))
                if h["text"]:
                    parts.append(f"(OCR: {h['text'].strip()[:200]})")
                imgs_added += 1
                continue
            except Exception as e:
                print(f"Vision load failed [{h['thumb_path']}]: {e}")
        text = (h["text"] or "").strip().replace("\n", " ")[:400]
        parts.append(f"{meta}: {text}")
    parts.append(f"\nQuestion: {question}")
    resp = client.models.generate_content(model="gemini-2.0-flash", contents=parts)
    return resp.text


def _chat_groq(question: str, hits: List[Dict], api_key: str) -> str:
    ctx = "\n".join(
        f"[#{i}] ({h['timestamp'][:16]}, {h['type']}, {h['source'] or 'note'}): "
        f"{(h['text'] or '').strip().replace(chr(10), ' ')[:400]}"
        for i, h in enumerate(hits, 1)
    )
    client = Groq(api_key=api_key)
    resp = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Memories:\n{ctx}\n\nQuestion: {question}"},
        ],
        temperature=0.3,
        max_tokens=600,
    )
    return resp.choices[0].message.content


def _pick_chat_backend(user_keys: Dict[str, str]):
    """Prefer the user's own key, else the Space-owner env fallback. Returns (provider, key) or (None, None)."""
    uk = user_keys or {}
    gemini_key = (uk.get("gemini") or "").strip() or OWNER_GEMINI_KEY
    groq_key = (uk.get("groq") or "").strip() or OWNER_GROQ_KEY
    if GEMINI_SDK and gemini_key:
        return "gemini", gemini_key
    if GROQ_SDK and groq_key:
        return "groq", groq_key
    return None, None


NO_KEY_MSG = (
    "⚠️ No chat API key set. Paste your own Gemini or Groq API key in the **🔑 Keys** tab — "
    "it's used only for your session and never written to disk. "
    "(Free Gemini key: https://aistudio.google.com/apikey · Free Groq key: https://console.groq.com/keys)"
)


def chat(
    question: str,
    history,
    user_keys: Optional[Dict[str, str]],
    profile: Optional[gr.OAuthProfile] = None,
):
    if not question.strip():
        return "", history
    store = get_store(profile)
    if store is None:
        return "", history + [
            {"role": "user", "content": question},
            {"role": "assistant", "content": LOGIN_REQUIRED_MD},
        ]
    provider, key = _pick_chat_backend(user_keys or {})
    if provider is None:
        answer = NO_KEY_MSG
    else:
        hits = _search(store, question, top_k=6)
        if not hits:
            answer = "I don't have any memories matching that yet."
        else:
            try:
                if provider == "gemini":
                    answer = _chat_gemini(question, hits, key)
                else:
                    answer = _chat_groq(question, hits, key)
            except Exception as e:
                answer = f"Chat error ({provider}): {e}"
    return "", history + [
        {"role": "user", "content": question},
        {"role": "assistant", "content": answer},
    ]


# ---------- manage ----------
def stats(profile: Optional[gr.OAuthProfile] = None) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    with store.db() as conn:
        total = conn.execute("SELECT COUNT(*) FROM items").fetchone()[0]
        by_type = conn.execute("SELECT type, COUNT(*) FROM items GROUP BY type").fetchall()
    who = f" · signed in as **{profile.username}**" if profile else ""
    if total == 0:
        return f"**0 memories** — add some in the ➕ tab.{who}"
    return f"**{total} memories** · " + " · ".join(f"{t}: {c}" for t, c in by_type) + who


def clear_all(profile: Optional[gr.OAuthProfile] = None) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    store.reset()
    return "🗑️ Cleared."


def export_data(profile: Optional[gr.OAuthProfile] = None):
    store = get_store(profile)
    if store is None:
        return None
    out = store.export_dir / f"memory_export_{int(time.time())}.zip"
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        if store.db_path.exists():
            z.write(store.db_path, "memory.db")
        if store.index_path.exists():
            z.write(store.index_path, "faiss.index")
        for f in store.thumbs.glob("*.jpg"):
            z.write(f, f"thumbs/{f.name}")
    return str(out)


def import_data(zip_path, profile: Optional[gr.OAuthProfile] = None) -> str:
    store = get_store(profile)
    if store is None:
        return LOGIN_REQUIRED_MD
    if not zip_path:
        return "No file."
    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(store.root)
    with store.lock:
        store.index = store._load_index()
    return f"✅ Restored. Total: {store.index.ntotal}."


# ---------- UI ----------
def chat_backend_badge(user_keys: Optional[Dict[str, str]] = None) -> str:
    provider, _ = _pick_chat_backend(user_keys or {})
    if provider == "gemini":
        source = "your key" if (user_keys or {}).get("gemini") else "Space-owner fallback"
        return f"💬 Chat: **Gemini 2.0 Flash** (vision) · using {source}"
    if provider == "groq":
        source = "your key" if (user_keys or {}).get("groq") else "Space-owner fallback"
        return f"💬 Chat: **Groq Llama-3.3-70B** (text) · using {source}"
    return NO_KEY_MSG


def save_keys(gemini_key: str, groq_key: str):
    """Store the user's own keys in session state. Never written to disk."""
    keys = {
        "gemini": (gemini_key or "").strip(),
        "groq": (groq_key or "").strip(),
    }
    return keys, chat_backend_badge(keys), "✅ Saved for this session. Refresh clears them."


def clear_keys():
    empty: Dict[str, str] = {"gemini": "", "groq": ""}
    return empty, "", "", chat_backend_badge(empty), "🗑️ Cleared."


with gr.Blocks(title="Personal Memory Search", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🧠 Personal Memory Search\n"
        "Semantic search & chat over your screenshots, photos and notes. "
        "Sign in with Hugging Face — each user gets a private, persistent store."
    )
    with gr.Row():
        # gr.LoginButton is a no-op locally; on Spaces with hf_oauth: true it triggers HF OAuth.
        gr.LoginButton()
    stats_md = gr.Markdown(LOGIN_REQUIRED_MD)

    # Per-session BYOK: keys live in server RAM tied to this browser session, never on disk.
    keys_state = gr.State({"gemini": "", "groq": ""})

    with gr.Tabs():
        with gr.Tab("🔍 Search"):
            with gr.Row():
                q = gr.Textbox(
                    label="Query",
                    placeholder="e.g. 'that pricing table' or 'sunset from the balcony'",
                    scale=4,
                )
                k = gr.Slider(1, 20, value=TOP_K_DEFAULT, step=1, label="Top K", scale=1)
            search_btn = gr.Button("Search", variant="primary")
            gallery = gr.Gallery(label="Results", columns=4, height=380, object_fit="contain")
            details = gr.Markdown()
            search_btn.click(search_ui, [q, k], [gallery, details])
            q.submit(search_ui, [q, k], [gallery, details])

        with gr.Tab("💬 Chat"):
            backend_md = gr.Markdown(chat_backend_badge())
            chatbot = gr.Chatbot(height=440, type="messages")
            msg = gr.Textbox(placeholder="Ask something about your memories…", label="")
            msg.submit(chat, [msg, chatbot, keys_state], [msg, chatbot])
            gr.Button("Clear chat").click(lambda: [], None, chatbot)

        with gr.Tab("📅 Timeline"):
            timeline_plot = gr.Plot(value=_empty_plot("Sign in to see your memory timeline"))
            timeline_md = gr.Markdown(LOGIN_REQUIRED_MD)
            refresh_btn = gr.Button("Refresh")
            refresh_btn.click(build_timeline_plot, None, timeline_plot).then(
                timeline_summary, None, timeline_md
            )

        with gr.Tab("➕ Add memories"):
            with gr.Row():
                with gr.Column():
                    imgs = gr.File(
                        label="Screenshots / photos",
                        file_count="multiple",
                        file_types=["image"],
                        type="filepath",
                    )
                    ingest_btn = gr.Button("Index images", variant="primary")
                    ingest_out = gr.Markdown()
                    ingest_btn.click(ingest_images, imgs, ingest_out).then(
                        stats, None, stats_md
                    ).then(build_timeline_plot, None, timeline_plot).then(
                        timeline_summary, None, timeline_md
                    )
                with gr.Column():
                    note = gr.Textbox(
                        label="Paste a note",
                        lines=6,
                        placeholder="A quote, an idea, a thought…",
                    )
                    note_btn = gr.Button("Save note", variant="primary")
                    note_out = gr.Markdown()
                    note_btn.click(ingest_text, note, note_out).then(
                        stats, None, stats_md
                    ).then(build_timeline_plot, None, timeline_plot).then(
                        timeline_summary, None, timeline_md
                    ).then(lambda: "", None, note)

        with gr.Tab("🔑 Keys"):
            gr.Markdown(
                "### Bring your own API key\n"
                "Chat lets a language model answer over your memories. To avoid charging the "
                "Space owner's quota, paste **your own** Gemini or Groq key below.\n\n"
                "- **Session-only.** Keys live in server memory tied to this browser session and are "
                "never written to disk. Refreshing this page clears them.\n"
                "- **Priority.** If you paste a key it's used; otherwise the Space falls back to any "
                "key the owner configured.\n"
                "- Free keys: [Gemini (AI Studio)](https://aistudio.google.com/apikey) · "
                "[Groq](https://console.groq.com/keys)"
            )
            with gr.Row():
                gemini_in = gr.Textbox(
                    label="Gemini API key (vision + text)",
                    placeholder="AIza…",
                    type="password",
                )
                groq_in = gr.Textbox(
                    label="Groq API key (text only)",
                    placeholder="gsk_…",
                    type="password",
                )
            with gr.Row():
                save_keys_btn = gr.Button("Save keys for this session", variant="primary")
                clear_keys_btn = gr.Button("Clear keys")
            keys_status = gr.Markdown()
            save_keys_btn.click(
                save_keys,
                [gemini_in, groq_in],
                [keys_state, backend_md, keys_status],
            )
            clear_keys_btn.click(
                clear_keys,
                None,
                [keys_state, gemini_in, groq_in, backend_md, keys_status],
            )

        with gr.Tab("⚙️ Manage"):
            gr.Markdown(
                "### Storage\n"
                "Your memories live under `"
                + str(BASE_DIR)
                + "/<your-username>/` on the Space's persistent volume."
            )
            gr.Markdown("### Backup / restore")
            with gr.Row():
                export_btn = gr.Button("Export as zip")
                export_out = gr.File(label="Download")
                export_btn.click(export_data, None, export_out)
            with gr.Row():
                import_zip = gr.File(label="Restore zip", file_types=[".zip"], type="filepath")
                import_btn = gr.Button("Restore")
                import_out = gr.Markdown()
                import_btn.click(import_data, import_zip, import_out).then(
                    stats, None, stats_md
                ).then(build_timeline_plot, None, timeline_plot).then(
                    timeline_summary, None, timeline_md
                )
            gr.Markdown("### Danger zone")
            clear_btn = gr.Button("Delete all my memories", variant="stop")
            clear_out = gr.Markdown()
            clear_btn.click(clear_all, None, clear_out).then(stats, None, stats_md).then(
                build_timeline_plot, None, timeline_plot
            ).then(timeline_summary, None, timeline_md)

    # Populate per-user data on connection (fires again after OAuth redirect completes).
    demo.load(stats, None, stats_md)
    demo.load(build_timeline_plot, None, timeline_plot)
    demo.load(timeline_summary, None, timeline_md)


if __name__ == "__main__":
    demo.launch(ssr_mode=False)