# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# DISCLAIMER: This software is provided "as is" without any warranty,
# express or implied, including but not limited to the warranties of
# merchantability, fitness for a particular purpose, and non-infringement.
#
# In no event shall the authors or copyright holders be liable for any
# claim, damages, or other liability, whether in an action of contract,
# tort, or otherwise, arising from, out of, or in connection with the
# software or the use or other dealings in the software.
# -----------------------------------------------------------------------------
 
# @Author  : Tek Raj Chhetri
# @Email   : tekraj@mit.edu
# @Web     : https://tekrajchhetri.com/
# @File    : lightweightfts5.py
# @Software: PyCharm

#!/usr/bin/env python3
"""
Ontology Semantic Search — Lightweight Edition
===============================================
• NO ChromaDB — eliminates the 7 GB+ memory overhead
• Semantic search via SentenceTransformers + numpy cosine similarity
• Embeddings computed on-the-fly per query (fast with small model)
• SQLite FTS5 for fast text pre-filtering to keep memory low
• Only downloads the SQLite DB (~few hundred MB), not the Chroma index
• Fits comfortably in a 16 GB HuggingFace Space

Dataset : https://huggingface.co/datasets/sensein/ontology-sqlite-vectorstore
Run     : python app_gradio.py
"""

from __future__ import annotations

import gc
import inspect
import os
import re
import shutil
import sqlite3
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import gradio as gr

# ─────────────────────────────────────────────────────────────────────────────
# Paths & defaults
# ─────────────────────────────────────────────────────────────────────────────

CACHE_DIR          = Path(os.getenv("ONTOLOGY_CACHE_DIR", "cache"))
DEFAULT_DB_PATH    = str(CACHE_DIR / os.getenv("ONTOLOGY_DB_NAME", "bioportal.db"))
DEFAULT_MODEL      = os.getenv("ONTOLOGY_EMBED_MODEL", "BAAI/bge-small-en-v1.5")

MODEL_CHOICES = [
    "BAAI/bge-small-en-v1.5",
    "BAAI/bge-large-en-v1.5",
    "sentence-transformers/all-MiniLM-L6-v2",
    "BAAI/bge-base-en-v1.5",
]

HF_DATASET_REPO = "sensein/ontology-sqlite-vectorstore"
HF_DB_FILENAME  = "bioportal.db"

LOGO_URL     = "https://avatars.githubusercontent.com/u/47326880"
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "

CACHE_DIR.mkdir(parents=True, exist_ok=True)
IS_HF_SPACE = os.getenv("SPACE_ID") is not None

# ─────────────────────────────────────────────────────────────────────────────
# CSS (compact)
# ─────────────────────────────────────────────────────────────────────────────

CUSTOM_CSS = """
body, .gradio-container { font-family: 'Inter', system-ui, sans-serif !important; }
#app-header { display:flex;align-items:center;gap:20px;background:linear-gradient(135deg,#1e3a5f 0%,#0f2540 60%,#162d4a 100%);border-radius:12px;padding:20px 28px;margin-bottom:16px;box-shadow:0 4px 20px rgba(0,0,0,.25); }
#app-header img { width:56px;height:56px;border-radius:10px;border:2px solid rgba(255,255,255,.2);flex-shrink:0; }
#app-header-text h1 { margin:0;font-size:1.6rem;font-weight:700;color:#fff; }
#app-header-text p { margin:4px 0 0;font-size:.875rem;color:rgba(255,255,255,.65); }
#status-bar textarea { background:#f0fdf4!important;border:1.5px solid #86efac!important;border-radius:8px!important;color:#166534!important;font-size:.82rem!important;font-family:'JetBrains Mono',monospace!important;padding:6px 12px!important;min-height:36px!important; }
#sidebar { background:#f8fafc;border:1px solid #e2e8f0;border-radius:12px;padding:16px; }
#query-box textarea { font-size:1rem!important;border:2px solid #cbd5e1!important;border-radius:10px!important;padding:12px!important; }
#query-box textarea:focus { border-color:#3b82f6!important; }
#result-table table { font-size:.83rem!important; }
#result-table th { background:#1e3a5f!important;color:#fff!important;font-weight:600!important;position:sticky;top:0; }
#result-table tr:hover td { background:#eff6ff!important; }
#sql-editor textarea { font-family:'JetBrains Mono',monospace!important;font-size:.88rem!important;background:#0f172a!important;color:#e2e8f0!important;border-radius:8px!important;border:1.5px solid #334155!important; }
#dl-log textarea { background:#fffbeb!important;border:1.5px solid #fcd34d!important;border-radius:8px!important;color:#78350f!important;font-size:.82rem!important;font-family:'JetBrains Mono',monospace!important; }
"""

