app.py

#!/usr/bin/env python3
"""
Zetema — semantic search over the Perseus Greek corpus. Gradio web UI.
Run: python app.py
Then open http://localhost:7861 in your browser.
"""

import json
import os
from threading import Thread

# Hugging Face Spaces (ZeroGPU): the `spaces` package must be imported before
# torch so it can patch CUDA initialization for the main process. Locally the
# env var is absent and nothing changes.
IS_SPACE = os.environ.get("SPACE_ID") is not None
if IS_SPACE:
    import spaces

import duckdb
import numpy as np
import torch
import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


def gpu(**kwargs):
    """@spaces.GPU on a ZeroGPU Space, no-op decorator locally."""
    if IS_SPACE:
        return spaces.GPU(**kwargs)
    return lambda fn: fn


DB_PATH = os.path.join(os.path.dirname(__file__), "data", "perseus_rag_qwen.duckdb")
# Dataset repo holding the prebuilt DuckDB index, downloaded at startup when
# data/perseus_rag_qwen.duckdb is not on disk (i.e. on a Space).
INDEX_REPO = os.environ.get("ZETEMA_INDEX_REPO", "Jacobo/zetema-index")
DATES_PATH = os.path.join(os.path.dirname(__file__), "data", "tlg_dates.json")
ENG_PATH = os.path.join(os.path.dirname(__file__), "data", "eng_translations.json")
UNKNOWN_CENTURY = "Unknown"
UNKNOWN_GENRE = "Unknown"
RERANKER_NAME = "Qwen/Qwen3-Reranker-0.6B"
RERANK_CANDIDATES = 50   # dense-retrieval pool handed to the reranker
RERANK_BATCH = 16
# Reranker P("yes") below which a search counts as having no real match in
# the corpus (anachronistic concepts like "racism" land here): the UI warns
# and synthesis is skipped instead of confabulating over near-noise passages.
WEAK_MATCH_THRESHOLD = 0.30
LLM_MODELS = ["google/gemma-3-12b-it", "Qwen/Qwen3-4B-Instruct-2507", "Qwen/Qwen3-8B"]
DEFAULT_LLM = LLM_MODELS[0]
if IS_SPACE or torch.cuda.is_available():
    DEVICE = "cuda"
elif torch.backends.mps.is_available():
    DEVICE = "mps"
else:
    DEVICE = "cpu"

# ---------------------------------------------------------------------------
# Index loading (once at startup)
# ---------------------------------------------------------------------------

_model: SentenceTransformer | None = None
_embeddings: np.ndarray | None = None   # (N, dim) float32
_meta: list[dict] | None = None         # list of {author, work, section_ref, text}
_authors: list[str] = []
_cent_min: np.ndarray | None = None     # (N,) int16 signed centuries; 0 = unknown
_cent_max: np.ndarray | None = None
_century_labels: list[str] = []         # dropdown choices, chronological + Unknown
_century_by_label: dict[str, int] = {}
_genre_masks: dict[str, np.ndarray] = {}  # genre label -> (N,) bool section mask
_genre_labels: list[str] = []             # dropdown choices, alphabetical + Unknown

# TLG Canon epithets -> display genre labels. Epithets absent from this map
# fall through as their raw canon abbreviation, so an expanded corpus still
# gets a (less pretty) dropdown entry rather than being dropped.
GENRE_NAMES = {
    "Alchem.": "Alchemy",
    "Apol.": "Apologetics",
    "Astrol.": "Astrology",
    "Astron.": "Astronomy",
    "Biogr.": "Biography",
    "Bucol.": "Bucolic poetry",
    "Chronogr.": "Chronography",
    "Comic.": "Comedy",
    "Eleg.": "Elegy",
    "Epic.": "Epic poetry",
    "Epigr.": "Epigram",
    "Epist.": "Epistolography",
    "Geogr.": "Geography",
    "Geom.": "Geometry",
    "Gnom.": "Gnomology",
    "Gramm.": "Grammar",
    "Hist.": "History",
    "Iamb.": "Iambic poetry",
    "Lyr.": "Lyric poetry",
    "Math.": "Mathematics",
    "Mech.": "Mechanics",
    "Med.": "Medicine",
    "Mus.": "Music",
    "Myth.": "Mythography",
    "Orat.": "Oratory",
    "Paradox.": "Paradoxography",
    "Parodius": "Parody",
    "Perieg.": "Periegesis",
    "Phil.": "Philosophy",
    "Philol.": "Philology",
    "Poet. Med.": "Medical poetry",
    "Poet. Phil.": "Philosophical poetry",
    "Poeta": "Poetry",
    "Rhet.": "Rhetoric",
    "Scr. Eccl.": "Ecclesiastical",
    "Scr. Erot.": "Novel",
    "Soph.": "Sophistic",
    "Tact.": "Tactics",
    "Theol.": "Theology",
    "Trag.": "Tragedy",
}


