app.py

# Dependencies (add to requirements.txt on Spaces)
# gradio pypdf pypdfium2 langchain-text-splitters sentence-transformers faiss-cpu transformers requests

import base64
import hmac
import hashlib
import io
import json
import os
import pickle
import re
import textwrap
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from threading import Lock
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import parse_qs, urlparse
from uuid import uuid4

import faiss
import gradio as gr
import numpy as np
import pypdfium2 as pdfium
import requests
try:
    from gradio_rangeslider import RangeSlider
except ImportError:
    RangeSlider = None
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from starlette.middleware import Middleware


# ---------------------------------------------------------------------------
# Persistent storage paths
# ---------------------------------------------------------------------------
def _can_write_to_dir(path: Path) -> bool:
    try:
        path.mkdir(parents=True, exist_ok=True)
        probe = path / ".write_probe"
        with probe.open("w", encoding="utf-8") as fh:
            fh.write("")
        probe.unlink(missing_ok=True)
        return True
    except OSError:
        return False


def resolve_data_dir() -> Path:
    configured_dir = os.getenv("PERSIST_DIR", "").strip()
    candidates: List[Path] = []
    if configured_dir:
        candidates.append(Path(configured_dir))
    candidates.extend(
        [
            Path("/data"),
            Path.cwd() / ".data",
            Path("/tmp/rp-chatbot-data"),
        ]
    )

    seen_paths = set()
    for candidate in candidates:
        candidate_key = str(candidate)
        if candidate_key in seen_paths:
            continue
        seen_paths.add(candidate_key)
        if _can_write_to_dir(candidate):
            return candidate

    raise RuntimeError(
        "No writable data directory found. Set PERSIST_DIR to a writable path."
    )


DATA_DIR = resolve_data_dir()
UPLOAD_DIR = DATA_DIR / "uploads"
INDEX_PATH_TEMPLATE = "docs_{strategy}.faiss"
META_PATH_TEMPLATE = "meta_{strategy}.pkl"
LEGACY_INDEX_PATH = DATA_DIR / "docs.faiss"
LEGACY_META_PATH = DATA_DIR / "meta.pkl"
REGISTRY_PATH = DATA_DIR / "registry.json"
FEEDBACK_PATH = DATA_DIR / "feedback.json"

UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

# Thread-safety when running on Spaces
INDEX_LOCK = Lock()
REGISTRY_LOCK = Lock()
FEEDBACK_LOCK = Lock()
EMBEDDER_LOCK = Lock()
MIGRATION_LOCK = Lock()
MIGRATION_DONE = False

# Embedding strategies (cosine similarity via normalized embeddings)
EMBEDDING_STRATEGIES: Dict[str, Dict[str, str]] = {
    "fast_minilm_l6": {
        "label": "Fast MiniLM-L6",
        "model": os.getenv(
            "EMBED_MODEL_FAST", "sentence-transformers/all-MiniLM-L6-v2"
        ),
    },
    "quality_mpnet": {
        "label": "Quality MPNet",
        "model": os.getenv("EMBED_MODEL_QUALITY", "sentence-transformers/all-mpnet-base-v2"),
    },
    "qa_multilingual": {
        "label": "QA Multilingual",
        "model": os.getenv(
            "EMBED_MODEL_QA", "sentence-transformers/multi-qa-mpnet-base-dot-v1"
        ),
    },
}
EMBEDDING_OPTIONS = [cfg["label"] for cfg in EMBEDDING_STRATEGIES.values()]
EMBEDDING_LABEL_TO_KEY = {
    cfg["label"]: key for key, cfg in EMBEDDING_STRATEGIES.items()
}
DEFAULT_EMBEDDING_KEY = next(iter(EMBEDDING_STRATEGIES.keys()))
DEFAULT_EMBEDDING_LABEL = EMBEDDING_STRATEGIES[DEFAULT_EMBEDDING_KEY]["label"]
EMBEDDERS: Dict[str, SentenceTransformer] = {}

# Text splitter config (tweak via env if needed)
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "800"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "120"))
MIN_CHUNK_TEXT_LENGTH = int(os.getenv("MIN_CHUNK_TEXT_LENGTH", "80"))
ACCESS_QUERY_PARAM = os.getenv("APP_ACCESS_QUERY_PARAM", "token")
URL_ACCESS_TOKEN = os.getenv("APP_URL_TOKEN", "").strip()
ACCESS_COOKIE_NAME = os.getenv("APP_ACCESS_COOKIE_NAME", "rag_access")
ALLOW_LOOPBACK_BYPASS = os.getenv("APP_ALLOW_LOOPBACK_BYPASS", "1").strip().lower() in {
    "1",
    "true",
    "yes",
    "on",
}
IS_HF_SPACE = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID"))
ENABLE_TOKEN_GATE_ON_SPACES = os.getenv(
    "APP_ENABLE_TOKEN_GATE_ON_SPACES", "0"
).strip().lower() in {"1", "true", "yes", "on"}
TOKEN_GATE_ENABLED = bool(URL_ACCESS_TOKEN) and (
    not IS_HF_SPACE or ENABLE_TOKEN_GATE_ON_SPACES
)

def _get_env_float(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, default))
    except (TypeError, ValueError):
        return default


RAG_MIN_SCORE = _get_env_float("RAG_MIN_SCORE", 0.15)


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


def utc_year() -> int:
    return utc_now().year


def utc_now_iso_z() -> str:
    return utc_now().isoformat(timespec="seconds").replace("+00:00", "Z")


DEFAULT_CHAT_MODEL = os.getenv(
    "DEEPINFRA_RAG_MODEL", "meta-llama/Meta-Llama-3.1-8B-Instruct"
)

REASONING_PRESETS = {
    "Fast (8B)": {
        "model": os.getenv(
            "DEEPINFRA_FAST_MODEL", "meta-llama/Meta-Llama-3.1-8B-Instruct"
        ),
        "temperature": _get_env_float("DEEPINFRA_FAST_TEMP", 0.1),
        "extra": {},
    },
    "Balanced (70B)": {
        "model": os.getenv(
            "DEEPINFRA_BALANCED_MODEL", "meta-llama/Meta-Llama-3.1-70B-Instruct"
        ),
        "temperature": _get_env_float("DEEPINFRA_BALANCED_TEMP", 0.15),
        "extra": {},
    },
    "Deep Reasoning (R1 Distill)": {
        "model": os.getenv(
            "DEEPINFRA_REASONING_MODEL", "deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
        ),
        "temperature": _get_env_float("DEEPINFRA_REASONING_TEMP", 0.2),
        "extra": {"reasoning": {"effort": "medium"}},
    },
}


# ---------------------------------------------------------------------------
# Registry helpers
# ---------------------------------------------------------------------------
def _load_registry() -> Dict:
    if not REGISTRY_PATH.exists():
        return {"documents": {}}
    with REGISTRY_PATH.open("r", encoding="utf-8") as fh:
        return json.load(fh)