# ─────────────────────────────────────────────────────────────────────────────
# Device
# ─────────────────────────────────────────────────────────────────────────────

def pick_device(user_device: Optional[str] = None) -> str:
    if user_device and user_device != "auto":
        return user_device
    if IS_HF_SPACE:
        return "cpu"
    try:
        import torch
        if torch.cuda.is_available(): return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "mps"
    except ImportError:
        pass
    return "cpu"


def make_embedder(model_name: str, device: str):
    from sentence_transformers import SentenceTransformer
    emb = SentenceTransformer(model_name, device=device)
    gc.collect()
    return emb


# ─────────────────────────────────────────────────────────────────────────────
# Safe encode
# ─────────────────────────────────────────────────────────────────────────────

_encode_sig_cache: Dict[int, set] = {}

def _safe_encode(emb, sentences: List[str], *,
                 normalize_embeddings: bool = True, batch_size: int = 64):
    key = id(emb)
    if key not in _encode_sig_cache:
        _encode_sig_cache[key] = set(inspect.signature(emb.encode).parameters)
    accepted = _encode_sig_cache[key]
    kw: Dict[str, Any] = {}
    if "normalize_embeddings" in accepted: kw["normalize_embeddings"] = normalize_embeddings
    if "show_progress_bar"    in accepted: kw["show_progress_bar"]    = False
    if "batch_size"           in accepted: kw["batch_size"]           = batch_size
    if "num_workers"          in accepted: kw["num_workers"]          = 0
    return emb.encode(sentences, **kw)


# ─────────────────────────────────────────────────────────────────────────────
# Download — only the SQLite DB (skip the 7 GB Chroma index entirely)
# ─────────────────────────────────────────────────────────────────────────────

def _has_huggingface_hub() -> bool:
    try:
        import huggingface_hub
        return True
    except ImportError:
        return False


def download_assets(force=False, progress_cb=None) -> Tuple[bool, str]:
    def _log(msg):
        if progress_cb: progress_cb(msg)
        print(msg)

    db_path = Path(DEFAULT_DB_PATH)

    if db_path.exists() and not force:
        _log(f"SQLite DB already present: {db_path}")
        return True, "Assets ready."

    db_path.parent.mkdir(parents=True, exist_ok=True)

    if _has_huggingface_hub():
        from huggingface_hub import hf_hub_download
        _log(f"Downloading SQLite DB from {HF_DATASET_REPO} via huggingface_hub...")
        try:
            hf_hub_download(repo_id=HF_DATASET_REPO, filename=HF_DB_FILENAME,
                            repo_type="dataset", local_dir=str(CACHE_DIR),
                            force_download=force)
            _log(f"SQLite DB saved -> {db_path}")
            return True, "Assets ready."
        except Exception as e:
            return False, f"Failed to download: {e}"
    else:
        import urllib.request
        url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/{HF_DB_FILENAME}"
        _log(f"Downloading SQLite DB...\n  {url}")
        tmp = db_path.with_suffix(".tmp")
        try:
            def _hook(count, block, total):
                if total > 0:
                    pct = min(count * block / total * 100, 100)
                    _log(f"  {pct:.1f}% ({count*block/1e6:.1f}/{total/1e6:.1f} MB)")
            urllib.request.urlretrieve(url, tmp, reporthook=_hook)
            tmp.rename(db_path)
            _log(f"SQLite DB saved -> {db_path}")
            return True, "Assets ready."
        except Exception as e:
            if tmp.exists(): tmp.unlink()
            return False, f"Failed: {e}"


# ─────────────────────────────────────────────────────────────────────────────
# FTS5 setup — create full-text search index in SQLite if not present
# ─────────────────────────────────────────────────────────────────────────────