def century_label(c: int) -> str:
    n = abs(c)
    suffix = {1: "st", 2: "nd", 3: "rd"}.get(n % 10 if n % 100 not in (11, 12, 13) else 0, "th")
    return f"{n}{suffix} c. {'BC' if c < 0 else 'AD'}"


def load_index() -> None:
    global _model, _embeddings, _meta, _authors
    global _cent_min, _cent_max, _century_labels, _century_by_label
    global _genre_masks, _genre_labels

    db_path = DB_PATH
    if not os.path.exists(db_path):
        if IS_SPACE:
            from huggingface_hub import hf_hub_download
            print(f"Downloading index from dataset {INDEX_REPO}…")
            db_path = hf_hub_download(
                repo_id=INDEX_REPO,
                filename="perseus_rag_qwen.duckdb",
                repo_type="dataset",
            )
        else:
            raise FileNotFoundError(
                f"Index not found: {db_path}\n"
                "Run `python ingest.py` first to build the index."
            )

    con = duckdb.connect(db_path, read_only=True)
    model_name = con.execute(
        "SELECT value FROM meta WHERE key = 'model'"
    ).fetchone()[0]

    print(f"Loading embedding model: {model_name}")
    _model = SentenceTransformer(
        model_name,
        model_kwargs={"dtype": torch.float16} if DEVICE != "cpu" else {},
        processor_kwargs={"padding_side": "left"},
    )

    print("Loading embeddings from DuckDB...")
    rows = con.execute(
        "SELECT author, work, section_ref, lang_tag, text, "
        "tlg_id, work_id, cts_ref, embedding "
        "FROM sections ORDER BY id"
    ).fetchall()
    con.close()

    _meta = [
        {
            "author": r[0],
            "work": r[1],
            "section_ref": r[2],
            "lang_tag": r[3],
            "text": r[4],
            "tlg_id": r[5],
            "work_id": r[6],
            "cts_ref": r[7],
        }
        for r in rows
    ]
    _embeddings = np.array([r[8] for r in rows], dtype=np.float32)
    # Embeddings are already L2-normalized from ingest; ensure unit norm
    norms = np.linalg.norm(_embeddings, axis=1, keepdims=True)
    _embeddings = _embeddings / np.where(norms == 0, 1, norms)

    _authors = sorted({m["author"] for m in _meta if m["author"]})

    # Century metadata: join sections to data/tlg_dates.json on tlg_id
    # (derived from the TLG Canon; 5th c. BC = -5, there is no century 0,
    # so 0 serves as the "unknown" sentinel).
    dates: dict[str, dict] = {}
    if os.path.exists(DATES_PATH):
        with open(DATES_PATH, encoding="utf-8") as f:
            dates = json.load(f)
    else:
        print(f"Warning: {DATES_PATH} not found — century filter disabled.")
    cmin, cmax = [], []
    genre_sets: list[set[str]] = []  # per-section genres (from author epithets)
    for m in _meta:
        info = dates.get(m["tlg_id"]) or {}
        lo, hi = info.get("century_min"), info.get("century_max")
        m["date"] = info.get("date") if lo is not None else None
        cmin.append(lo if lo is not None else 0)
        cmax.append(hi if hi is not None else 0)
        genre_sets.append(
            {GENRE_NAMES.get(e, e) for e in info.get("epithets", [])}
        )
    _cent_min = np.array(cmin, dtype=np.int16)
    _cent_max = np.array(cmax, dtype=np.int16)
    present = sorted({
        c
        for lo, hi in zip(cmin, cmax) if lo != 0
        for c in range(lo, hi + 1) if c != 0
    })
    _century_by_label = {century_label(c): c for c in present}
    _century_labels = list(_century_by_label) + [UNKNOWN_CENTURY]

    # Genre masks: one boolean vector per genre present in the corpus, so
    # filtering at query time is a few array ORs. Authors can carry several
    # epithets; sections with none fall into the "Unknown" bucket.
    n = len(genre_sets)
    _genre_masks = {
        g: np.fromiter((g in s for s in genre_sets), dtype=bool, count=n)
        for g in sorted({g for s in genre_sets for g in s})
    }
    _genre_labels = list(_genre_masks) + [UNKNOWN_GENRE]
    _genre_masks[UNKNOWN_GENRE] = np.fromiter(
        (not s for s in genre_sets), dtype=bool, count=n
    )

    print(f"Ready — {len(_meta):,} sections, {len(_authors)} authors.")