def _save_registry(data: Dict) -> None:
    with REGISTRY_PATH.open("w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2, ensure_ascii=True)


def load_registry_threadsafe() -> Dict:
    with REGISTRY_LOCK:
        return _load_registry()


def update_registry(doc_id: str, payload: Dict) -> None:
    with REGISTRY_LOCK:
        registry = _load_registry()
        registry.setdefault("documents", {})
        registry["documents"][doc_id] = payload
        _save_registry(registry)


def remove_from_registry(doc_id: str) -> Optional[Dict]:
    with REGISTRY_LOCK:
        registry = _load_registry()
        documents = registry.get("documents", {})
        document = documents.pop(doc_id, None)
        registry["documents"] = documents
        _save_registry(registry)
    return document


def _load_feedback_registry() -> Dict:
    if not FEEDBACK_PATH.exists():
        return {"feedback": {}}
    with FEEDBACK_PATH.open("r", encoding="utf-8") as fh:
        data = json.load(fh)
    if not isinstance(data, dict):
        return {"feedback": {}}
    feedback = data.get("feedback")
    if not isinstance(feedback, dict):
        data["feedback"] = {}
    return data


def _save_feedback_registry(data: Dict) -> None:
    with FEEDBACK_PATH.open("w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2, ensure_ascii=True)


def load_feedback_registry_threadsafe() -> Dict:
    with FEEDBACK_LOCK:
        return _load_feedback_registry()


def create_feedback_entry(payload: Dict) -> Dict:
    with FEEDBACK_LOCK:
        store = _load_feedback_registry()
        feedback_map = store.setdefault("feedback", {})
        feedback_id = payload.get("feedback_id")
        if not feedback_id:
            feedback_id = f"fb_{utc_now().strftime('%Y%m%d%H%M%S')}_{uuid4().hex[:8]}"
        payload["feedback_id"] = feedback_id
        feedback_map[feedback_id] = payload
        _save_feedback_registry(store)
    return payload


def get_feedback_entry(feedback_id: str) -> Optional[Dict]:
    feedback_id = (feedback_id or "").strip()
    if not feedback_id:
        return None
    store = load_feedback_registry_threadsafe()
    feedback_map = store.get("feedback", {})
    entry = feedback_map.get(feedback_id)
    return dict(entry) if isinstance(entry, dict) else None


def delete_feedback_entry(feedback_id: str) -> Dict:
    feedback_id = (feedback_id or "").strip()
    if not feedback_id:
        return {"status": "error", "message": "Feedback ID is missing."}
    with FEEDBACK_LOCK:
        store = _load_feedback_registry()
        feedback_map = store.get("feedback", {})
        removed = feedback_map.pop(feedback_id, None)
        store["feedback"] = feedback_map
        _save_feedback_registry(store)
    if not removed:
        return {"status": "error", "message": f"Feedback {feedback_id} not found."}
    return {"status": "deleted", "message": f"Deleted feedback {feedback_id}."}


def mark_feedback_resolved(feedback_id: str) -> Dict:
    feedback_id = (feedback_id or "").strip()
    if not feedback_id:
        return {"status": "error", "message": "Feedback ID is missing."}
    with FEEDBACK_LOCK:
        store = _load_feedback_registry()
        feedback_map = store.get("feedback", {})
        entry = feedback_map.get(feedback_id)
        if not isinstance(entry, dict):
            return {"status": "error", "message": f"Feedback {feedback_id} not found."}
        if entry.get("status") == "resolved":
            return {"status": "noop", "message": f"Feedback {feedback_id} is already resolved."}
        entry["status"] = "resolved"
        entry["resolved_at"] = utc_now_iso_z()
        feedback_map[feedback_id] = entry
        store["feedback"] = feedback_map
        _save_feedback_registry(store)
    return {"status": "resolved", "message": f"Marked feedback {feedback_id} as resolved."}


def compute_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as fh:
        for block in iter(lambda: fh.read(8192), b""):
            digest.update(block)
    return digest.hexdigest()


# ---------------------------------------------------------------------------
# Vector index helpers
# ---------------------------------------------------------------------------
def strategy_paths(strategy_key: str) -> Tuple[Path, Path]:
    index_path = DATA_DIR / INDEX_PATH_TEMPLATE.format(strategy=strategy_key)
    meta_path = DATA_DIR / META_PATH_TEMPLATE.format(strategy=strategy_key)
    return index_path, meta_path


def load_embedder(strategy_key: str) -> SentenceTransformer:
    with EMBEDDER_LOCK:
        existing = EMBEDDERS.get(strategy_key)
        if existing is not None:
            return existing
        strategy = EMBEDDING_STRATEGIES.get(strategy_key)
        if not strategy:
            raise ValueError(f"Unknown embedding strategy: {strategy_key}")
        model_id = strategy["model"]
        embedder = SentenceTransformer(model_id)
        EMBEDDERS[strategy_key] = embedder
        return embedder


def encode_texts(texts: List[str], strategy_key: str) -> np.ndarray:
    embedder = load_embedder(strategy_key)
    vectors = embedder.encode(texts, normalize_embeddings=True)
    return np.asarray(vectors, dtype="float32")


def load_index(strategy_key: str) -> Tuple[Optional[faiss.Index], List[Dict]]:
    index_path, meta_path = strategy_paths(strategy_key)
    if not index_path.exists() or not meta_path.exists():
        return None, []
    index = faiss.read_index(str(index_path))
    with meta_path.open("rb") as fh:
        metadata = pickle.load(fh)
    return index, metadata


def save_index(strategy_key: str, index: faiss.Index, metadata: List[Dict]) -> None:
    index_path, meta_path = strategy_paths(strategy_key)
    faiss.write_index(index, str(index_path))
    with meta_path.open("wb") as fh:
        pickle.dump(metadata, fh)


def clear_legacy_index_files() -> None:
    for path in (LEGACY_INDEX_PATH, LEGACY_META_PATH):
        try:
            path.unlink()
        except FileNotFoundError:
            continue


def get_registry_doc_ids() -> set[str]:
    registry = load_registry_threadsafe()
    documents = registry.get("documents", {})
    return set(documents.keys())


def ensure_strategy_indexes_initialized() -> None:
    """
    Migrates legacy single-index storage into all configured embedding strategies.
    """
    global MIGRATION_DONE
    with MIGRATION_LOCK:
        if MIGRATION_DONE:
            return
        MIGRATION_DONE = True

        if not (LEGACY_INDEX_PATH.exists() and LEGACY_META_PATH.exists()):
            return

        has_new_layout = False
        for strategy_key in EMBEDDING_STRATEGIES:
            index_path, meta_path = strategy_paths(strategy_key)
            if index_path.exists() and meta_path.exists():
                has_new_layout = True
                break
        if has_new_layout:
            return

        try:
            with LEGACY_META_PATH.open("rb") as fh:
                legacy_metadata = pickle.load(fh)
        except Exception:
            return

        if not isinstance(legacy_metadata, list) or not legacy_metadata:
            clear_legacy_index_files()
            return

        active_doc_ids = get_registry_doc_ids()
        if active_doc_ids:
            legacy_metadata = [
                item for item in legacy_metadata if item.get("doc_id") in active_doc_ids
            ]
        else:
            legacy_metadata = []

        if not legacy_metadata:
            clear_index_files()
            clear_legacy_index_files()
            return

        for item in legacy_metadata:
            try:
                int(item.get("year"))
            except (TypeError, ValueError):
                item["year"] = get_registry_doc_year(item.get("doc_id")) or utc_year()

        for strategy_key in EMBEDDING_STRATEGIES:
            strategy_metadata = [dict(item) for item in legacy_metadata]
            rebuild_index_from_metadata(strategy_key, strategy_metadata)
        clear_legacy_index_files()


# ---------------------------------------------------------------------------
# OCR client for DeepInfra DeepSeek OCR
# ---------------------------------------------------------------------------
class DeepInfraOCR:
    def __init__(self) -> None:
        self.api_token = os.getenv("DEEPINFRA")
        self.model_id = os.getenv("DEEPINFRA_OCR_MODEL", "deepseek-ai/deepseek-ocr")
        base_url = os.getenv("DEEPINFRA_API_URL", "https://api.deepinfra.com/v1/inference")
        self.endpoint = base_url.rstrip("/") + f"/{self.model_id}"
        self.timeout = int(os.getenv("DEEPINFRA_TIMEOUT", "120"))

    def available(self) -> bool:
        return bool(self.api_token)

    def infer(self, image) -> str:
        """
        Runs OCR on a single PIL image. Expects DeepInfra API to respond with JSON
        containing the recognized text.
        """
        if not self.available():
            raise RuntimeError("DeepInfra token not configured (DEEPINFRA).")

        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

        payload = {"input": {"image": image_base64}}
        headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json",
        }

        response = requests.post(
            self.endpoint,
            headers=headers,
            json=payload,
            timeout=self.timeout,
        )
        response.raise_for_status()
        data = response.json()
        text = self._extract_text(data)
        if not text:
            raise ValueError(f"DeepInfra OCR returned an empty result: {data}")
        return text.strip()

    @staticmethod
    def _extract_text(data: Dict) -> Optional[str]:
        """
        DeepInfra responses may vary; try a few common response shapes.
        """
        if data is None:
            return None
        if isinstance(data, str):
            return data

        for key in ("text", "output", "result"):
            value = data.get(key)
            if isinstance(value, str):
                return value
            if isinstance(value, dict):
                nested = (
                    value.get("text")
                    or value.get("output")
                    or value.get("result")
                    or value.get("content")
                )
                if isinstance(nested, str):
                    return nested

        results = data.get("results") or data.get("outputs") or data.get("choices")
        if isinstance(results, list) and results:
            first = results[0]
            if isinstance(first, str):
                return first
            if isinstance(first, dict):
                return (
                    first.get("text")
                    or first.get("output")
                    or first.get("result")
                    or first.get("content")
                    or first.get("message", {}).get("content")
                )

        return None


OCR_CLIENT = DeepInfraOCR()


# ---------------------------------------------------------------------------
# PDF text extraction & chunking
# ---------------------------------------------------------------------------
def extract_pdf_text(pdf_path: Path) -> List[Dict]:
    """
    Returns a list of {"page": int, "text": str} dictionaries.
    Falls back to OCR for pages where text extraction fails.
    """
    reader = PdfReader(str(pdf_path))
    pages: List[Dict] = []
    ocr_candidates: List[int] = []

    for idx, page in enumerate(reader.pages):
        raw_text = (page.extract_text() or "").strip()
        if len(raw_text) > 30:
            pages.append({"page": idx + 1, "text": raw_text})
        else:
            pages.append({"page": idx + 1, "text": ""})
            ocr_candidates.append(idx)

    if ocr_candidates and not OCR_CLIENT.available():
        # Leave blank text; user will be notified about missing OCR capability.
        return pages

    if ocr_candidates:
        pdf = pdfium.PdfDocument(str(pdf_path))
        for page_index in ocr_candidates:
            page = pdf.get_page(page_index)
            bitmap = page.render(scale=2.0, rotation=0)
            image = bitmap.to_pil()
            try:
                text = OCR_CLIENT.infer(image)
            except Exception as exc:
                pages[page_index]["text"] = ""
                pages[page_index]["error"] = str(exc)
            else:
                pages[page_index]["text"] = text

    return pages


def extract_candidate_years(text: str) -> List[int]:
    current_year = utc_year()
    max_year = current_year + 1
    years: List[int] = []
    for match in re.findall(r"\b(19\d{2}|20\d{2}|21\d{2})\b", text or ""):
        year = int(match)
        if 1900 <= year <= max_year:
            years.append(year)
    return years


def infer_publication_year(pdf_path: Path, doc_name: str, pages: List[Dict]) -> Tuple[int, str]:
    """
    Extract publication year from metadata and front matter first.
    If unclear, use the most recent year mentioned in the document as a proxy.
    """
    current_year = utc_year()

    metadata_years: List[int] = []
    try:
        reader = PdfReader(str(pdf_path))
        metadata = reader.metadata or {}
        for key, value in metadata.items():
            if value is None:
                continue
            combined = f"{key} {value}"
            metadata_years.extend(extract_candidate_years(str(combined)))
    except Exception:
        pass

    filename_years = extract_candidate_years(doc_name)
    cover_text = "\n".join((p.get("text") or "") for p in pages[:3])
    front_matter_years = extract_candidate_years(cover_text)

    explicit_years = metadata_years + filename_years + front_matter_years
    if explicit_years:
        year_counts = Counter(explicit_years)
        highest_count = max(year_counts.values())
        tied_years = [year for year, count in year_counts.items() if count == highest_count]
        return max(tied_years), "explicit"

    full_text = "\n".join((p.get("text") or "") for p in pages)
    proxy_years = extract_candidate_years(full_text)
    if proxy_years:
        return max(proxy_years), "proxy_latest_mentioned"

    return current_year, "proxy_current_year"


def chunk_pages(
    pages: List[Dict], doc_id: str, doc_name: str, doc_year: int
) -> List[Dict]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
    )
    chunks: List[Dict] = []
    for page in pages:
        text = page.get("text", "").strip()
        if not text:
            continue
        page_chunks = splitter.split_text(text)
        for idx, chunk_text in enumerate(page_chunks, start=1):
            normalized = " ".join(chunk_text.split())
            if len(normalized) < MIN_CHUNK_TEXT_LENGTH:
                continue
            chunks.append(
                {
                    "text": normalized,
                    "metadata": {
                        "doc_id": doc_id,
                        "doc_name": doc_name,
                        "year": doc_year,
                        "page": page["page"],
                        "chunk_number": idx,
                    },
                }
            )
    return chunks