def ensure_fts(con: sqlite3.Connection):
    """Create FTS5 virtual table for fast text pre-filtering if not exists."""
    try:
        con.execute("SELECT 1 FROM classes_fts LIMIT 1")
        return  # already exists
    except sqlite3.OperationalError:
        pass

    print("Building FTS5 index (one-time)...")
    con.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS classes_fts USING fts5(
            preferred_label, definition,
            content='classes', content_rowid='id'
        )
    """)
    con.execute("""
        INSERT INTO classes_fts(rowid, preferred_label, definition)
        SELECT id, COALESCE(preferred_label,''), COALESCE(definition,'')
        FROM classes
    """)
    con.commit()
    print("FTS5 index built.")


# ─────────────────────────────────────────────────────────────────────────────
# Runtime singleton
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class Runtime:
    embedder: Any
    db_con:   sqlite3.Connection
    device:   str
    model:    str


_runtime:     Optional[Runtime] = None
_runtime_err: Optional[str]     = None
_loading:     bool              = False
_load_log:    List[str]         = []
_runtime_lock = threading.Lock()
_db_lock      = threading.Lock()


def _log_startup(msg: str):
    _load_log.append(msg)
    print(msg)


def _open_sqlite(path: str) -> sqlite3.Connection:
    con = sqlite3.connect(path, check_same_thread=False)
    cache_kb = 16384 if IS_HF_SPACE else 65536
    for p in ("PRAGMA journal_mode=WAL", "PRAGMA synchronous=NORMAL",
              "PRAGMA temp_store=MEMORY", f"PRAGMA cache_size=-{cache_kb}"):
        con.execute(p)
    return con


def _do_load(db_path: str, model_name: str, device: str):
    global _runtime, _runtime_err, _loading
    try:
        t0 = time.time()
        _log_startup(f"Loading model: {model_name} on {device}")
        embedder = make_embedder(model_name, device)

        _log_startup("Warming up model...")
        _safe_encode(embedder, ["warmup"], normalize_embeddings=True, batch_size=1)
        gc.collect()
        _log_startup("Model warm-up done.")

        _log_startup("Opening SQLite...")
        con = _open_sqlite(db_path)

        _log_startup("Ensuring FTS5 index...")
        ensure_fts(con)

        with _runtime_lock:
            _runtime = Runtime(embedder=embedder, db_con=con,
                               device=device, model=model_name)
        _log_startup(f"Ready in {time.time()-t0:.1f}s | model={model_name}")
    except Exception as e:
        with _runtime_lock:
            _runtime_err = f"{type(e).__name__}: {e}"
        _log_startup(f"Load failed: {e}")
    finally:
        _loading = False
        gc.collect()


def get_runtime() -> Tuple[Optional[Runtime], Optional[str]]:
    if _loading:
        return None, "Loading — please wait..."
    with _runtime_lock:
        if _runtime_err: return None, f"Error: {_runtime_err}"
        if _runtime is None: return None, "No data loaded. Go to Data & Setup tab."
        return _runtime, None


def reload_runtime(db_path: str, model_name: str, device: str) -> str:
    global _runtime, _runtime_err, _loading
    _loading = True
    with _runtime_lock:
        if _runtime:
            try: _runtime.db_con.close()
            except: pass
        _runtime = None; _runtime_err = None; _load_log.clear()
    gc.collect()
    threading.Thread(target=_do_load, args=(db_path, model_name, device), daemon=True).start()
    for _ in range(80):
        time.sleep(0.25)
        if not _loading: break
    rt, err = get_runtime()
    return err if err else f"Ready | model={rt.model} | device={rt.device}"


# ─────────────────────────────────────────────────────────────────────────────
# Startup
# ─────────────────────────────────────────────────────────────────────────────

def _startup():
    global _loading, _runtime_err
    _loading = True

    if not os.path.exists(DEFAULT_DB_PATH):
        _log_startup("SQLite DB not found — downloading from HuggingFace...")
        ok, msg = download_assets(force=False, progress_cb=_log_startup)
        if not ok:
            _runtime_err = msg; _loading = False; return

    if os.path.exists(DEFAULT_DB_PATH):
        _do_load(DEFAULT_DB_PATH, DEFAULT_MODEL, pick_device())
    else:
        _runtime_err = f"SQLite DB not found at '{DEFAULT_DB_PATH}'."
        _loading = False

threading.Thread(target=_startup, daemon=True).start()


# ─────────────────────────────────────────────────────────────────────────────
# SQLite helpers
# ─────────────────────────────────────────────────────────────────────────────

def fetch_synonyms(con: sqlite3.Connection, sqlite_id: int) -> List[str]:
    return [r[0] for r in con.execute(
        "SELECT synonym FROM synonyms WHERE class_id=? ORDER BY synonym",
        (sqlite_id,)).fetchall()]


def concept_details(con: sqlite3.Connection, class_uri: str) -> Optional[Dict[str, Any]]:
    row = con.execute(
        "SELECT id, ontology_id, preferred_label, definition, notation, obsolete "
        "FROM classes WHERE class_uri=?", (class_uri,)).fetchone()
    if not row: return None
    cid = int(row[0])
    syns = con.execute(
        "SELECT synonym, syn_type FROM synonyms WHERE class_id=? ORDER BY synonym",
        (cid,)).fetchall()
    parents = con.execute(
        "SELECT parent_uri FROM parents WHERE class_id=? ORDER BY parent_uri",
        (cid,)).fetchall()
    return {
        "sqlite_id": cid, "ontology": row[1], "label": row[2],
        "definition": row[3], "notation": row[4], "obsolete": bool(row[5]),
        "synonyms": [{"term": s[0], "type": s[1]} for s in syns],
        "parents": [p[0] for p in parents],
    }


# ─────────────────────────────────────────────────────────────────────────────
# Semantic search — FTS5 pre-filter + embedding re-rank
#
# Strategy:
#   1. Use FTS5 to find top N_CANDIDATES matching the query text
#   2. Encode query + candidates with the embedding model
#   3. Re-rank by cosine similarity
#   This avoids loading millions of vectors into RAM.
# ─────────────────────────────────────────────────────────────────────────────

N_CANDIDATES = 500  # FTS candidates to fetch before re-ranking


def _fts_tokenize(query: str) -> str:
    """Convert a natural-language query to an FTS5 OR query."""
    tokens = re.findall(r"[a-zA-Z0-9]+", query.lower())
    if not tokens: return query
    return " OR ".join(tokens)


def _fetch_candidates(con: sqlite3.Connection, query: str,
                      ontology: Optional[str], include_obsolete: bool,
                      limit: int) -> List[Dict[str, Any]]:
    """Fetch candidate rows via FTS5 text matching."""
    fts_query = _fts_tokenize(query)

    # Build the SQL — join FTS results back to classes
    sql = """
        SELECT c.id, c.ontology_id, c.class_uri, c.preferred_label,
               c.definition, c.notation, c.obsolete
        FROM classes_fts fts
        JOIN classes c ON c.id = fts.rowid
        WHERE classes_fts MATCH ?
    """
    params: List[Any] = [fts_query]

    if ontology:
        sql += " AND c.ontology_id = ?"
        params.append(ontology)
    if not include_obsolete:
        sql += " AND (c.obsolete IS NULL OR c.obsolete = 0)"

    sql += f" LIMIT {int(limit)}"

    rows = con.execute(sql, params).fetchall()
    candidates = []
    for r in rows:
        candidates.append({
            "sqlite_id": int(r[0]), "ontology": r[1] or "",
            "uri": r[2] or "", "label": r[3] or "",
            "definition": r[4] or "", "notation": r[5] or "",
            "obsolete": str(int(r[6] or 0)),
        })
    return candidates


def _build_doc_text(cand: Dict[str, Any], synonyms: List[str]) -> str:
    """Build embedding text for a candidate."""
    parts = []
    if cand["label"]: parts.append(f"LABEL: {cand['label']}")
    if synonyms: parts.append("SYNONYMS: " + " | ".join(synonyms[:32]))
    if cand["definition"]: parts.append(f"DEFINITION: {cand['definition']}")
    return "\n".join(parts) or cand["label"] or cand["uri"]


def semantic_search(
    query: str, *, rt: Runtime, n_results: int,
    ontology: Optional[str], include_obsolete: bool,
    min_score: Optional[float], details: bool,
) -> List[Dict[str, Any]]:

    # Step 1: FTS pre-filter
    with _db_lock:
        candidates = _fetch_candidates(
            rt.db_con, query, ontology, include_obsolete,
            limit=max(N_CANDIDATES, n_results * 10))

    if not candidates:
        return []

    # Step 2: build doc texts and encode
    with _db_lock:
        doc_texts = []
        for c in candidates:
            syns = fetch_synonyms(rt.db_con, c["sqlite_id"])
            c["_synonyms"] = syns
            doc_texts.append(_build_doc_text(c, syns))

    query_vec = _safe_encode(
        rt.embedder, [QUERY_PREFIX + query],
        normalize_embeddings=True, batch_size=1)

    doc_vecs = _safe_encode(
        rt.embedder, doc_texts,
        normalize_embeddings=True, batch_size=64)

    # Step 3: cosine similarity (vectors are already normalized)
    scores = np.dot(doc_vecs, query_vec[0]).flatten()

    # Step 4: rank and return top-K
    top_idx = np.argsort(scores)[::-1][:n_results]

    out: List[Dict[str, Any]] = []
    for i in top_idx:
        score = round(float(scores[i]), 4)
        if min_score is not None and score < min_score:
            continue
        c = candidates[i]
        item = {
            "score": score, "ontology": c["ontology"],
            "label": c["label"], "notation": c["notation"],
            "uri": c["uri"], "sqlite_id": c["sqlite_id"],
            "obsolete": c["obsolete"],
            "synonyms": c.get("_synonyms", []),
        }
        if details and c["uri"]:
            with _db_lock:
                item["details"] = concept_details(rt.db_con, c["uri"])
        out.append(item)

    return out


# ─────────────────────────────────────────────────────────────────────────────
# Result formatters
# ─────────────────────────────────────────────────────────────────────────────

def results_to_df(results: List[Dict[str, Any]]) -> pd.DataFrame:
    return pd.DataFrame([{
        "score": r["score"], "label": r["label"], "ontology": r["ontology"],
        "notation": r["notation"], "obsolete": r["obsolete"],
        "synonyms": ", ".join((r.get("synonyms") or [])[:5]),
        "uri": r["uri"], "sqlite_id": r["sqlite_id"],
    } for r in results])


def results_to_md(results: List[Dict[str, Any]], show_details: bool) -> str:
    if not results: return "*No results found.*"
    lines: List[str] = []
    for idx, r in enumerate(results, 1):
        bar = "█" * int(r["score"] * 10) + "░" * (10 - int(r["score"] * 10))
        lines.append(f"### {idx}. {r.get('label','')}  `{r.get('score')}`  `{bar}`")
        lines.append("| Field | Value |\n|---|---|")
        lines.append(f"| **Ontology** | {r.get('ontology','')} |")
        if r.get("notation"): lines.append(f"| **Notation** | {r['notation']} |")
        lines.append(f"| **URI** | `{r.get('uri','')}` |")
        lines.append(f"| **Obsolete** | {r.get('obsolete','')} |")
        syns = r.get("synonyms") or []
        if syns:
            head = ", ".join(syns[:15])
            tail = f" *...+{len(syns)-15} more*" if len(syns) > 15 else ""
            lines.append(f"| **Synonyms** | {head}{tail} |")
        if show_details and r.get("details"):
            d = r["details"]
            if d.get("definition"):
                lines += ["", "> **Definition**", f"> {d['definition']}"]
            if d.get("parents"):
                ps = d["parents"][:20]
                tail = f"\n*...+{len(d['parents'])-20} more*" if len(d["parents"]) > 20 else ""
                lines += ["", "**Parents**", "\n".join(f"- `{p}`" for p in ps) + tail]
        lines.append("")
    return "\n".join(lines)


# ─────────────────────────────────────────────────────────────────────────────
# SQL console
# ─────────────────────────────────────────────────────────────────────────────

_SELECT_RE = re.compile(r"^\s*(with\s+[\s\S]+?\s+)?select\b", re.IGNORECASE)
_FORBID_RE = re.compile(
    r"\b(attach|detach|pragma|vacuum|reindex|drop|alter"
    r"|create|replace|truncate)\b", re.IGNORECASE)


def run_sql(sql: str, con: sqlite3.Connection, *,
            select_only: bool = True, max_rows: int = 500) -> Tuple[pd.DataFrame, str]:
    sql = (sql or "").strip().rstrip(";")
    if not sql: return pd.DataFrame(), "Enter a SQL query."
    if select_only:
        if not _SELECT_RE.match(sql):
            return pd.DataFrame(), "Blocked: only SELECT allowed."
        if _FORBID_RE.search(sql):
            return pd.DataFrame(), "Blocked: forbidden keywords."
        if max_rows > 0 and not re.search(r"\blimit\b", sql, re.IGNORECASE):
            sql = f"{sql} LIMIT {int(max_rows)}"
    with _db_lock:
        cur = con.execute(sql)
        cols = [d[0] for d in (cur.description or [])]
        rows = cur.fetchall()
    if not cols: return pd.DataFrame(), f"OK — rows affected: {cur.rowcount}"
    return pd.DataFrame(rows, columns=cols), f"OK — {len(rows)} rows returned."


# ─────────────────────────────────────────────────────────────────────────────
# UI callbacks
# ─────────────────────────────────────────────────────────────────────────────

def ui_status_poll() -> str:
    if _loading:
        steps = ["[ ] Loading embedding model...", "[ ] Opening SQLite database...",
                 "[ ] Building FTS5 index..."]
        done = 0
        log_text = " ".join(_load_log).lower()
        if "model warm-up done" in log_text: done = 1
        if "opening sqlite"     in log_text: done = 2
        if "ready in"           in log_text: done = 3
        lines = []
        for i, step in enumerate(steps):
            if   i < done:  lines.append(step.replace("[ ]", "[x]"))
            elif i == done: lines.append(step.replace("[ ]", "-->"))
            else:           lines.append(step)
        last = _load_log[-1] if _load_log else "Starting..."
        lines.append(f"\nLast: {last}")
        return "\n".join(lines)
    rt, err = get_runtime()
    if err: return f"Not ready: {err}"
    return f"Ready  |  model={rt.model}  |  device={rt.device}"


def ui_download_and_reload(force, device_choice, model_name):
    log: List[str] = []
    ok, msg = download_assets(force=force, progress_cb=log.append)
    log.append(msg)
    if not ok:
        joined = "\n".join(log)
        return joined, joined
    log.append("Loading runtime...")
    status = reload_runtime(DEFAULT_DB_PATH, model_name.strip(), pick_device(device_choice))
    log.append(status)
    return "\n".join(log), status


def ui_search(query, n_results, ontology, include_obsolete,
              use_min_score, min_score, details):
    t0 = time.time()
    query = (query or "").strip()
    if not query: return pd.DataFrame(), "Enter a query.", ""
    rt, err = get_runtime()
    if err: return pd.DataFrame(), err, ""
    try:
        results = semantic_search(
            query, rt=rt, n_results=int(n_results),
            ontology=ontology.strip() or None,
            include_obsolete=include_obsolete,
            min_score=(min_score if use_min_score else None),
            details=details)
        dt = (time.time() - t0) * 1000
        return (results_to_df(results),
                f"{len(results)} results | {dt:.0f} ms | device={rt.device}",
                results_to_md(results, details))
    except Exception as e:
        return pd.DataFrame(), f"{type(e).__name__}: {e}", ""


def ui_sql(sql, select_only, max_rows):
    rt, err = get_runtime()
    if err: return pd.DataFrame(), err
    try:
        return run_sql(sql, rt.db_con, select_only=select_only, max_rows=int(max_rows))
    except Exception as e:
        return pd.DataFrame(), f"{type(e).__name__}: {e}"


# ─────────────────────────────────────────────────────────────────────────────
# Theme
# ─────────────────────────────────────────────────────────────────────────────

try:
    _theme = gr.themes.Base(
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.slate,
        neutral_hue=gr.themes.colors.slate,
        font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
        font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
    )
except Exception:
    _theme = "default"


# ─────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ─────────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="Ontology Semantic Search") as demo:

    gr.HTML(f"""
    <div id="app-header">
        <img src="{LOGO_URL}" alt="Sensein logo" />
        <div id="app-header-text">
            <h1>Ontology Semantic Search</h1>
            <p>Lightweight semantic search &mdash;
               SQLite FTS5 + SentenceTransformers re-ranking</p>
                 <p>Data source: sensein/ontology-sqlite-vectorstore</p>
        </div>
    </div>
    """)

    runtime_status = gr.Textbox(
        label="Runtime Status", value="Checking local assets...",
        interactive=False, elem_id="status-bar")

    with gr.Tabs():

        # ── Search ────────────────────────────────────────────────────────
        with gr.Tab("Search"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1, min_width=260, elem_id="sidebar"):
                    gr.Markdown("#### Filters")
                    c_n       = gr.Slider(label="Top-K results", minimum=1, maximum=50, value=20, step=1)
                    c_onto    = gr.Textbox(label="Ontology filter", placeholder="e.g. NCIT, HP, GO", value="")
                    c_obs     = gr.Checkbox(label="Include obsolete terms", value=False)
                    c_usemsco = gr.Checkbox(label="Enable min-score filter", value=False)
                    c_mscore  = gr.Slider(label="Min score", minimum=0.0, maximum=1.0, value=0.5, step=0.01)
                    c_detail  = gr.Checkbox(label="Show definitions & parents", value=False)
                with gr.Column(scale=4):
                    c_query = gr.Textbox(
                        label="Search query",
                        placeholder="e.g.  lung cancer  ·  astrocyte  ·  myocardial infarction",
                        lines=2, elem_id="query-box")
                    with gr.Row():
                        btn_search = gr.Button("Search", variant="primary", scale=3)
                        btn_clear  = gr.Button("Clear", scale=1)
                    search_status = gr.Textbox(label="", value="", interactive=False,
                                               show_label=False, elem_id="status-bar")
                    search_table  = gr.Dataframe(label="Results", interactive=False, wrap=True,
                                                  row_count=(0, "dynamic"), elem_id="result-table")
                    search_md     = gr.Markdown(value="")

        # ── SQL Console ───────────────────────────────────────────────────
        with gr.Tab("SQL Console"):
            gr.Markdown("Query the SQLite database directly. **SELECT-only** by default.")
            with gr.Row():
                with gr.Column(scale=3):
                    c_sql = gr.Textbox(
                        label="SQL Query",
                        value="SELECT ontology_id, COUNT(*) AS n\nFROM classes\nGROUP BY ontology_id\nORDER BY n DESC\nLIMIT 20;",
                        lines=8, elem_id="sql-editor")
                    with gr.Row():
                        sql_run       = gr.Button("Execute", variant="primary")
                        c_select_only = gr.Checkbox(label="SELECT-only", value=True)
                        c_max_rows    = gr.Slider(label="Row limit", minimum=10, maximum=5000, value=500, step=10)
                with gr.Column(scale=1, elem_id="sidebar"):
                    gr.Markdown("#### Examples")
                    gr.Markdown("""```sql
SELECT ontology_id, COUNT(*) n FROM classes
GROUP BY ontology_id ORDER BY n DESC;

SELECT * FROM classes
WHERE preferred_label LIKE '%cancer%' LIMIT 50;

SELECT name FROM sqlite_master WHERE type='table';
```""")
            sql_status = gr.Textbox(label="Status", value="", interactive=False)
            sql_table  = gr.Dataframe(label="Results", interactive=False, wrap=True,
                                       row_count=(0, "dynamic"))

        # ── Data & Setup ──────────────────────────────────────────────────
        with gr.Tab("Data & Setup"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown(f"""### Download from HuggingFace