# ---------------------------------------------------------------------------
# Reranker (Qwen3-Reranker: causal LM scored on yes/no logits)
# ---------------------------------------------------------------------------

_rr_tokenizer = None
_rr_model = None
_rr_yes_id = None
_rr_no_id = None

RERANK_PREFIX = (
    "<|im_start|>system\nJudge whether the Document meets the requirements "
    "based on the Query and the Instruct provided. Note that the answer can "
    'only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
)
RERANK_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
RERANK_TASK = (
    "Given a search query about ancient Greek literature (in English or "
    "Greek), judge whether the ancient Greek passage is relevant to it"
)


def load_reranker() -> None:
    global _rr_tokenizer, _rr_model, _rr_yes_id, _rr_no_id
    print(f"Loading reranker: {RERANKER_NAME}")
    _rr_tokenizer = AutoTokenizer.from_pretrained(RERANKER_NAME, padding_side="left")
    _rr_model = AutoModelForCausalLM.from_pretrained(
        RERANKER_NAME,
        dtype=torch.float16 if DEVICE != "cpu" else torch.float32,
    ).to(DEVICE).eval()
    _rr_yes_id = _rr_tokenizer.convert_tokens_to_ids("yes")
    _rr_no_id = _rr_tokenizer.convert_tokens_to_ids("no")


@torch.no_grad()
def rerank_scores(query: str, texts: list[str]) -> list[float]:
    """P(relevant) for each passage under the Qwen3 reranker."""
    prompts = [
        f"{RERANK_PREFIX}<Instruct>: {RERANK_TASK}\n<Query>: {query}\n"
        f"<Document>: {t}{RERANK_SUFFIX}"
        for t in texts
    ]
    scores: list[float] = []
    for start in range(0, len(prompts), RERANK_BATCH):
        batch = prompts[start : start + RERANK_BATCH]
        inputs = _rr_tokenizer(
            batch, padding=True, truncation=True, max_length=2048,
            return_tensors="pt",
        ).to(DEVICE)
        # logits_to_keep=1: only the last position's logits are needed;
        # without it the full (batch, seq, vocab) tensor is ~9 GiB at fp16.
        logits = _rr_model(**inputs, logits_to_keep=1).logits[:, -1, :]
        pair = torch.stack([logits[:, _rr_no_id], logits[:, _rr_yes_id]], dim=1)
        scores.extend(torch.softmax(pair.float(), dim=1)[:, 1].tolist())
    return scores


# ---------------------------------------------------------------------------
# Search
# ---------------------------------------------------------------------------

@gpu()
def semantic_search(
    query: str,
    top_k: int = 10,
    author_filter: list[str] | None = None,
    century_filter: list[str] | None = None,
    genre_filter: list[str] | None = None,
) -> list[dict]:
    """Dense-retrieve a candidate pool, rerank it, return top-k matches."""
    # Qwen3-Embedding embeds queries with an instruction prompt ("query"
    # prompt from the model config); passages were embedded without one.
    q_emb = _model.encode(
        [query],
        prompt_name="query",
        normalize_embeddings=True,
        convert_to_numpy=True,
    )[0]  # (dim,)

    scores = (_embeddings @ q_emb).astype(np.float32)  # (N,) cosine similarity

    if author_filter:
        mask = np.array(
            [m["author"] in author_filter for m in _meta], dtype=bool
        )
        scores = np.where(mask, scores, -2.0)

    if century_filter:
        # A section matches if any selected century falls inside its
        # author's range; ANDs with the author filter via repeated masking.
        match = np.zeros(len(scores), dtype=bool)
        for label in century_filter:
            if label == UNKNOWN_CENTURY:
                match |= _cent_min == 0
            elif label in _century_by_label:
                c = _century_by_label[label]
                match |= (_cent_min <= c) & (c <= _cent_max)
        scores = np.where(match, scores, -2.0)

    if genre_filter:
        # A section matches if its author carries any selected genre;
        # ANDs with the other filters via repeated masking.
        match = np.zeros(len(scores), dtype=bool)
        for label in genre_filter:
            if label in _genre_masks:
                match |= _genre_masks[label]
        scores = np.where(match, scores, -2.0)

    pool = min(max(RERANK_CANDIDATES, top_k), len(scores))
    top_indices = np.argpartition(scores, -pool)[-pool:]
    top_indices = top_indices[np.argsort(scores[top_indices])[::-1]]
    # Author filter can leave fewer valid rows than the pool size
    top_indices = [i for i in top_indices if scores[i] > -2.0]

    candidates = [
        {**_meta[i], "dense_score": float(scores[i]), "idx": int(i)}
        for i in top_indices
    ]
    rr = rerank_scores(query, [c["text"] for c in candidates])
    for c, s in zip(candidates, rr):
        c["score"] = s
    candidates.sort(key=lambda c: c["score"], reverse=True)
    return candidates[:top_k]