def add_chunks_to_indices(chunks: List[Dict]) -> Tuple[int, Dict[str, int]]:
    if not chunks:
        return 0, {}

    texts = [chunk["text"] for chunk in chunks]
    vectors_by_strategy: Dict[str, np.ndarray] = {}
    for strategy_key in EMBEDDING_STRATEGIES:
        vectors_by_strategy[strategy_key] = encode_texts(texts, strategy_key)

    with INDEX_LOCK:
        end_ids: Dict[str, int] = {}
        for strategy_key, vectors in vectors_by_strategy.items():
            index, metadata = load_index(strategy_key)
            dim = vectors.shape[1]
            if index is None:
                index = faiss.IndexFlatIP(dim)
                metadata = []
            elif index.d != dim:
                existing_texts = [item["text"] for item in metadata]
                index = faiss.IndexFlatIP(dim)
                if existing_texts:
                    reencoded = encode_texts(existing_texts, strategy_key)
                    index.add(reencoded)
                for row_idx, item in enumerate(metadata):
                    item["vector_id"] = row_idx

            start_idx = len(metadata)
            index.add(vectors)
            for offset, chunk in enumerate(chunks):
                chunk_meta = {
                    "vector_id": start_idx + offset,
                    "text": chunk["text"],
                }
                chunk_meta.update(chunk["metadata"])
                metadata.append(chunk_meta)
            save_index(strategy_key, index, metadata)
            end_ids[strategy_key] = start_idx + len(chunks)

    added = len(chunks)
    return added, end_ids


# ---------------------------------------------------------------------------
# Ingestion workflow
# ---------------------------------------------------------------------------
def ingest_pdf(path_str: str) -> Dict:
    ensure_strategy_indexes_initialized()
    pdf_path = Path(path_str)
    if not pdf_path.exists():
        return {"status": "error", "message": f"File not found: {pdf_path}"}

    sha256 = compute_sha256(pdf_path)
    doc_id = sha256[:12]
    original_name = pdf_path.name
    registry = load_registry_threadsafe()
    existing = registry.get("documents", {}).get(doc_id)
    if existing:
        return {
            "status": "skipped",
            "doc_id": doc_id,
            "message": f"Already indexed ({existing['title']})",
        }

    persisted_path = UPLOAD_DIR / f"{doc_id}.pdf"
    if not persisted_path.exists():
        persisted_path.write_bytes(pdf_path.read_bytes())

    pages = extract_pdf_text(persisted_path)
    blank_pages = [p["page"] for p in pages if not p.get("text")]

    publication_year, year_source = infer_publication_year(
        persisted_path, original_name, pages
    )

    chunks = chunk_pages(
        pages, doc_id=doc_id, doc_name=original_name, doc_year=publication_year
    )
    added_chunks, end_ids = add_chunks_to_indices(chunks)

    document_record = {
        "doc_id": doc_id,
        "sha256": sha256,
        "title": original_name,
        "stored_path": str(persisted_path.relative_to(DATA_DIR)),
        "year": publication_year,
        "year_source": year_source,
        "pages": len(pages),
        "blank_pages": blank_pages,
        "chunks": added_chunks,
        "first_vector_ids": {
            key: (end_ids[key] - added_chunks) for key in end_ids
        },
        "ingested_at": utc_now_iso_z(),
    }
    update_registry(doc_id, document_record)

    if blank_pages and OCR_CLIENT.available():
        message = (
            f"Indexed {original_name} ({publication_year}, {added_chunks} chunks). "
            f"OCR ran on pages: {blank_pages}"
        )
    elif blank_pages:
        message = (
            f"Indexed {original_name} ({publication_year}, {added_chunks} chunks) but "
            f"{len(blank_pages)} pages lacked text (set the DEEPINFRA token to enable OCR)."
        )
    else:
        message = f"Indexed {original_name} ({publication_year}, {added_chunks} chunks)."

    return {"status": "indexed", "doc_id": doc_id, "message": message}


def clear_index_files(strategy_key: Optional[str] = None) -> None:
    strategy_keys = [strategy_key] if strategy_key else list(EMBEDDING_STRATEGIES.keys())
    for key in strategy_keys:
        index_path, meta_path = strategy_paths(key)
        for path in (index_path, meta_path):
            try:
                path.unlink()
            except FileNotFoundError:
                continue


def rebuild_index_from_metadata(strategy_key: str, metadata: List[Dict]) -> None:
    if not metadata:
        clear_index_files(strategy_key=strategy_key)
        return
    texts = [m["text"] for m in metadata]
    vectors = encode_texts(texts, strategy_key)
    dim = vectors.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(vectors)
    for idx, meta in enumerate(metadata):
        meta["vector_id"] = idx
    save_index(strategy_key, index, metadata)


def rebuild_index_after_delete(
    strategy_key: str, index: faiss.Index, metadata: List[Dict], removed_doc_id: str
) -> None:
    if index is None or not metadata:
        clear_index_files(strategy_key=strategy_key)
        return

    total_vectors = int(getattr(index, "ntotal", 0))
    usable_count = min(len(metadata), total_vectors)
    if usable_count <= 0:
        clear_index_files(strategy_key=strategy_key)
        return

    remaining_metadata: List[Dict] = []
    remaining_vectors: List[np.ndarray] = []

    for vector_idx in range(usable_count):
        item = metadata[vector_idx]
        if item.get("doc_id") == removed_doc_id:
            continue

        remaining_metadata.append(dict(item))
        vector = index.reconstruct(vector_idx)
        remaining_vectors.append(np.asarray(vector, dtype="float32"))

    if not remaining_metadata:
        clear_index_files(strategy_key=strategy_key)
        return

    matrix = np.vstack(remaining_vectors).astype("float32")
    dim = matrix.shape[1]
    rebuilt_index = faiss.IndexFlatIP(dim)
    rebuilt_index.add(matrix)
    for new_idx, item in enumerate(remaining_metadata):
        item["vector_id"] = new_idx
    save_index(strategy_key, rebuilt_index, remaining_metadata)