Dataset: [`{HF_DATASET_REPO}`](https://huggingface.co/datasets/{HF_DATASET_REPO})

**Lightweight mode**: Only downloads the SQLite DB (~few hundred MB).
No Chroma index needed — uses FTS5 pre-filtering + embedding re-ranking.

{"**Running on HuggingFace Space** — assets auto-download on startup." if IS_HF_SPACE else "**Running locally** — assets auto-download on startup."}""")
                    c_force    = gr.Checkbox(label="Force re-download", value=False)
                    c_dl_model = gr.Dropdown(label="Embedding model", choices=MODEL_CHOICES, value=DEFAULT_MODEL)
                    c_dl_dev   = gr.Dropdown(label="Device", choices=["auto", "cpu", "cuda", "mps"], value="auto")
                    btn_dl     = gr.Button("Download & Load", variant="primary")
                    dl_log     = gr.Textbox(label="Download log", value="", lines=10, interactive=False, elem_id="dl-log")
                with gr.Column(scale=1, elem_id="sidebar"):
                    gr.Markdown(f"""#### How it works
1. **FTS5 pre-filter**: SQLite full-text search finds ~500 text-matching candidates
2. **Embed & re-rank**: Encodes query + candidates with `{DEFAULT_MODEL}`
3. **Cosine similarity**: Returns top-K by semantic similarity

This uses ~1–2 GB RAM total vs 10+ GB for the full Chroma index.""")

        # ── Advanced Config ───────────────────────────────────────────────
        with gr.Tab("Advanced Config"):
            gr.Markdown("Change model or device at runtime.")
            with gr.Row():
                with gr.Column():
                    adv_db     = gr.Textbox(label="SQLite DB path", value=DEFAULT_DB_PATH)
                    adv_model  = gr.Dropdown(label="Embedding model", choices=MODEL_CHOICES, value=DEFAULT_MODEL)
                    adv_device = gr.Dropdown(label="Device", choices=["auto", "cpu", "cuda", "mps"], value="auto")
                    adv_btn    = gr.Button("Apply & Reload", variant="primary")
                    adv_status = gr.Textbox(label="Status", value="", interactive=False)
                with gr.Column(elem_id="sidebar"):
                    gr.Markdown("""#### Model options
| Model | RAM | Speed | Quality |
|---|---|---|---|
| `bge-small-en-v1.5` | ~130 MB | Fast | Good |
| `all-MiniLM-L6-v2`  | ~90 MB  | Fastest | Fair |
| `bge-base-en-v1.5`  | ~440 MB | Med | Better |""")

    # ── Event wiring ──────────────────────────────────────────────────────
    btn_search.click(fn=ui_search,
        inputs=[c_query, c_n, c_onto, c_obs, c_usemsco, c_mscore, c_detail],
        outputs=[search_table, search_status, search_md])
    c_query.submit(fn=ui_search,
        inputs=[c_query, c_n, c_onto, c_obs, c_usemsco, c_mscore, c_detail],
        outputs=[search_table, search_status, search_md])

    def _clear(): return "", pd.DataFrame(), "", ""
    btn_clear.click(fn=_clear, inputs=[], outputs=[c_query, search_table, search_status, search_md])

    sql_run.click(fn=ui_sql, inputs=[c_sql, c_select_only, c_max_rows], outputs=[sql_table, sql_status])

    btn_dl.click(fn=ui_download_and_reload, inputs=[c_force, c_dl_dev, c_dl_model], outputs=[dl_log, runtime_status])

    def _adv_reload(db_path, model_name, device_choice):
        db_path = db_path.strip(); model_name = model_name.strip()
        if not os.path.exists(db_path): return f"SQLite DB not found: {db_path}"
        return reload_runtime(db_path, model_name, pick_device(device_choice))

    adv_btn.click(fn=_adv_reload, inputs=[adv_db, adv_model, adv_device],
                  outputs=[adv_status]
    ).then(fn=ui_status_poll, inputs=[], outputs=[runtime_status])

    demo.load(fn=ui_status_poll, inputs=[], outputs=[runtime_status])

    try:
        _timer = gr.Timer(value=2)
        _timer.tick(fn=ui_status_poll, inputs=[], outputs=[runtime_status])
    except AttributeError:
        pass


# ─────────────────────────────────────────────────────────────────────────────
# Launch
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    launch_kwargs: Dict[str, Any] = {
        "server_name": "0.0.0.0",
        "server_port": int(os.getenv("PORT", 7860)),
    }

    import inspect as _inspect
    _launch_sig = set(_inspect.signature(demo.launch).parameters)
    if "theme" in _launch_sig:
        launch_kwargs["theme"] = _theme
        launch_kwargs["css"]   = CUSTOM_CSS
    if "show_api" in _launch_sig:
        launch_kwargs["show_api"] = False

    demo.queue().launch(mcp_server=True, **launch_kwargs)