def author_summary(results: list[dict], top_n: int = 5) -> str:
    """Return a ranked list of authors by average score in results."""
    by_author: dict[str, list[float]] = {}
    for r in results:
        by_author.setdefault(r["author"], []).append(r["score"])
    ranked = sorted(
        by_author.items(),
        key=lambda kv: (len(kv[1]), sum(kv[1]) / len(kv[1])),
        reverse=True,
    )[:top_n]
    lines = []
    for author, scores in ranked:
        avg = sum(scores) / len(scores)
        lines.append(f"**{author}** — {len(scores)} passage{'s' if len(scores)>1 else ''}, avg score {avg:.3f}")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# LLM synthesis (in-process transformers, streamed)
# ---------------------------------------------------------------------------

_llm_cache: dict[str, tuple] = {}  # model name -> (tokenizer, model)


def get_llm(model_name: str) -> tuple:
    """Lazy-load and cache a chat LLM (first request pays the load time)."""
    if model_name not in _llm_cache:
        print(f"Loading LLM: {model_name}")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            # bf16 on CUDA; fp16 on MPS (Apple-silicon bf16 support is spotty)
            dtype=torch.bfloat16 if DEVICE == "cuda"
            else torch.float16 if DEVICE == "mps"
            else torch.float32,
        ).to(DEVICE).eval()
        _llm_cache[model_name] = (tokenizer, model)
    return _llm_cache[model_name]


@gpu(duration=120)
def llm_chat_stream(messages: list[dict], model_name: str):
    """Stream a chat completion, yielding the accumulated text."""
    try:
        tokenizer, model = get_llm(model_name)
        # enable_thinking is honored by hybrid Qwen3 templates and ignored
        # by instruct-only ones; without it Qwen3-8B emits <think> blocks.
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
            enable_thinking=False,
        )
        inputs = tokenizer([text], return_tensors="pt").to(model.device)
        streamer = TextIteratorStreamer(
            tokenizer, skip_prompt=True, skip_special_tokens=True
        )
        thread = Thread(
            target=model.generate,
            kwargs=dict(
                **inputs, streamer=streamer, max_new_tokens=1024,
                do_sample=True, temperature=0.7, top_p=0.8, top_k=20,
            ),
        )
        thread.start()
        acc = ""
        for chunk in streamer:
            acc += chunk
            yield acc
        thread.join()
    except Exception as e:
        yield f"⚠️ LLM error: {e}"


CONTEXT_PASSAGES = 8  # how many retrieved passages the LLM sees


def passages_context(passages: list[dict]) -> str:
    """Format retrieved passages as numbered context for LLM prompts.

    Carries the reranker's relevance score with each passage. Without it the
    synthesis and chat treat a 0.99 match and a 0.31 near-miss as equally
    authoritative, discarding the ranking the retrieval path just computed.
    """
    return "\n\n".join(
        f"[{i}] {p['author']}, {p['work']} §{p['section_ref']} "
        f"(relevance {p['score']:.2f})\n{p['text'][:600]}"
        for i, p in enumerate(passages[:CONTEXT_PASSAGES], 1)
    )


def evidence_note(passages: list[dict]) -> str:
    """Shared preamble telling the LLM how to read the relevance scores."""
    best = max((p["score"] for p in passages[:CONTEXT_PASSAGES]), default=0.0)
    note = (
        "Passages are listed best-first, and each carries the reranker's "
        "relevance score from 0 to 1. Weight them accordingly: a passage "
        f"scoring below about {WEAK_MATCH_THRESHOLD:.2f} is a weak match that "
        "may not bear on the query at all."
    )
    if best < WEAK_MATCH_THRESHOLD:
        note += (
            " Note that NO passage here scored above "
            f"{WEAK_MATCH_THRESHOLD:.2f}, so the corpus probably does not "
            "address this query. Say so plainly rather than assembling an "
            "answer out of near-misses."
        )
    return note