def delete_document(doc_id: str) -> Dict:
    ensure_strategy_indexes_initialized()
    document = remove_from_registry(doc_id)
    if not document:
        return {"status": "error", "doc_id": doc_id, "message": "Document not found."}

    title = document.get("title", doc_id)

    try:
        with INDEX_LOCK:
            for strategy_key in EMBEDDING_STRATEGIES:
                index, metadata = load_index(strategy_key)
                if not metadata or index is None:
                    clear_index_files(strategy_key=strategy_key)
                    continue
                rebuild_index_after_delete(
                    strategy_key=strategy_key,
                    index=index,
                    metadata=metadata,
                    removed_doc_id=doc_id,
                )

            if not get_registry_doc_ids():
                clear_index_files()
                clear_legacy_index_files()
    except Exception as exc:
        # Keep registry and index state consistent when deletion fails.
        update_registry(doc_id, document)
        return {
            "status": "error",
            "doc_id": doc_id,
            "message": f"Failed to remove {title}: {exc}",
        }

    stored_path = document.get("stored_path")
    if stored_path:
        stored_file = DATA_DIR / stored_path
        if stored_file.exists():
            try:
                stored_file.unlink()
            except Exception:
                pass

    message = f"Removed {title}."
    return {"status": "deleted", "doc_id": doc_id, "message": message}


def get_registry_doc_year(doc_id: Optional[str]) -> Optional[int]:
    if not doc_id:
        return None
    registry = load_registry_threadsafe()
    doc = registry.get("documents", {}).get(doc_id, {})
    year_value = doc.get("year")
    try:
        return int(year_value)
    except (TypeError, ValueError):
        return None


def resolve_hit_year(hit: Dict) -> int:
    try:
        return int(hit.get("year"))
    except (TypeError, ValueError):
        pass
    registry_year = get_registry_doc_year(hit.get("doc_id"))
    if registry_year is not None:
        hit["year"] = registry_year
        return registry_year
    fallback = utc_year()
    hit["year"] = fallback
    return fallback


def select_diverse_hits(hits: List[Dict], top_k: int) -> List[Dict]:
    if len(hits) <= top_k:
        return hits
    per_doc_cap = max(2, top_k // 2)
    selected: List[Dict] = []
    selected_ids = set()
    per_doc_counts: Dict[str, int] = {}

    for hit in hits:
        hit_id = (
            hit.get("doc_id"),
            hit.get("page"),
            hit.get("chunk_number"),
            hit.get("vector_id"),
        )
        if hit_id in selected_ids:
            continue
        doc_id = hit.get("doc_id", "")
        if per_doc_counts.get(doc_id, 0) >= per_doc_cap:
            continue
        selected.append(hit)
        selected_ids.add(hit_id)
        per_doc_counts[doc_id] = per_doc_counts.get(doc_id, 0) + 1
        if len(selected) >= top_k:
            return selected

    for hit in hits:
        if len(selected) >= top_k:
            break
        hit_id = (
            hit.get("doc_id"),
            hit.get("page"),
            hit.get("chunk_number"),
            hit.get("vector_id"),
        )
        if hit_id in selected_ids:
            continue
        selected.append(hit)
        selected_ids.add(hit_id)

    return selected


def _parse_year_range(year_from: Optional[int], year_to: Optional[int]) -> Tuple[int, int]:
    current_year = utc_year()
    from_year = int(year_from) if year_from is not None else 1900
    to_year = int(year_to) if year_to is not None else current_year + 1
    if from_year > to_year:
        from_year, to_year = to_year, from_year
    return from_year, to_year


def parse_year_range_input(year_range: Any) -> Tuple[int, int]:
    year_from: Optional[int] = None
    year_to: Optional[int] = None

    if isinstance(year_range, (list, tuple)):
        if len(year_range) >= 1:
            try:
                year_from = int(year_range[0])
            except (TypeError, ValueError):
                year_from = None
        if len(year_range) >= 2:
            try:
                year_to = int(year_range[1])
            except (TypeError, ValueError):
                year_to = None
        elif year_from is not None:
            year_to = year_from
    elif year_range is not None:
        try:
            year_from = int(year_range)
            year_to = year_from
        except (TypeError, ValueError):
            year_from, year_to = None, None

    return _parse_year_range(year_from, year_to)


# ---------------------------------------------------------------------------
# Retrieval & LLM answer generation
# ---------------------------------------------------------------------------
def resolve_embedding_selection(selected_label: Optional[str]) -> Tuple[str, str]:
    strategy_key = EMBEDDING_LABEL_TO_KEY.get(
        (selected_label or "").strip(), DEFAULT_EMBEDDING_KEY
    )
    strategy_label = EMBEDDING_STRATEGIES[strategy_key]["label"]
    return strategy_key, strategy_label


def search_index(
    query: str,
    top_k: int = 5,
    strategy_key: str = DEFAULT_EMBEDDING_KEY,
    year_from: Optional[int] = None,
    year_to: Optional[int] = None,
) -> List[Dict]:
    ensure_strategy_indexes_initialized()
    active_doc_ids = get_registry_doc_ids()
    if not active_doc_ids:
        return []

    from_year, to_year = _parse_year_range(year_from, year_to)
    with INDEX_LOCK:
        index, metadata = load_index(strategy_key)
    if index is None or not metadata:
        return []

    query_vector = encode_texts([query], strategy_key)
    fetch_k = min(len(metadata), max(top_k * 8, 100))
    if fetch_k < top_k:
        fetch_k = min(len(metadata), top_k)

    scores, ids = index.search(query_vector, fetch_k)
    candidates: List[Dict] = []
    for score, idx in zip(scores[0], ids[0]):
        if idx == -1:
            continue
        meta = metadata[idx]
        doc_id = meta.get("doc_id")
        if doc_id not in active_doc_ids:
            continue
        hit = {**meta, "score": float(score)}
        year_value = resolve_hit_year(hit)
        if from_year <= year_value <= to_year:
            candidates.append(hit)

    if not candidates:
        return []

    ranked = sorted(candidates, key=lambda item: item.get("score", 0.0), reverse=True)
    high_confidence = [item for item in ranked if item.get("score", 0.0) >= RAG_MIN_SCORE]
    retrieval_pool = high_confidence if high_confidence else ranked
    return select_diverse_hits(retrieval_pool, top_k)[:top_k]


class DeepInfraChat:
    def __init__(self) -> None:
        self.api_token = os.getenv("DEEPINFRA")
        self.model_id = DEFAULT_CHAT_MODEL
        self.endpoint = os.getenv(
            "DEEPINFRA_CHAT_URL", "https://api.deepinfra.com/v1/openai/chat/completions"
        )
        self.timeout = int(os.getenv("DEEPINFRA_TIMEOUT", "120"))

    def available(self) -> bool:
        return bool(self.api_token)

    def generate(
        self,
        messages: List[Dict],
        temperature: float = 0.1,
        model_override: Optional[str] = None,
        extra_payload: Optional[Dict] = None,
    ) -> str:
        if not self.available():
            raise RuntimeError("Configure DEEPINFRA to enable answering.")
        payload = {
            "model": model_override or self.model_id,
            "messages": messages,
            "temperature": temperature,
        }

        if extra_payload:
            payload.update(extra_payload)

        headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json",
        }
        response = requests.post(
            self.endpoint, headers=headers, json=payload, timeout=self.timeout
        )
        response.raise_for_status()
        data = response.json()
        choices = data.get("choices")
        if not choices:
            raise ValueError(f"DeepInfra chat returned no choices: {data}")
        content = choices[0].get("message", {}).get("content")
        if not content:
            raise ValueError(f"No content in response: {data}")
        return content.strip()


CHAT_CLIENT = DeepInfraChat()


def _token_matches(candidate: str) -> bool:
    if not URL_ACCESS_TOKEN:
        return False
    if not candidate:
        return False
    return hmac.compare_digest(candidate.strip(), URL_ACCESS_TOKEN)


def _query_tokens_from_raw_query(raw_query: str) -> List[str]:
    tokens: List[str] = []

    parsed = parse_qs(raw_query, keep_blank_values=True)
    tokens.extend(token for token in parsed.get(ACCESS_QUERY_PARAM, []) if token)

    # Keep raw values too so '+' and similar characters do not break matching.
    prefix = f"{ACCESS_QUERY_PARAM}="
    for pair in raw_query.split("&"):
        if pair.startswith(prefix):
            raw_value = pair[len(prefix) :].strip()
            if raw_value:
                tokens.append(raw_value)

    return tokens


def _query_tokens_from_scope(scope: Dict[str, Any]) -> List[str]:
    raw_query = (scope.get("query_string") or b"").decode("utf-8", errors="ignore")
    if not raw_query:
        return []
    return _query_tokens_from_raw_query(raw_query)


def _referer_query_tokens(scope: Dict[str, Any]) -> List[str]:
    for key, value in scope.get("headers", []):
        if key == b"referer":
            referer = value.decode("latin-1")
            parsed = urlparse(referer)
            return _query_tokens_from_raw_query(parsed.query)
    return []


def _cookies_from_scope(scope: Dict[str, Any]) -> Dict[str, str]:
    cookies: Dict[str, str] = {}
    for key, value in scope.get("headers", []):
        if key == b"cookie":
            header = value.decode("latin-1")
            for chunk in header.split(";"):
                if "=" not in chunk:
                    continue
                cookie_key, cookie_value = chunk.strip().split("=", 1)
                cookies[cookie_key] = cookie_value
    return cookies


def _decode_cookie_token(cookie_value: str) -> List[str]:
    if not cookie_value:
        return []

    tokens = [cookie_value]
    if not cookie_value.startswith("b64."):
        return tokens

    encoded = cookie_value[4:]
    padded = encoded + ("=" * (-len(encoded) % 4))
    try:
        decoded = base64.urlsafe_b64decode(padded.encode("ascii")).decode("utf-8")
    except Exception:
        return tokens

    if decoded:
        tokens.append(decoded)
    return tokens


def _is_loopback_client(scope: Dict[str, Any]) -> bool:
    client = scope.get("client")
    if isinstance(client, (tuple, list)) and client:
        host = str(client[0]).lower()
        if host in {"127.0.0.1", "::1", "localhost"}:
            return True
    for key, value in scope.get("headers", []):
        if key == b"host":
            host_header = value.decode("latin-1").lower()
            if host_header.startswith("localhost:") or host_header.startswith("127.0.0.1:"):
                return True
    return False


class QueryTokenGateMiddleware:
    def __init__(self, app) -> None:
        self.app = app

    async def __call__(self, scope, receive, send):
        scope_type = scope.get("type")
        if scope_type not in {"http", "websocket"}:
            await self.app(scope, receive, send)
            return

        if ALLOW_LOOPBACK_BYPASS and _is_loopback_client(scope):
            await self.app(scope, receive, send)
            return

        if not URL_ACCESS_TOKEN:
            await self.app(scope, receive, send)
            return

        # Gate only the landing page to avoid blocking internal Gradio assets/API
        # requests on Spaces, which can lead to a blank screen.
        path = str(scope.get("path") or "/")
        if path not in {"", "/"}:
            await self.app(scope, receive, send)
            return

        query_tokens = _query_tokens_from_scope(scope)
        referer_tokens = _referer_query_tokens(scope)
        cookies = _cookies_from_scope(scope)
        cookie_tokens = _decode_cookie_token(cookies.get(ACCESS_COOKIE_NAME, ""))

        token_candidates = query_tokens + referer_tokens + cookie_tokens
        token_valid = any(_token_matches(candidate) for candidate in token_candidates)
        if not token_valid:
            if scope_type == "websocket":
                await send({"type": "websocket.close", "code": 1008})
                return

            body = f"""<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1" />
    <title>Token Required</title>
    <style>
      body {{
        font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
        margin: 0;
        padding: 24px;
        background: #f7f7f8;
        color: #111827;
      }}
      .card {{
        max-width: 720px;
        margin: 48px auto;
        background: #fff;
        border: 1px solid #e5e7eb;
        border-radius: 12px;
        padding: 20px;
      }}
      code {{
        background: #f3f4f6;
        padding: 2px 6px;
        border-radius: 6px;
      }}
    </style>
  </head>
  <body>
    <div class="card">
      <h2>Token is required</h2>
      <p>Open this Space with a valid token in the URL query:</p>
      <p><code>?{ACCESS_QUERY_PARAM}=YOUR_TOKEN</code></p>
    </div>
  </body>
</html>""".encode("utf-8")
            await send(
                {
                    "type": "http.response.start",
                    "status": 200,
                    "headers": [
                        (b"content-type", b"text/html; charset=utf-8"),
                        (b"cache-control", b"no-store"),
                        (b"content-length", str(len(body)).encode("utf-8")),
                    ],
                }
            )
            await send({"type": "http.response.body", "body": body})
            return

        should_set_cookie = scope_type == "http" and any(
            _token_matches(token) for token in query_tokens
        )
        if not should_set_cookie:
            await self.app(scope, receive, send)
            return

        async def send_wrapper(message):
            if message.get("type") == "http.response.start":
                headers = list(message.get("headers") or [])
                secure_part = "; Secure" if IS_HF_SPACE else ""
                cookie_token = base64.urlsafe_b64encode(
                    URL_ACCESS_TOKEN.encode("utf-8")
                ).decode("ascii")
                cookie_value = (
                    f"{ACCESS_COOKIE_NAME}=b64.{cookie_token}; Path=/; HttpOnly; SameSite=Lax"
                    f"{secure_part}"
                )
                headers.append((b"set-cookie", cookie_value.encode("utf-8")))
                message["headers"] = headers
            await send(message)

        await self.app(scope, receive, send_wrapper)


def build_context_prompt(hits: List[Dict]) -> str:
    blocks = []
    for idx, hit in enumerate(hits, start=1):
        snippet = hit["text"].strip().replace("\n", " ")
        snippet = " ".join(snippet.split())
        snippet = textwrap.shorten(snippet, width=700, placeholder="…")
        doc_name = hit.get("doc_name", "Document")
        page = hit.get("page", "?")
        year = hit.get("year", "?")
        blocks.append(f"[{idx}] {doc_name} ({year}, page {page}): {snippet}")
    return "\n".join(blocks)


def format_references(hits: List[Dict]) -> str:
    references_lines = []
    for idx, hit in enumerate(hits, start=1):
        doc_name = hit.get("doc_name", "Document")
        page = hit.get("page", "?")
        year = hit.get("year", "?")
        score = hit.get("score", 0.0)
        snippet = textwrap.shorten(hit["text"].strip(), width=220, placeholder="…")
        references_lines.append(
            f"**[{idx}] {doc_name} ({year}) — page {page} (score {score:.3f})**\n{snippet}"
        )
    return "\n\n".join(references_lines)


def resolve_reasoning_selection(
    preset_label: Optional[str], custom_model: Optional[str]
) -> Tuple[str, float, Dict]:
    preset = REASONING_PRESETS.get(preset_label or "")
    model_id = (custom_model or "").strip()
    if not model_id and preset:
        model_id = preset.get("model", DEFAULT_CHAT_MODEL)
    if not model_id:
        model_id = DEFAULT_CHAT_MODEL
    temperature = preset.get("temperature", 0.1) if preset else 0.1
    extra_payload = preset.get("extra", {}).copy() if preset else {}
    return model_id, temperature, extra_payload


def tuples_to_messages(history: List[Tuple[str, str]]) -> List[Dict[str, str]]:
    messages: List[Dict[str, str]] = []
    for user_text, bot_text in history:
        messages.append({"role": "user", "content": user_text})
        messages.append({"role": "assistant", "content": bot_text})
    return messages


def condense_hits_for_trace(hits: List[Dict]) -> List[Dict]:
    condensed: List[Dict] = []
    for hit in hits:
        condensed.append(
            {
                "doc_id": hit.get("doc_id"),
                "doc_name": hit.get("doc_name", "Document"),
                "year": hit.get("year"),
                "page": hit.get("page"),
                "chunk_number": hit.get("chunk_number"),
                "score": round(float(hit.get("score", 0.0)), 6),
                "snippet": textwrap.shorten(
                    " ".join(str(hit.get("text", "")).split()),
                    width=240,
                    placeholder="…",
                ),
            }
        )
    return condensed