def synthesize_stream(query: str, passages: list[dict], model_name: str):
    """Stream an LLM synthesis of the retrieved passages."""
    context = passages_context(passages)

    prompt = (
        f"You are a scholar of ancient Greek literature. "
        f"The following passages from the Perseus corpus were retrieved for the query:\n"
        f"  \"{query}\"\n\n"
        f"{evidence_note(passages)}\n\n"
        f"Passages (in ancient Greek):\n{context}\n\n"
        f"Please:\n"
        f"1. Identify the main themes or concepts related to the query found in these passages.\n"
        f"2. Note which authors or works engage most directly with the topic.\n"
        f"3. If helpful, briefly translate or paraphrase key phrases.\n"
        f"Answer concisely in English."
    )

    yield from llm_chat_stream([{"role": "user", "content": prompt}], model_name)


def chat_system_prompt(query: str, passages: list[dict],
                       synthesis: str = "") -> str:
    prompt = (
        "You are a scholar of ancient Greek literature helping a user explore "
        "passages retrieved from the Perseus corpus for the search query "
        f"\"{query}\".\n\n"
        f"{evidence_note(passages)}\n\n"
        f"Retrieved passages (in ancient Greek):\n{passages_context(passages)}\n\n"
    )
    if synthesis.strip():
        # The user is reading this summary on screen, so the chat has to build
        # on it. Without it the chat answers as though the summary never
        # existed and can flatly contradict what the user just read.
        prompt += (
            "You have already given the user this summary of these passages:\n"
            f"---\n{synthesis.strip()}\n---\n"
            "Stay consistent with it. If you now think part of it was wrong, "
            "say so explicitly rather than quietly changing your account.\n\n"
        )
    return prompt + (
        "Answer the user's questions about these passages: interpret, translate, "
        "compare, and point to the relevant authors and works. Refer to passages "
        "by their bracketed numbers. If a question goes beyond what the passages "
        "support, say so rather than inventing sources. Answer in English unless "
        "asked otherwise."
    )


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

# Works that have an English translation in the corpus, mapped to its version
# label (e.g. "tlg0010.tlg018" -> "perseus-eng2"); regenerate with
# extract_eng_translations.py.
_eng_translations: dict[str, str] = {}
if os.path.exists(ENG_PATH):
    with open(ENG_PATH, encoding="utf-8") as _f:
        _eng_translations = json.load(_f)


def scaife_url(r: dict) -> str:
    """Scaife Viewer link for a result — reader page when the passage has a
    CTS reference, the work's library page otherwise. When the work has a
    known English translation, it is opened alongside via ?right=."""
    work = f"{r['tlg_id']}.{r['work_id']}"
    urn = f"urn:cts:greekLit:{work}"
    if r.get("cts_ref"):
        url = f"https://scaife.perseus.org/reader/{urn}.{r['lang_tag']}:{r['cts_ref']}/"
        eng = _eng_translations.get(work)
        return f"{url}?right={eng}" if eng else url
    return f"https://scaife.perseus.org/library/{urn}/"