def chat_with_rag(
    user_message: str,
    history: Optional[List[Tuple[str, str]]],
    top_k: int,
    year_range: Any,
    embedding_label: Optional[str],
    preset_label: Optional[str],
    custom_model: Optional[str],
    chat_trace: Optional[List[Dict]],
) -> Tuple[List[Dict[str, str]], str, List[Tuple[str, str]], Any, List[Dict]]:
    history = history or []
    chat_trace = chat_trace or []
    message = user_message.strip()
    if not message:
        return (
            tuples_to_messages(history),
            "",
            history,
            gr.update(value=user_message),
            chat_trace,
        )

    strategy_key, strategy_name = resolve_embedding_selection(embedding_label)
    year_from, year_to = parse_year_range_input(year_range)
    model_id, temperature, extra_payload = resolve_reasoning_selection(
        preset_label, custom_model
    )
    turn_settings = {
        "top_k": int(top_k),
        "year_range": [int(year_from), int(year_to)],
        "embedding_strategy": strategy_name,
        "embedding_key": strategy_key,
        "reasoning_profile": preset_label or "",
        "custom_model_input": (custom_model or "").strip(),
        "llm_model": model_id,
        "temperature": temperature,
    }
    hits = search_index(
        message,
        top_k=top_k,
        strategy_key=strategy_key,
        year_from=year_from,
        year_to=year_to,
    )
    if not hits:
        reply = (
            "No matching context was found for the selected year range and embedding strategy. "
            "Try widening the year filter or switching embeddings."
        )
        updated_history = history + [(user_message, reply)]
        turn_record = {
            "timestamp": utc_now_iso_z(),
            "user_message": user_message,
            "assistant_reply": reply,
            "settings": turn_settings,
            "rag_hits": [],
        }
        updated_trace = chat_trace + [turn_record]
        return (
            tuples_to_messages(updated_history),
            "",
            updated_history,
            gr.update(value=""),
            updated_trace,
        )

    context_prompt = build_context_prompt(hits)
    system_message = {
        "role": "system",
        "content": (
            "You are a research assistant. Use ONLY the provided context blocks when "
            "answering. Cite sources inline using square brackets matching the context "
            "indices (e.g., [1], [2]). If the information is unavailable, state that "
            "you do not know."
        ),
    }

    messages = [system_message]
    for prev_user, prev_bot in history:
        messages.append({"role": "user", "content": prev_user})
        messages.append({"role": "assistant", "content": prev_bot})

    augmented_user = (
        f"Context:\n{context_prompt}\n\n"
        f"Current user message: {message}\n"
        "Respond with a concise answer followed by bullet points of key supporting facts. "
        f"Use citations [n] and preserve publication-year awareness (current retrieval embedding: {strategy_name})."
    )
    messages.append({"role": "user", "content": augmented_user})

    try:
        answer_text = CHAT_CLIENT.generate(
            messages,
            temperature=temperature,
            model_override=model_id,
            extra_payload=extra_payload,
        )
    except Exception as exc:
        answer_text = (
            "Failed to call the language model. "
            f"Verify the DEEPINFRA environment variable and model access. Details: {exc}"
        )

    updated_history = history + [(user_message, answer_text)]
    references_text = format_references(hits)
    turn_record = {
        "timestamp": utc_now_iso_z(),
        "user_message": user_message,
        "assistant_reply": answer_text,
        "settings": turn_settings,
        "rag_hits": condense_hits_for_trace(hits),
    }
    updated_trace = chat_trace + [turn_record]
    return (
        tuples_to_messages(updated_history),
        references_text,
        updated_history,
        gr.update(value=""),
        updated_trace,
    )


def reset_chat():
    return [], "", [], gr.update(value=""), []


# ---------------------------------------------------------------------------
# Gradio UI helpers
# ---------------------------------------------------------------------------
def registry_snapshot() -> Tuple[List[List], List[str]]:
    registry = load_registry_threadsafe()
    documents = registry.get("documents", {})
    rows: List[List] = []
    entries: List[str] = []
    sorted_docs = sorted(
        documents.items(),
        key=lambda item: item[1].get("ingested_at", ""),
        reverse=True,
    )
    for doc_id, doc in sorted_docs:
        year_value = doc.get("year")
        try:
            year_display = str(int(year_value))
        except (TypeError, ValueError):
            year_display = "—"
        rows.append(
            [
                "🗑️",
                doc_id,
                doc["title"],
                year_display,
                doc["pages"],
                doc.get("chunks", 0),
                ", ".join(map(str, doc.get("blank_pages", []))) or "—",
                doc.get("ingested_at", ""),
            ]
        )
        entries.append(doc_id)
    return rows, entries


def get_year_bounds() -> Tuple[int, int]:
    registry = load_registry_threadsafe()
    documents = registry.get("documents", {})
    years: List[int] = []
    for doc in documents.values():
        try:
            years.append(int(doc.get("year")))
        except (TypeError, ValueError):
            continue
    if years:
        return min(years), max(years)
    current_year = utc_year()
    return current_year - 20, current_year


def year_slider_updates(selected_range: Any = None) -> Any:
    min_year, max_year = get_year_bounds()
    if min_year > max_year:
        min_year, max_year = max_year, min_year

    selected_from, selected_to = parse_year_range_input(selected_range)
    selected_from = max(min_year, min(max_year, int(selected_from)))
    selected_to = max(min_year, min(max_year, int(selected_to)))

    if selected_from > selected_to:
        selected_from, selected_to = selected_to, selected_from

    return gr.update(
        minimum=min_year,
        maximum=max_year,
        value=(selected_from, selected_to),
    )


def registry_outputs(
    selected_year_range: Any = None,
):
    table, _ = registry_snapshot()
    year_range_update = year_slider_updates(selected_year_range)
    return (
        table,
        year_range_update,
    )


def handle_upload(files: List[str], selected_year_range: Any):
    if not files:
        table, year_range_update = registry_outputs(
            selected_year_range=selected_year_range
        )
        return (
            "Please upload at least one PDF.",
            table,
            year_range_update,
        )
    messages = []
    for file_path in files:
        try:
            result = ingest_pdf(file_path)
            messages.append(result["message"])
        except Exception as exc:
            messages.append(f"Failed to index {Path(file_path).name}: {exc}")
    table, year_range_update = registry_outputs(
        selected_year_range=selected_year_range
    )
    return "\n".join(messages), table, year_range_update


def handle_delete(doc_id: Optional[str], selected_year_range: Any):
    doc_id = (doc_id or "").strip()
    if not doc_id:
        table, year_range_update = registry_outputs(selected_year_range=selected_year_range)
        return "Select a row in the document table to delete.", table, year_range_update
    try:
        outcome = delete_document(doc_id)
    except Exception as exc:
        outcome = {
            "status": "error",
            "doc_id": doc_id,
            "message": f"Delete failed for {doc_id}: {exc}",
        }
    table, year_range_update = registry_outputs(selected_year_range=selected_year_range)
    return outcome["message"], table, year_range_update


def handle_refresh_registry(selected_year_range: Any):
    return registry_outputs(selected_year_range=selected_year_range)


def _event_row_col(evt: gr.SelectData) -> Tuple[Optional[int], Optional[int]]:
    row_index: Optional[int] = None
    col_index: Optional[int] = None
    event_index = getattr(evt, "index", None)
    if isinstance(event_index, tuple):
        if len(event_index) >= 1:
            row_index = event_index[0]
        if len(event_index) >= 2:
            col_index = event_index[1]
    elif isinstance(event_index, list):
        if len(event_index) >= 1:
            row_index = event_index[0]
        if len(event_index) >= 2:
            col_index = event_index[1]
    elif isinstance(event_index, int):
        row_index = event_index

    return row_index, col_index


def _row_from_table_data(table_data: Any, row_index: int) -> Optional[List[Any]]:
    if row_index < 0:
        return None

    # Gradio Dataframe can provide pandas-like or list-like structures.
    if hasattr(table_data, "iloc"):
        try:
            row = table_data.iloc[row_index]
            if hasattr(row, "tolist"):
                return list(row.tolist())
            return list(row)
        except Exception:
            return None

    if isinstance(table_data, np.ndarray):
        try:
            row = table_data[row_index]
            return row.tolist() if hasattr(row, "tolist") else list(row)
        except Exception:
            return None

    if isinstance(table_data, list):
        if row_index >= len(table_data):
            return None
        row = table_data[row_index]
        if isinstance(row, (list, tuple)):
            return list(row)
        return [row]

    return None


def handle_registry_row_select(table_data: Any, selected_year_range: Any, evt: gr.SelectData):
    try:
        row_index: Optional[int] = None
        col_index: Optional[int] = None
        row_index, col_index = _event_row_col(evt)
        selected_value = getattr(evt, "value", None)

        if row_index is None:
            return gr.update(), gr.update(), year_slider_updates(selected_year_range)
        if row_index < 0:
            return gr.update(), gr.update(), year_slider_updates(selected_year_range)

        row = _row_from_table_data(table_data, row_index)
        if not row or len(row) < 2:
            return gr.update(), gr.update(), year_slider_updates(selected_year_range)

        # First column acts like a delete action button.
        if col_index == 0 or selected_value == "🗑️":
            return handle_delete(doc_id=str(row[1]), selected_year_range=selected_year_range)

        return (
            "Click the 🗑️ icon in the first column to delete this entry.",
            gr.update(),
            year_slider_updates(selected_year_range),
        )
    except Exception as exc:
        return (
            f"Delete action failed: {exc}",
            gr.update(),
            year_slider_updates(selected_year_range),
        )


def normalize_chat_history(history: Optional[List[Tuple[str, str]]]) -> List[Dict[str, str]]:
    normalized: List[Dict[str, str]] = []
    for turn in history or []:
        if not isinstance(turn, (list, tuple)) or len(turn) < 2:
            continue
        normalized.append(
            {
                "user": str(turn[0]),
                "assistant": str(turn[1]),
            }
        )
    return normalized


def collect_rag_papers_from_trace(chat_trace: Optional[List[Dict]]) -> List[Dict]:
    papers: Dict[str, Dict] = {}
    for turn in chat_trace or []:
        if not isinstance(turn, dict):
            continue
        for hit in turn.get("rag_hits", []):
            if not isinstance(hit, dict):
                continue
            doc_id = str(hit.get("doc_id") or "").strip()
            doc_name = str(hit.get("doc_name") or "Document").strip()
            year_value = hit.get("year")
            key = doc_id or f"{doc_name}:{year_value}"
            if key in papers:
                continue
            papers[key] = {
                "doc_id": doc_id,
                "title": doc_name,
                "year": year_value,
            }
    return list(papers.values())


def feedback_snapshot() -> List[List]:
    store = load_feedback_registry_threadsafe()
    feedback_map = store.get("feedback", {})
    if not isinstance(feedback_map, dict):
        return []

    rows: List[List] = []
    sorted_entries = sorted(
        feedback_map.items(),
        key=lambda item: item[1].get("created_at", "") if isinstance(item[1], dict) else "",
        reverse=True,
    )
    for feedback_id, entry in sorted_entries:
        if not isinstance(entry, dict):
            continue
        settings = entry.get("settings", {})
        if not isinstance(settings, dict):
            settings = {}
        chat_history = entry.get("chat_history", [])
        if not isinstance(chat_history, list):
            chat_history = []
        rag_papers = entry.get("rag_papers", [])
        if not isinstance(rag_papers, list):
            rag_papers = []
        rag_titles = [
            str(paper.get("title") or "Document")
            for paper in rag_papers
            if isinstance(paper, dict)
        ]
        rag_display = ", ".join(rag_titles) if rag_titles else "—"
        year_range = settings.get("year_range")
        if isinstance(year_range, (list, tuple)) and len(year_range) >= 2:
            year_display = f"{year_range[0]}-{year_range[1]}"
        else:
            year_display = "—"
        status_value = str(entry.get("status") or "open")
        resolve_icon = "✅" if status_value != "resolved" else "☑️"
        rows.append(
            [
                "🗑️",
                resolve_icon,
                str(feedback_id),
                status_value,
                str(entry.get("created_at") or ""),
                str(entry.get("resolved_at") or "—"),
                textwrap.shorten(
                    str(entry.get("feedback_text") or ""),
                    width=140,
                    placeholder="…",
                ),
                len(chat_history),
                year_display,
                str(settings.get("embedding_strategy") or ""),
                str(settings.get("llm_model") or ""),
                textwrap.shorten(rag_display, width=120, placeholder="…"),
            ]
        )
    return rows


def feedback_entry_json(entry: Optional[Dict]) -> str:
    if not isinstance(entry, dict):
        return ""
    return json.dumps(entry, indent=2, ensure_ascii=True)


def open_feedback_dialog():
    return gr.update(visible=True), gr.update(value="")


def close_feedback_dialog():
    return gr.update(visible=False)


def handle_submit_feedback(
    feedback_text: str,
    history: Optional[List[Tuple[str, str]]],
    chat_trace: Optional[List[Dict]],
    top_k: int,
    year_range: Any,
    embedding_label: Optional[str],
    preset_label: Optional[str],
    custom_model: Optional[str],
):
    message = (feedback_text or "").strip()
    if not message:
        feedback_rows = feedback_snapshot()
        warning = "Please enter feedback before submitting."
        return (
            warning,
            gr.update(visible=True),
            gr.update(),
            feedback_rows,
            gr.update(),
            warning,
        )

    strategy_key, strategy_name = resolve_embedding_selection(embedding_label)
    model_id, temperature, _ = resolve_reasoning_selection(preset_label, custom_model)
    year_from, year_to = parse_year_range_input(year_range)
    normalized_history = normalize_chat_history(history)
    rag_papers = collect_rag_papers_from_trace(chat_trace)

    llm_models_used = sorted(
        {
            str(turn.get("settings", {}).get("llm_model") or "").strip()
            for turn in (chat_trace or [])
            if isinstance(turn, dict)
        }
        | {model_id}
    )
    llm_models_used = [model for model in llm_models_used if model]

    feedback_payload = {
        "feedback_id": "",
        "status": "open",
        "feedback_text": message,
        "created_at": utc_now_iso_z(),
        "resolved_at": "",
        "chat_history": normalized_history,
        "chat_trace": chat_trace or [],
        "settings": {
            "top_k": int(top_k),
            "year_range": [int(year_from), int(year_to)],
            "embedding_strategy": strategy_name,
            "embedding_key": strategy_key,
            "reasoning_profile": preset_label or "",
            "custom_model_input": (custom_model or "").strip(),
            "llm_model": model_id,
            "llm_models_used_in_chat": llm_models_used,
            "temperature": temperature,
        },
        "rag_papers": rag_papers,
    }

    saved_entry = create_feedback_entry(feedback_payload)
    feedback_rows = feedback_snapshot()
    success = f"Feedback saved as {saved_entry['feedback_id']}."
    return (
        success,
        gr.update(visible=False),
        gr.update(value=""),
        feedback_rows,
        gr.update(value=feedback_entry_json(saved_entry)),
        success,
    )


def handle_refresh_feedback():
    return feedback_snapshot()


def handle_feedback_row_select(table_data: Any, evt: gr.SelectData):
    try:
        row_index, col_index = _event_row_col(evt)
        selected_value = getattr(evt, "value", None)

        if row_index is None or row_index < 0:
            return "Select a feedback row.", gr.update(), gr.update()

        row = _row_from_table_data(table_data, row_index)
        if not row or len(row) < 3:
            return "Could not read selected feedback row.", gr.update(), gr.update()

        feedback_id = str(row[2]).strip()
        if not feedback_id:
            return "Feedback ID is missing in selected row.", gr.update(), gr.update()

        if col_index == 0 or selected_value == "🗑️":
            result = delete_feedback_entry(feedback_id)
            return result["message"], feedback_snapshot(), gr.update(value="")

        if col_index == 1 or selected_value in {"✅", "☑️"}:
            result = mark_feedback_resolved(feedback_id)
            entry = get_feedback_entry(feedback_id)
            detail_json = feedback_entry_json(entry)
            return result["message"], feedback_snapshot(), gr.update(value=detail_json)

        entry = get_feedback_entry(feedback_id)
        if not entry:
            return f"Feedback {feedback_id} not found.", feedback_snapshot(), gr.update(value="")
        return (
            f"Selected feedback {feedback_id}.",
            gr.update(),
            gr.update(value=feedback_entry_json(entry)),
        )
    except Exception as exc:
        return f"Feedback action failed: {exc}", gr.update(), gr.update()