# Palette as CSS variables set per scheme by a class on <body>; a shared
# mapping block below feeds the same variables into Gradio's own theme
# variables (:root-level, so body-level overrides win) — buttons, filters,
# inputs and page background follow the cards. Default (no class) =
# "Parchment"; the alternatives are switchable at runtime via the scheme
# picker at the bottom of the page. body.dark (Gradio dark mode) reuses the
# Dark scholarly values so ?__theme=dark stays coherent.
APP_CSS = """
body {
    --sem-card-bg: #faf6ec;
    --sem-card-border: #d9cdb4;
    --sem-text: #3d2f1e;
    --sem-author: #8b2e2e;
    --sem-link: #8b2e2e;
    --sem-meta: #7a6a52;
    --sem-score: #a08c5a;
    --sem-page-bg: #efe7d3;
    --sem-block-bg: #faf6ec;
    --sem-input-bg: #fdfbf4;
    --sem-primary: #8b2e2e;
    --sem-primary-hover: #722525;
    --sem-primary-text: #faf6ec;
    --sem-secondary-bg: #e7dcc2;
    --sem-secondary-bg-hover: #ddcfae;
    --sem-accent-soft: #ecdfc4;
    --sem-subdued: #7a6a52;
}
body.sem-aegean {
    --sem-card-bg: #f7fafc;
    --sem-card-border: #cfdde8;
    --sem-text: #17202a;
    --sem-author: #1b4f72;
    --sem-link: #1b4f72;
    --sem-meta: #5d7285;
    --sem-score: #c0603d;
    --sem-page-bg: #e8eff5;
    --sem-block-bg: #ffffff;
    --sem-input-bg: #ffffff;
    --sem-primary: #1b4f72;
    --sem-primary-hover: #163f5c;
    --sem-primary-text: #ffffff;
    --sem-secondary-bg: #dbe7f0;
    --sem-secondary-bg-hover: #cdddea;
    --sem-accent-soft: #d9e7f1;
    --sem-subdued: #5d7285;
}
body.sem-academic {
    --sem-card-bg: #ffffff;
    --sem-card-border: #e0e0e0;
    --sem-text: #1c1c1c;
    --sem-author: #1c1c1c;
    --sem-link: #1a56a0;
    --sem-meta: #666666;
    --sem-score: #9a9a9a;
    --sem-page-bg: #f4f5f7;
    --sem-block-bg: #ffffff;
    --sem-input-bg: #ffffff;
    --sem-primary: #1a56a0;
    --sem-primary-hover: #144781;
    --sem-primary-text: #ffffff;
    --sem-secondary-bg: #e8eaee;
    --sem-secondary-bg-hover: #dcdfe5;
    --sem-accent-soft: #dce8f5;
    --sem-subdued: #666666;
}
body.sem-darkscholar, body.dark {
    --sem-card-bg: #20242b;
    --sem-card-border: #3a4048;
    --sem-text: #e8e2d5;
    --sem-author: #d4a656;
    --sem-link: #d4a656;
    --sem-meta: #8f9aab;
    --sem-score: #d4a656;
    --sem-page-bg: #14171c;
    --sem-block-bg: #20242b;
    --sem-input-bg: #262b34;
    --sem-primary: #d4a656;
    --sem-primary-hover: #c2933f;
    --sem-primary-text: #14171c;
    --sem-secondary-bg: #2c323c;
    --sem-secondary-bg-hover: #363d49;
    --sem-accent-soft: #3a3420;
    --sem-subdued: #8f9aab;
}
body {
    --body-background-fill: var(--sem-page-bg);
    --background-fill-primary: var(--sem-block-bg);
    --background-fill-secondary: var(--sem-accent-soft);
    --block-background-fill: var(--sem-block-bg);
    --body-text-color: var(--sem-text);
    --body-text-color-subdued: var(--sem-subdued);
    --block-title-text-color: var(--sem-text);
    --block-label-text-color: var(--sem-subdued);
    --border-color-primary: var(--sem-card-border);
    --input-background-fill: var(--sem-input-bg);
    --button-primary-background-fill: var(--sem-primary);
    --button-primary-background-fill-hover: var(--sem-primary-hover);
    --button-primary-text-color: var(--sem-primary-text);
    --button-secondary-background-fill: var(--sem-secondary-bg);
    --button-secondary-background-fill-hover: var(--sem-secondary-bg-hover);
    --button-secondary-text-color: var(--sem-text);
    --color-accent: var(--sem-primary);
    --color-accent-soft: var(--sem-accent-soft);
    --link-text-color: var(--sem-link);
    --slider-color: var(--sem-primary);
    --checkbox-background-color-selected: var(--sem-primary);
    --loader-color: var(--sem-primary);
}
.gradio-container { padding-top: 24px !important; }
.sem-title h1 {
    font-size: 2.4em;
    line-height: 1.25;
    margin: 0.2em 0 0.3em 0;
    color: var(--sem-author);
}
.sem-card {
    border: 1px solid var(--sem-card-border);
    border-radius: 8px;
    padding: 14px 16px;
    margin-bottom: 12px;
    background: var(--sem-card-bg);
    font-family: serif;
}
.sem-score {
    font-family: monospace;
    font-size: 0.8em;
    color: var(--sem-score);
    margin-bottom: 6px;
}
.sem-author {
    font-weight: bold;
    font-size: 1.05em;
    color: var(--sem-author);
}
.sem-meta {
    color: var(--sem-meta);
    font-size: 0.9em;
    margin-bottom: 8px;
}
.sem-meta a { color: var(--sem-link); }
.sem-lang { font-size: 0.85em; }
.sem-greek {
    font-size: 1.2em;
    line-height: 1.8;
    color: var(--sem-text);
}
.sem-warn {
    border: 1px solid #d4a017;
    background: #fff8e1;
    border-radius: 8px;
    padding: 10px 14px;
    margin-bottom: 12px;
    color: #7a5c00;
    font-size: 0.9em;
}
body.dark .sem-warn, body.sem-darkscholar .sem-warn {
    background: #33290e;
    border-color: #8a6d12;
    color: #e8c96a;
}
"""