def build_interface() -> gr.Blocks:
    ensure_strategy_indexes_initialized()
    initial_table, _ = registry_snapshot()
    initial_feedback_table = feedback_snapshot()
    initial_year_min, initial_year_max = get_year_bounds()

    with gr.Blocks(
        title="Document RAG",
        css="""
        .feedback-dialog-panel {
            border: 1px solid #d1d5db;
            border-radius: 12px;
            padding: 12px;
            margin-top: 10px;
            background: #f9fafb;
        }
        """,
    ) as demo:
        gr.Markdown(
            """
            # Document RAG

            Use the tabs:
            - **Upload PDFs**: index/delete RAG documents.
            - **Chat**: query the indexed documents with year-range and embedding controls.
            - **Feedback**: review, resolve, and delete feedback submissions.
            - **Config**: environment variable reference.
            """
        )

        with gr.Tab("Upload PDFs"):
            with gr.Row():
                file_input = gr.File(
                    label="PDF files",
                    file_types=[".pdf"],
                    file_count="multiple",
                    type="filepath",
                )
            upload_button = gr.Button("Process and Index")
            upload_status = gr.Textbox(
                label="Status", lines=4, interactive=False, show_label=True
            )
            registry_table = gr.Dataframe(
                headers=[
                    "Delete",
                    "Doc ID",
                    "Title",
                    "Year",
                    "Pages",
                    "Chunks",
                    "OCR pages",
                    "Indexed at",
                ],
                datatype=["str", "str", "str", "str", "number", "number", "str", "str"],
                interactive=False,
                value=initial_table,
            )
            with gr.Row():
                refresh_button = gr.Button("Refresh List")

        with gr.Tab("Chat"):
            reasoning_options = list(REASONING_PRESETS.keys())
            default_preset = reasoning_options[0] if reasoning_options else None

            with gr.Row():
                chatbot = gr.Chatbot(
                    label="Conversation",
                    height=420,
                    show_label=True,
                    sanitize_html=False,
                    type="messages",
                )
                references_md = gr.Markdown(label="References (last reply)")

            with gr.Row():
                chat_input = gr.Textbox(
                    label="Message",
                    placeholder="Ask about the indexed PDFs…",
                    lines=2,
                )
                send_button = gr.Button("Send", variant="primary")
                clear_button = gr.Button("Clear", variant="secondary")
                feedback_button = gr.Button("feedback", variant="secondary")

            with gr.Row():
                top_k_slider = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=10,
                    step=1,
                    label="Top-k documents",
                )
                if RangeSlider is not None:
                    year_range_slider = RangeSlider(
                        minimum=initial_year_min,
                        maximum=initial_year_max,
                        value=(initial_year_min, initial_year_max),
                        step=1,
                        label="Year range",
                        interactive=True,
                    )
                else:
                    # Fallback to built-in slider when the range-slider plugin is unavailable.
                    year_range_slider = gr.Slider(
                        minimum=initial_year_min,
                        maximum=initial_year_max,
                        value=(initial_year_min, initial_year_max),
                        step=1,
                        label="Year range",
                        interactive=True,
                    )

            with gr.Row():
                embedding_dropdown = gr.Dropdown(
                    label="Retrieval embedding strategy",
                    choices=EMBEDDING_OPTIONS,
                    value=DEFAULT_EMBEDDING_LABEL,
                    allow_custom_value=False,
                )
                reasoning_dropdown = gr.Dropdown(
                    label="Reasoning profile",
                    choices=reasoning_options,
                    value=default_preset,
                    allow_custom_value=False,
                )
                custom_model_box = gr.Textbox(
                    label="Custom model ID (optional)",
                    placeholder=f"Overrides preset (default: {DEFAULT_CHAT_MODEL})",
                )

            chat_feedback_status = gr.Textbox(
                label="Feedback status",
                lines=1,
                interactive=False,
                show_label=True,
                value="",
            )

            with gr.Group(visible=False, elem_classes=["feedback-dialog-panel"]) as feedback_dialog:
                gr.Markdown("### Submit feedback")
                feedback_input = gr.Textbox(
                    label="What should be fixed?",
                    lines=4,
                    placeholder="Describe what happened and what should change.",
                )
                with gr.Row():
                    submit_feedback_button = gr.Button("Submit feedback", variant="primary")
                    cancel_feedback_button = gr.Button("Cancel", variant="secondary")

            history_state = gr.State([])
            chat_trace_state = gr.State([])

        with gr.Tab("Feedback"):
            feedback_tab_status = gr.Textbox(
                label="Status",
                lines=2,
                interactive=False,
                show_label=True,
                value="",
            )
            feedback_table = gr.Dataframe(
                headers=[
                    "Delete",
                    "Resolve",
                    "Feedback ID",
                    "Status",
                    "Created at",
                    "Resolved at",
                    "Feedback",
                    "Turns",
                    "Year range",
                    "Embedding",
                    "LLM",
                    "RAG papers",
                ],
                datatype=[
                    "str",
                    "str",
                    "str",
                    "str",
                    "str",
                    "str",
                    "str",
                    "number",
                    "str",
                    "str",
                    "str",
                    "str",
                ],
                interactive=False,
                value=initial_feedback_table,
            )
            with gr.Row():
                feedback_refresh_button = gr.Button("Refresh Feedback")
            gr.Markdown(
                "Click `✅` in the Resolve column to mark an item as resolved, or `🗑️` to delete it."
            )
            feedback_detail_box = gr.Textbox(
                label="Selected feedback payload (JSON)",
                lines=18,
                interactive=False,
            )

        with gr.Tab("Config"):
            gr.Markdown(
                """
                ## Runtime Configuration

                - `APP_URL_TOKEN` (optional): URL token for app access (used as `?token=...` by default).
                - `APP_ENABLE_TOKEN_GATE_ON_SPACES` (optional): defaults to `0`; set to `1` only if you explicitly want query-token gating on Spaces.
                - `APP_ALLOW_LOOPBACK_BYPASS` (optional): defaults to `1`; keeps localhost startup/health checks working.
                - `DEEPINFRA`: required for OCR and answer generation.
                - `DEEPINFRA_OCR_MODEL` (optional): defaults to `deepseek-ai/deepseek-ocr`.
                - `DEEPINFRA_RAG_MODEL` (optional): defaults to `meta-llama/Meta-Llama-3.1-8B-Instruct`.
                - `DEEPINFRA_FAST_MODEL`, `DEEPINFRA_BALANCED_MODEL`, `DEEPINFRA_REASONING_MODEL` (optional): override preset models.
                - `EMBED_MODEL_FAST`, `EMBED_MODEL_QUALITY`, `EMBED_MODEL_QA` (optional): override retrieval embedding models.
                """
            )
            if RangeSlider is None:
                gr.Markdown(
                    "Year range uses fallback slider because `gradio_rangeslider` is not installed."
                )

        send_inputs = [
            chat_input,
            history_state,
            top_k_slider,
            year_range_slider,
            embedding_dropdown,
            reasoning_dropdown,
            custom_model_box,
            chat_trace_state,
        ]
        send_outputs = [chatbot, references_md, history_state, chat_input, chat_trace_state]
        registry_outputs_all = [upload_status, registry_table, year_range_slider]
        registry_refresh_outputs = [registry_table, year_range_slider]
        feedback_submit_outputs = [
            chat_feedback_status,
            feedback_dialog,
            feedback_input,
            feedback_table,
            feedback_detail_box,
            feedback_tab_status,
        ]
        feedback_row_outputs = [
            feedback_tab_status,
            feedback_table,
            feedback_detail_box,
        ]

        upload_button.click(
            fn=handle_upload,
            inputs=[file_input, year_range_slider],
            outputs=registry_outputs_all,
        )
        refresh_button.click(
            fn=handle_refresh_registry,
            inputs=[year_range_slider],
            outputs=registry_refresh_outputs,
        )
        registry_table.select(
            fn=handle_registry_row_select,
            inputs=[registry_table, year_range_slider],
            outputs=registry_outputs_all,
        )

        send_button.click(
            fn=chat_with_rag,
            inputs=send_inputs,
            outputs=send_outputs,
        )
        chat_input.submit(
            fn=chat_with_rag,
            inputs=send_inputs,
            outputs=send_outputs,
        )
        clear_button.click(
            fn=reset_chat,
            outputs=send_outputs,
            queue=False,
        )

        feedback_button.click(
            fn=open_feedback_dialog,
            outputs=[feedback_dialog, chat_feedback_status],
            queue=False,
        )
        cancel_feedback_button.click(
            fn=close_feedback_dialog,
            outputs=[feedback_dialog],
            queue=False,
        )
        submit_feedback_button.click(
            fn=handle_submit_feedback,
            inputs=[
                feedback_input,
                history_state,
                chat_trace_state,
                top_k_slider,
                year_range_slider,
                embedding_dropdown,
                reasoning_dropdown,
                custom_model_box,
            ],
            outputs=feedback_submit_outputs,
        )

        feedback_refresh_button.click(
            fn=handle_refresh_feedback,
            outputs=[feedback_table],
        )
        feedback_table.select(
            fn=handle_feedback_row_select,
            inputs=[feedback_table],
            outputs=feedback_row_outputs,
        )

        demo.load(
            fn=handle_refresh_registry,
            inputs=[year_range_slider],
            outputs=registry_refresh_outputs,
            queue=False,
        )
        demo.load(
            fn=handle_refresh_feedback,
            outputs=[feedback_table],
            queue=False,
        )

    return demo


if __name__ == "__main__":
    app = build_interface()
    launch_kwargs: Dict[str, Any] = {"ssr_mode": False}
    if URL_ACCESS_TOKEN and not TOKEN_GATE_ENABLED:
        print(
            "APP_URL_TOKEN is set but token gating is disabled on Spaces by default. "
            "Set APP_ENABLE_TOKEN_GATE_ON_SPACES=1 to enable it."
        )
    if TOKEN_GATE_ENABLED:
        launch_kwargs["app_kwargs"] = {
            "middleware": [Middleware(QueryTokenGateMiddleware)]
        }
    app.launch(**launch_kwargs)