# Scheme picker (bottom of the page): label -> body class ("" = default
# Parchment).
SCHEME_CLASSES = {
    "Parchment": "",
    "Aegean": "sem-aegean",
    "Clean academic": "sem-academic",
    "Dark scholarly": "sem-darkscholar",
}
SCHEME_SWITCH_JS = """
(choice) => {
    const map = {
        "Parchment": "",
        "Aegean": "sem-aegean",
        "Clean academic": "sem-academic",
        "Dark scholarly": "sem-darkscholar",
    };
    document.body.classList.remove("sem-aegean", "sem-academic", "sem-darkscholar");
    if (map[choice]) document.body.classList.add(map[choice]);
}
"""

def format_results_html(results: list[dict]) -> str:
    if not results:
        return "<p>No results.</p>"
    cards = []
    for r in results:
        score_pct = int(r["score"] * 100)
        score_bar = "█" * (score_pct // 5) + "░" * (20 - score_pct // 5)
        text_preview = r["text"][:400] + ("…" if len(r["text"]) > 400 else "")
        date_note = f" &nbsp;·&nbsp; {r['date']}" if r.get("date") else ""
        cards.append(f"""
<div class="sem-card">
  <div class="sem-score">{score_bar} {r['score']:.3f}</div>
  <div class="sem-author">{r['author']}</div>
  <div class="sem-meta">
    <em>{r['work']}</em> &nbsp;·&nbsp; §{r['section_ref']}{date_note}
    &nbsp;·&nbsp; <span class="sem-lang">{r['lang_tag']}</span>
    &nbsp;·&nbsp; <a href="{scaife_url(r)}" target="_blank" rel="noopener">Read in context ↗</a>
  </div>
  <div class="sem-greek">{text_preview}</div>
</div>""")
    return "\n".join(cards)


def run_search(query, top_k, author_filter, century_filter, genre_filter,
               llm_model, do_synthesize):
    if not query.strip():
        yield "<p>Enter a search query.</p>", "", "", [], "", []
        return

    results = semantic_search(
        query,
        top_k=int(top_k),
        author_filter=author_filter or None,
        century_filter=century_filter or None,
        genre_filter=genre_filter or None,
    )

    html = format_results_html(results)
    top_authors = author_summary(results)

    weak_match = bool(results) and results[0]["score"] < WEAK_MATCH_THRESHOLD
    if weak_match:
        html = (
            '<div class="sem-warn">'
            f'⚠️ No strong matches — the best passage scored only '
            f'{results[0]["score"]:.2f}. The corpus may not contain this '
            'concept; the results below are the nearest retrieved text, not '
            'confirmed matches.</div>'
        ) + html

    # New search grounds a fresh chat: store results/query, clear history.
    # Results render immediately; synthesis (if requested) streams in after.
    yield html, top_authors, "", results, query, []
    if do_synthesize and results:
        if weak_match:
            note = (
                f"*Synthesis skipped: no retrieved passage scored above "
                f"{WEAK_MATCH_THRESHOLD:.2f}, so the corpus likely does not "
                f"address this concept directly.*"
            )
            yield html, top_authors, note, results, query, []
        else:
            for acc in synthesize_stream(query, results, llm_model):
                yield html, top_authors, acc, results, query, []


def chat_respond(message, history, llm_model, query, results, synthesis):
    """Chat about the last search's passages; streams the assistant reply."""
    message = (message or "").strip()
    if not message:
        yield history, ""
        return
    history = history + [{"role": "user", "content": message}]
    if not results:
        history.append({
            "role": "assistant",
            "content": "Run a search first — then I can discuss the retrieved passages with you.",
        })
        yield history, ""
        return

    messages = [
        {"role": "system",
         "content": chat_system_prompt(query, results, synthesis or "")}
    ] + history
    history.append({"role": "assistant", "content": ""})
    for acc in llm_chat_stream(messages, llm_model):
        history[-1]["content"] = acc
        yield history, ""


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="ζήτημα") as demo:
        gr.Markdown(
            "# Zetema\n"
            "Query the ancient Greek corpus using natural language. "
            "Passages are retrieved by embedding similarity, then reranked "
            "with a cross-encoder for final ordering.",
            elem_classes="sem-title",
        )

        with gr.Row():
            with gr.Column(scale=3):
                query_box = gr.Textbox(
                    label="Search query (in English or Greek)",
                    placeholder="e.g. 'the immortality of the soul', 'ψυχή', 'rhetoric and democracy'…",
                    lines=2,
                )
            with gr.Column(scale=1):
                search_btn = gr.Button("Search", variant="primary", size="md")

        with gr.Row():
            top_k = gr.Slider(5, 50, value=15, step=5, label="Number of results")
            author_filter = gr.Dropdown(
                choices=_authors,
                multiselect=True,
                label="Filter by author",
                value=None,
            )
            century_filter = gr.Dropdown(
                choices=_century_labels,
                multiselect=True,
                label="Filter by century",
                value=None,
            )
            genre_filter = gr.Dropdown(
                choices=_genre_labels,
                multiselect=True,
                label="Filter by genre",
                value=None,
            )

        with gr.Row():
            llm_model = gr.Dropdown(
                choices=LLM_MODELS,
                value=DEFAULT_LLM,
                label="LLM for synthesis / chat (loaded on first use)",
            )
            do_synthesize = gr.Checkbox(
                label="Synthesize results",
                value=True,
            )

        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("### Retrieved passages")
                results_html = gr.HTML()
            with gr.Column(scale=1):
                gr.Markdown("### Top authors on this topic")
                top_authors_md = gr.Markdown()
                gr.Markdown("### LLM synthesis")
                synthesis_md = gr.Markdown()

        gr.Markdown("### Chat about these results")
        chatbot = gr.Chatbot(
            height=400,
            label="Grounded in the retrieved passages — cleared on each new search",
        )
        with gr.Row():
            chat_input = gr.Textbox(
                placeholder="Ask about the retrieved passages — interpretation, translation, comparisons…",
                show_label=False,
                scale=5,
            )
            chat_send = gr.Button("Send", scale=1)

        results_state = gr.State([])
        query_state = gr.State("")

        search_outputs = [
            results_html, top_authors_md, synthesis_md,
            results_state, query_state, chatbot,
        ]
        search_inputs = [
            query_box, top_k, author_filter, century_filter, genre_filter,
            llm_model, do_synthesize,
        ]
        search_btn.click(fn=run_search, inputs=search_inputs, outputs=search_outputs)
        query_box.submit(fn=run_search, inputs=search_inputs, outputs=search_outputs)

        # synthesis_md is passed by value: the chat needs to see the summary
        # the user is currently reading.
        chat_inputs = [chat_input, chatbot, llm_model, query_state,
                       results_state, synthesis_md]
        chat_send.click(fn=chat_respond, inputs=chat_inputs, outputs=[chatbot, chat_input])
        chat_input.submit(fn=chat_respond, inputs=chat_inputs, outputs=[chatbot, chat_input])

        # Live scheme switcher: restyles cards AND app chrome instantly
        # (pure client-side; per-session, not persisted).
        scheme_picker = gr.Radio(
            choices=list(SCHEME_CLASSES),
            value="Parchment",
            label="Color scheme",
        )
        scheme_picker.change(fn=None, inputs=scheme_picker, outputs=None,
                             js=SCHEME_SWITCH_JS)

    return demo


def main() -> None:
    print("Loading index...")
    load_index()
    load_reranker()
    if IS_SPACE:
        # ZeroGPU functions have a time budget: pay the LLM download/load at
        # startup (CPU side) instead of inside the first @gpu call.
        get_llm(DEFAULT_LLM)
    demo = build_ui()
    # Spaces health-checks port 7860 and doesn't always set GRADIO_SERVER_PORT.
    default_port = "7860" if IS_SPACE else "7861"
    demo.launch(server_name="0.0.0.0",
                server_port=int(os.environ.get("GRADIO_SERVER_PORT", default_port)),
                inbrowser=not IS_SPACE,
                theme=gr.themes.Soft(), css=APP_CSS)


if __name__ == "__main__":
    main()