# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""RAG API: KB CRUD, document upload (KB + per-thread), ingestion SSE, search."""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
import queue as queue_module
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Literal, Optional
from urllib.parse import quote
from uuid import uuid4

import jwt
from fastapi import (
    APIRouter,
    Depends,
    Header,
    HTTPException,
    Query,
    Request,
    UploadFile,
)
from fastapi.responses import FileResponse, Response, StreamingResponse
from pydantic import BaseModel, Field

from auth.authentication import get_current_subject, get_current_subject_sse
from auth.storage import get_jwt_secret


async def _sse_auth(
    token: str | None = Query(None),
    authorization: str | None = Header(None),
) -> str:
    return await get_current_subject_sse(token, authorization)


from core.rag import embeddings, ingestion, reranker, retrieval, vector_store
from core.rag.authorization import document_for_subject_or_404
from core.rag.locators import backfill_document_locators
from core.rag.vector_store import kb_scope, thread_scope
from loggers import get_logger
from storage.studio_db import (
    get_connection,
    list_chat_settings,
    upsert_chat_settings_merge,
)
from utils.paths.storage_roots import ensure_dir, rag_uploads_root, resolve_under_root
from utils.rag.config import (
    RAG_MAX_UPLOAD_MB,
    RAG_RERANK_CANDIDATE_K,
    RAG_UPLOAD_EXTS,
)

router = APIRouter()
logger = get_logger(__name__)


# --- Pydantic schemas ---

ChunkingStrategy = Literal["standard", "late"]
KBMode = Literal["text", "multimodal"]


class CreateKBRequest(BaseModel):
    name: str = Field(min_length = 1, max_length = 200)
    description: str | None = None
    embedding_model: str | None = None
    chunking_strategy: ChunkingStrategy = "standard"
    mode: KBMode = "text"


class KBResponse(BaseModel):
    id: str
    name: str
    description: str | None
    embedding_model: str
    chunking_strategy: ChunkingStrategy
    mode: KBMode
    created_at: int


class KBListResponse(BaseModel):
    knowledge_bases: list[KBResponse]


class DocumentResponse(BaseModel):
    id: str
    kb_id: str | None
    thread_id: str | None
    filename: str
    content_type: str | None
    status: str
    num_chunks: int
    byte_size: int
    error: str | None
    created_at: int


class DocumentListResponse(BaseModel):
    documents: list[DocumentResponse]


class ThreadIndexSummary(BaseModel):
    thread_id: str
    title: str | None
    num_documents: int
    num_chunks: int


class ThreadIndexListResponse(BaseModel):
    threads: list[ThreadIndexSummary]


class UploadResponse(BaseModel):
    document_id: str
    job_id: str
    filename: str
    # True when an identical file (same content hash) was already indexed
    # in this scope, so no new ingestion job was started. job_id is "".
    already_indexed: bool = False


class SearchRequest(BaseModel):
    query: str = Field(min_length = 1, max_length = 4000)
    kb_id: str | None = None
    thread_id: str | None = None
    top_k: int = Field(default = 10, ge = 1, le = 100)
    mode: Literal["bm25", "dense", "hybrid"] = "hybrid"
    document_ids: list[str] | None = None
    enable_rerank: bool = False
    reranker_model: str | None = None
    min_score: float = Field(default = 0.0, ge = 0.0, le = 1.0)


class SearchHit(BaseModel):
    chunk_id: str
    document_id: str
    chunk_index: int
    text: str
    score: float
    page_number: int | None = None
    filename: str | None = None
    kind: str = "text"
    image_url: str | None = None
    source_page_index: int | None = None
    page_char_start: int | None = None
    page_char_end: int | None = None
    line_start: int | None = None
    line_end: int | None = None


class SearchResponse(BaseModel):
    hits: list[SearchHit]


class PrefetchRequest(BaseModel):
    """Prefetch RAG for an external-provider turn: studio decomposes the
    question (via the pre-cached helper) and retrieves, so the frontend can
    inject the chunks before calling the provider."""

    query: str = Field(min_length = 1, max_length = 4000)
    kb_id: str | None = None
    thread_id: str | None = None
    top_k: int = Field(default = 10, ge = 1, le = 100)
    mode: Literal["bm25", "dense", "hybrid"] = "hybrid"
    enable_rerank: bool = False
    reranker_model: str | None = None
    min_score: float = Field(default = 0.0, ge = 0.0, le = 1.0)


class PrefetchResponse(BaseModel):
    queries: list[str]
    hits: list[SearchHit]


# --- Helpers ---


def _sanitize_filename(filename: str) -> str:
    name = Path(filename).name.strip().replace("\x00", "")
    return name or "document"


def _now_ms() -> int:
    return int(time.time())


from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder  # noqa: E402


def _row_to_kb(row: Any) -> KBResponse:
    keys = row.keys() if hasattr(row, "keys") else ()
    chunking_strategy = (
        row["chunking_strategy"] if "chunking_strategy" in keys else "standard"
    )
    mode = row["mode"] if "mode" in keys else "text"
    return KBResponse(
        id = row["id"],
        name = row["name"],
        description = row["description"],
        embedding_model = row["embedding_model"],
        chunking_strategy = chunking_strategy,
        mode = mode,
        created_at = row["created_at"],
    )


def _validate_mode_combo(mode: KBMode, chunking_strategy: ChunkingStrategy) -> None:
    """Reject (multimodal, late) — no embedder supports both at once."""
    if mode == "multimodal" and chunking_strategy == "late":
        raise HTTPException(
            status_code = 400,
            detail = (
                "Late chunking is not supported in multimodal mode — "
                "the multimodal embedder does not expose per-token "
                "embeddings. Pick 'standard' chunking or 'text' mode."
            ),
        )


def _row_to_document(row: Any) -> DocumentResponse:
    return DocumentResponse(
        id = row["id"],
        kb_id = row["kb_id"],
        thread_id = row["thread_id"],
        filename = row["filename"],
        content_type = row["content_type"],
        status = row["status"],
        num_chunks = row["num_chunks"],
        byte_size = row["byte_size"],
        error = row["error"],
        created_at = row["created_at"],
    )


def _kb_or_404(kb_id: str) -> Any:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT * FROM rag_knowledge_bases WHERE id = ?",
            (kb_id,),
        ).fetchone()
    if not row:
        raise HTTPException(status_code = 404, detail = "Knowledge base not found")
    return row


def _thread_or_404(thread_id: str) -> None:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT id FROM chat_threads WHERE id = ?",
            (thread_id,),
        ).fetchone()
    if not row:
        raise HTTPException(status_code = 404, detail = "Thread not found")


def _document_or_404(document_id: str) -> Any:
    with get_connection() as conn:
        row = conn.execute(
            "SELECT * FROM rag_documents WHERE id = ?",
            (document_id,),
        ).fetchone()
    if not row:
        raise HTTPException(status_code = 404, detail = "Document not found")
    return row


async def _save_upload(file: UploadFile) -> tuple[Path, str, int, str]:
    import anyio

    filename = _sanitize_filename(file.filename or "document")
    ext = Path(filename).suffix.lower()
    if ext not in RAG_UPLOAD_EXTS:
        allowed = ", ".join(sorted(RAG_UPLOAD_EXTS))
        raise HTTPException(
            status_code = 400,
            detail = f"Unsupported file type: {ext}. Allowed: {allowed}",
        )
    upload_dir = ensure_dir(rag_uploads_root())
    stored_name = f"{uuid4().hex}_{Path(filename).stem}{ext}"
    stored_path = upload_dir / stored_name
    max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024
    written = 0
    # Hash the bytes as they stream so we can dedup identical re-uploads
    # within a scope without re-reading the file.
    hasher = hashlib.sha256()
    # Route writes through anyio worker thread so the event loop stays free.
    # Outer try/except cleans up partial files after async-with closes the fd
    # (Windows refuses unlink on an open fd).
    try:
        async with await anyio.open_file(stored_path, "wb") as f:
            while True:
                chunk = await file.read(1024 * 1024)
                if not chunk:
                    break
                written += len(chunk)
                if written > max_bytes:
                    raise HTTPException(
                        status_code = 413,
                        detail = f"File exceeds {RAG_MAX_UPLOAD_MB} MB limit",
                    )
                hasher.update(chunk)
                await f.write(chunk)
    except HTTPException:
        stored_path.unlink(missing_ok = True)
        raise
    if written == 0:
        stored_path.unlink(missing_ok = True)
        raise HTTPException(status_code = 400, detail = "Empty upload payload")
    return stored_path, filename, written, hasher.hexdigest()


def _start_ingestion(
    *,
    filename: str,
    stored_path: Path,
    byte_size: int,
    content_type: str | None,
    kb_id: str | None,
    thread_id: str | None,
    embedding_model: str,
    chunking_strategy: str = "standard",
    mode: str = "text",
    caption_images: bool = True,
    content_hash: str | None = None,
) -> UploadResponse:
    document_id = str(uuid4())
    with get_connection() as conn:
        # Dedup: if an identical file (same content hash) is already
        # indexed in this scope, skip re-ingestion. Only a 'completed'
        # row counts — a failed/in-flight prior attempt should be allowed
        # to retry. Scope is the same kb_id or thread_id the upload
        # targets (a file shared across two KBs is indexed in each).
        if content_hash:
            if kb_id is not None:
                existing = conn.execute(
                    "SELECT id, filename FROM rag_documents "
                    "WHERE kb_id = ? AND content_hash = ? AND status = 'completed' "
                    "LIMIT 1",
                    (kb_id, content_hash),
                ).fetchone()
            else:
                existing = conn.execute(
                    "SELECT id, filename FROM rag_documents "
                    "WHERE thread_id = ? AND content_hash = ? AND status = 'completed' "
                    "LIMIT 1",
                    (thread_id, content_hash),
                ).fetchone()
            if existing is not None:
                # Drop the redundant upload we just wrote to disk; the
                # already-indexed copy stays the source of truth.
                _unlink_if_under_uploads(stored_path)
                return UploadResponse(
                    document_id = existing["id"],
                    job_id = "",
                    filename = existing["filename"],
                    already_indexed = True,
                )
        conn.execute(
            """
            INSERT INTO rag_documents
            (id, kb_id, thread_id, filename, content_type, stored_path,
             status, num_chunks, byte_size, content_hash, created_at)
            VALUES (?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?, ?)
            """,
            (
                document_id,
                kb_id,
                thread_id,
                filename,
                content_type,
                str(stored_path),
                byte_size,
                content_hash,
                _now_ms(),
            ),
        )
        conn.commit()
    job_id = ingestion.enqueue_ingestion(
        document_id = document_id,
        stored_path = stored_path,
        kb_id = kb_id,
        thread_id = thread_id,
        embedding_model = embedding_model,
        chunking_strategy = chunking_strategy,
        mode = mode,
        enable_captions = caption_images,
    )
    return UploadResponse(document_id = document_id, job_id = job_id, filename = filename)


def _unlink_if_under_uploads(path: Path) -> None:
    try:
        real = Path(os.path.realpath(path))
        root = Path(os.path.realpath(rag_uploads_root()))
        real.relative_to(root)
    except (OSError, ValueError):
        return
    real.unlink(missing_ok = True)


# --- Knowledge bases ---


@router.post("/knowledge-bases", response_model = KBResponse)
def create_knowledge_base(
    payload: CreateKBRequest,
    current_subject: str = Depends(get_current_subject),
) -> KBResponse:
    from utils.rag.config import resolve_embedder

    _validate_mode_combo(payload.mode, payload.chunking_strategy)

    kb_id = str(uuid4())
    # No override: resolve from (mode, strategy) matrix.
    embedding_model = payload.embedding_model or resolve_embedder(
        payload.mode, payload.chunking_strategy
    )
    created_at = _now_ms()
    with get_connection() as conn:
        try:
            conn.execute(
                """
                INSERT INTO rag_knowledge_bases
                (id, name, description, owner_user_id, embedding_model,
                 chunking_strategy, mode, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    kb_id,
                    payload.name,
                    payload.description,
                    current_subject,
                    embedding_model,
                    payload.chunking_strategy,
                    payload.mode,
                    created_at,
                ),
            )
            conn.commit()
        except Exception as exc:
            raise HTTPException(
                status_code = 409,
                detail = f"Could not create KB: {exc}",
            ) from exc
    return KBResponse(
        id = kb_id,
        name = payload.name,
        description = payload.description,
        embedding_model = embedding_model,
        chunking_strategy = payload.chunking_strategy,
        mode = payload.mode,
        created_at = created_at,
    )


@router.get("/knowledge-bases", response_model = KBListResponse)
def list_knowledge_bases(
    current_subject: str = Depends(get_current_subject),
) -> KBListResponse:
    with get_connection() as conn:
        rows = conn.execute(
            "SELECT * FROM rag_knowledge_bases ORDER BY created_at DESC"
        ).fetchall()
    return KBListResponse(knowledge_bases = [_row_to_kb(r) for r in rows])


class RagDefaults(BaseModel):
    chunking_strategy: ChunkingStrategy = "standard"
    mode: KBMode = "text"
    embedding_model: str | None = None


class UpdateRagDefaultsRequest(BaseModel):
    """Patch shape — only fields present overwrite stored values."""

    chunking_strategy: ChunkingStrategy | None = None
    mode: KBMode | None = None
    embedding_model: str | None = None


_DEFAULTS_KEY = "rag.defaults"


def _load_rag_defaults() -> RagDefaults:
    settings = list_chat_settings()
    raw = settings.get(_DEFAULTS_KEY) or {}
    if not isinstance(raw, dict):
        raw = {}
    return RagDefaults(
        chunking_strategy = raw.get("chunking_strategy") or "standard",
        mode = raw.get("mode") or "text",
        embedding_model = raw.get("embedding_model"),
    )


@router.get("/defaults", response_model = RagDefaults)
def get_rag_defaults(
    current_subject: str = Depends(get_current_subject),
) -> RagDefaults:
    return _load_rag_defaults()


@router.post("/warmup")
def warmup_rag_embedder(
    current_subject: str = Depends(get_current_subject),
) -> dict:
    """Preload the configured default embedder so the first retrieval is warm.

    Called from the frontend when the user enables the RAG pill — moves
    the cold-load latency out of the first chat-completion path, where a
    multi-second load can race the llama-server prefill timeout.
    """
    from utils.rag.config import resolve_embedder

    defaults = _load_rag_defaults()
    model_name = defaults.embedding_model or resolve_embedder(
        defaults.mode,
        defaults.chunking_strategy,
    )
    try:
        embeddings.get_embedder(model_name)
    except Exception as exc:  # noqa: BLE001
        # Log the detailed exception server-side; return a generic message
        # so internal paths / stack info aren't exposed to the client.
        logger.warning("RAG warmup failed for %s: %s", model_name, exc)
        return {"ok": False, "model": model_name, "error": "Failed to load embedder"}
    return {"ok": True, "model": model_name}


@router.post("/reranker/precache")
def precache_rag_reranker(
    current_subject: str = Depends(get_current_subject),
) -> dict:
    """Download the reranker weights (~1.1 GB) into the HF cache.

    Called from the frontend the moment the user flips the "Use
    reranker" switch ON so the cost lands on the explicit toggle
    instead of the first chat turn — where a multi-minute download
    looks like a hung tool call.
    """
    from core.rag.reranker import precache_reranker
    from utils.rag.config import RAG_RERANKER_MODEL

    try:
        precache_reranker()
    except Exception as exc:  # noqa: BLE001
        # Detailed exception logged server-side; client gets a generic
        # message so internal paths / stack info aren't exposed.
        logger.warning(
            "RAG reranker precache failed",
            model = RAG_RERANKER_MODEL,
            error = str(exc),
        )
        return {
            "ok": False,
            "model": RAG_RERANKER_MODEL,
            "error": "Failed to download reranker",
        }
    return {"ok": True, "model": RAG_RERANKER_MODEL}


@router.put("/defaults", response_model = RagDefaults)
def set_rag_defaults(
    payload: UpdateRagDefaultsRequest,
    current_subject: str = Depends(get_current_subject),
) -> RagDefaults:
    current = _load_rag_defaults()
    new_strategy = payload.chunking_strategy or current.chunking_strategy
    new_mode = payload.mode or current.mode
    # PATCH-style: empty string clears, null/missing keeps current.
    if payload.embedding_model is None:
        new_embedder = current.embedding_model
    elif payload.embedding_model.strip() == "":
        new_embedder = None
    else:
        new_embedder = payload.embedding_model.strip()
    _validate_mode_combo(new_mode, new_strategy)

    upsert_chat_settings_merge(
        {
            _DEFAULTS_KEY: {
                "chunking_strategy": new_strategy,
                "mode": new_mode,
                "embedding_model": new_embedder,
            }
        }
    )
    return RagDefaults(
        chunking_strategy = new_strategy,
        mode = new_mode,
        embedding_model = new_embedder,
    )


class ThreadRagSettings(BaseModel):
    chunking_strategy: ChunkingStrategy = "standard"
    mode: KBMode = "text"
    embedding_model: str | None = None


class UpdateThreadRagSettingsRequest(BaseModel):
    chunking_strategy: ChunkingStrategy | None = None
    mode: KBMode | None = None
    embedding_model: str | None = None
    # Only consulted by reingest (not persisted as a thread setting); omit or
    # None keeps captioning on.
    caption_images: bool | None = None


def _thread_settings_key(thread_id: str) -> str:
    return f"thread:{thread_id}:rag"


def _load_thread_settings(thread_id: str) -> ThreadRagSettings:
    """Per-thread RAG settings (chat_settings['thread:<id>:rag']) with defaults fallback."""
    settings = list_chat_settings()
    raw = settings.get(_thread_settings_key(thread_id)) or {}
    if not isinstance(raw, dict):
        raw = {}
    fallback = _load_rag_defaults()
    return ThreadRagSettings(
        chunking_strategy = (raw.get("chunking_strategy") or fallback.chunking_strategy),
        mode = raw.get("mode") or fallback.mode,
        embedding_model = raw.get("embedding_model") or fallback.embedding_model,
    )


@router.get(
    "/threads/{thread_id}/settings",
    response_model = ThreadRagSettings,
)
def get_thread_rag_settings(
    thread_id: str,
    current_subject: str = Depends(get_current_subject),
) -> ThreadRagSettings:
    return _load_thread_settings(thread_id)


@router.put(
    "/threads/{thread_id}/settings",
    response_model = ThreadRagSettings,
)
def set_thread_rag_settings(
    thread_id: str,
    payload: UpdateThreadRagSettingsRequest,
    current_subject: str = Depends(get_current_subject),
) -> ThreadRagSettings:
    current = _load_thread_settings(thread_id)
    new_strategy = payload.chunking_strategy or current.chunking_strategy
    new_mode = payload.mode or current.mode
    if payload.embedding_model is None:
        new_embedder = current.embedding_model
    elif payload.embedding_model.strip() == "":
        new_embedder = None
    else:
        new_embedder = payload.embedding_model.strip()
    _validate_mode_combo(new_mode, new_strategy)

    upsert_chat_settings_merge(
        {
            _thread_settings_key(thread_id): {
                "chunking_strategy": new_strategy,
                "mode": new_mode,
                "embedding_model": new_embedder,
            }
        }
    )
    return ThreadRagSettings(
        chunking_strategy = new_strategy,
        mode = new_mode,
        embedding_model = new_embedder,
    )


class ReingestKBRequest(BaseModel):
    """All fields optional — omitting one keeps the KB's current value."""

    chunking_strategy: ChunkingStrategy | None = None
    mode: KBMode | None = None
    embedding_model: str | None = None
    # Not persisted on the KB; omit or None keeps captioning on for the rebuild.
    caption_images: bool | None = None


class ReingestResponse(BaseModel):
    job_ids: list[str]
    document_ids: list[str]


def _reingest_scope(
    *,
    kb_id: str | None,
    thread_id: str | None,
    chunking_strategy: str,
    mode: str,
    embedding_model: str,
    caption_images: bool = True,
) -> ReingestResponse:
    """Wipe scope artifacts and re-enqueue every document; metadata untouched."""
    scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id)  # type: ignore[arg-type]
    with get_connection() as conn:
        if kb_id:
            rows = conn.execute(
                "SELECT id, stored_path FROM rag_documents WHERE kb_id = ?",
                (kb_id,),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT id, stored_path FROM rag_documents WHERE thread_id = ?",
                (thread_id,),
            ).fetchall()
        # Drop rag_documents (chunks cascade); files on disk are reused below.
        doc_ids = [r["id"] for r in rows]
        if doc_ids:
            placeholders = ",".join("?" for _ in doc_ids)
            conn.execute(
                f"DELETE FROM rag_documents WHERE id IN ({placeholders})",
                doc_ids,
            )
            conn.commit()
    ingestion.delete_scope_artifacts(scope)

    job_ids: list[str] = []
    new_doc_ids: list[str] = []
    for row in rows:
        stored_path = Path(row["stored_path"])
        if not stored_path.is_file():
            continue
        filename = stored_path.name
        # Strip the upload-time UUID prefix; keep the original filename.
        if "_" in filename:
            _uuid_prefix, _, original = filename.partition("_")
            if original:
                filename = original
        upload = _start_ingestion(
            filename = filename,
            stored_path = stored_path,
            byte_size = stored_path.stat().st_size,
            content_type = None,
            kb_id = kb_id,
            thread_id = thread_id,
            embedding_model = embedding_model,
            chunking_strategy = chunking_strategy,
            mode = mode,
            caption_images = caption_images,
        )
        job_ids.append(upload.job_id)
        new_doc_ids.append(upload.document_id)
    return ReingestResponse(job_ids = job_ids, document_ids = new_doc_ids)


@router.post(
    "/knowledge-bases/{kb_id}/reingest",
    response_model = ReingestResponse,
)
def reingest_knowledge_base(
    kb_id: str,
    payload: ReingestKBRequest,
    current_subject: str = Depends(get_current_subject),
) -> ReingestResponse:
    from utils.rag.config import resolve_embedder

    kb_row = _kb_or_404(kb_id)
    keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
    current_strategy = (
        kb_row["chunking_strategy"] if "chunking_strategy" in keys else "standard"
    )
    current_mode = kb_row["mode"] if "mode" in keys else "text"
    current_embedder = kb_row["embedding_model"]

    new_strategy = payload.chunking_strategy or current_strategy
    new_mode = payload.mode or current_mode
    _validate_mode_combo(new_mode, new_strategy)

    new_embedder = payload.embedding_model or (
        current_embedder
        if (new_strategy == current_strategy and new_mode == current_mode)
        else resolve_embedder(new_mode, new_strategy)
    )

    with get_connection() as conn:
        conn.execute(
            """
            UPDATE rag_knowledge_bases
            SET chunking_strategy = ?, mode = ?, embedding_model = ?
            WHERE id = ?
            """,
            (new_strategy, new_mode, new_embedder, kb_id),
        )
        conn.commit()

    return _reingest_scope(
        kb_id = kb_id,
        thread_id = None,
        chunking_strategy = new_strategy,
        mode = new_mode,
        embedding_model = new_embedder,
        caption_images = payload.caption_images is not False,
    )


@router.post(
    "/threads/{thread_id}/reingest",
    response_model = ReingestResponse,
)
def reingest_thread_documents(
    thread_id: str,
    payload: UpdateThreadRagSettingsRequest | None = None,
    current_subject: str = Depends(get_current_subject),
) -> ReingestResponse:
    """Rebuild a thread's RAG index; optional body updates settings before reingest."""
    from utils.rag.config import resolve_embedder

    if payload is None:
        payload = UpdateThreadRagSettingsRequest()
    if (
        payload.chunking_strategy is not None
        or payload.mode is not None
        or payload.embedding_model is not None
    ):
        settings = set_thread_rag_settings(
            thread_id,
            payload,
            current_subject = current_subject,
        )
    else:
        settings = _load_thread_settings(thread_id)

    embedder = settings.embedding_model or resolve_embedder(
        settings.mode,
        settings.chunking_strategy,
    )
    return _reingest_scope(
        kb_id = None,
        thread_id = thread_id,
        chunking_strategy = settings.chunking_strategy,
        mode = settings.mode,
        embedding_model = embedder,
        caption_images = payload.caption_images is not False,
    )


@router.delete("/knowledge-bases/{kb_id}")
def delete_knowledge_base(
    kb_id: str,
    current_subject: str = Depends(get_current_subject),
) -> dict:
    _kb_or_404(kb_id)
    with get_connection() as conn:
        doc_rows = conn.execute(
            "SELECT stored_path FROM rag_documents WHERE kb_id = ?",
            (kb_id,),
        ).fetchall()
        conn.execute("DELETE FROM rag_knowledge_bases WHERE id = ?", (kb_id,))
        conn.commit()
    for row in doc_rows:
        _unlink_if_under_uploads(Path(row["stored_path"]))
    ingestion.delete_scope_artifacts(kb_scope(kb_id))
    return {"ok": True}


# --- Document upload (KB and per-thread) ---


@router.post("/knowledge-bases/{kb_id}/documents", response_model = UploadResponse)
async def upload_kb_document(
    kb_id: str,
    file: UploadFile,
    caption_images: bool = True,
    current_subject: str = Depends(get_current_subject),
) -> UploadResponse:
    kb_row = _kb_or_404(kb_id)
    stored_path, filename, byte_size, content_hash = await _save_upload(file)
    # Tolerate pre-Phase-3 rows missing chunking_strategy/mode.
    kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
    chunking_strategy = (
        kb_row["chunking_strategy"] if "chunking_strategy" in kb_keys else "standard"
    )
    mode = kb_row["mode"] if "mode" in kb_keys else "text"
    return _start_ingestion(
        filename = filename,
        stored_path = stored_path,
        byte_size = byte_size,
        content_type = file.content_type,
        kb_id = kb_id,
        thread_id = None,
        embedding_model = kb_row["embedding_model"],
        chunking_strategy = chunking_strategy,
        mode = mode,
        caption_images = caption_images,
        content_hash = content_hash,
    )


@router.post("/threads/{thread_id}/documents", response_model = UploadResponse)
async def upload_thread_document(
    thread_id: str,
    file: UploadFile,
    caption_images: bool = True,
    current_subject: str = Depends(get_current_subject),
) -> UploadResponse:
    from utils.rag.config import resolve_embedder

    # No chat_threads check — fresh threads aren't persisted until first run.
    stored_path, filename, byte_size, content_hash = await _save_upload(file)
    settings = _load_thread_settings(thread_id)
    embedder = settings.embedding_model or resolve_embedder(
        settings.mode,
        settings.chunking_strategy,
    )
    return _start_ingestion(
        filename = filename,
        stored_path = stored_path,
        byte_size = byte_size,
        content_type = file.content_type,
        kb_id = None,
        thread_id = thread_id,
        embedding_model = embedder,
        chunking_strategy = settings.chunking_strategy,
        mode = settings.mode,
        caption_images = caption_images,
        content_hash = content_hash,
    )


# --- Document list / delete ---


@router.get("/knowledge-bases/{kb_id}/documents", response_model = DocumentListResponse)
def list_kb_documents(
    kb_id: str,
    current_subject: str = Depends(get_current_subject),
) -> DocumentListResponse:
    _kb_or_404(kb_id)
    with get_connection() as conn:
        rows = conn.execute(
            "SELECT * FROM rag_documents WHERE kb_id = ? ORDER BY created_at DESC",
            (kb_id,),
        ).fetchall()
    return DocumentListResponse(documents = [_row_to_document(r) for r in rows])


@router.get("/threads/{thread_id}/documents", response_model = DocumentListResponse)
def list_thread_documents(
    thread_id: str,
    current_subject: str = Depends(get_current_subject),
) -> DocumentListResponse:
    with get_connection() as conn:
        rows = conn.execute(
            "SELECT * FROM rag_documents WHERE thread_id = ? ORDER BY created_at DESC",
            (thread_id,),
        ).fetchall()
    return DocumentListResponse(documents = [_row_to_document(r) for r in rows])


@router.get("/images/{document_id}/{filename}")
def get_rag_image(
    document_id: str,
    filename: str,
    current_subject: str = Depends(get_current_subject),
) -> FileResponse:
    """Serve an extracted image; realpath-check against the uploads root."""
    document_for_subject_or_404(document_id, current_subject)
    if "/" in filename or "\\" in filename or filename.startswith("."):
        raise HTTPException(status_code = 400, detail = "Invalid filename")
    root = Path(os.path.realpath(rag_uploads_root() / "images"))
    candidate = rag_uploads_root() / "images" / document_id / filename
    try:
        real = Path(os.path.realpath(candidate))
        real.relative_to(root)
    except (OSError, ValueError) as exc:
        raise HTTPException(status_code = 404, detail = "Image not found") from exc
    if not real.is_file():
        raise HTTPException(status_code = 404, detail = "Image not found")
    return FileResponse(str(real))


@router.delete("/documents/{document_id}")
def delete_document(
    document_id: str,
    current_subject: str = Depends(get_current_subject),
) -> dict:
    row = _document_or_404(document_id)
    scope = kb_scope(row["kb_id"]) if row["kb_id"] else thread_scope(row["thread_id"])
    with get_connection() as conn:
        conn.execute("DELETE FROM rag_documents WHERE id = ?", (document_id,))
        conn.commit()
    _unlink_if_under_uploads(Path(row["stored_path"]))
    ingestion.delete_document_artifacts(document_id, scope)
    return {"ok": True}


@router.get("/thread-indexes", response_model = ThreadIndexListResponse)
def list_thread_indexes(
    current_subject: str = Depends(get_current_subject),
) -> ThreadIndexListResponse:
    """List threads with >=1 RAG doc. LEFT JOIN keeps unpersisted threads (null title)."""
    with get_connection() as conn:
        rows = conn.execute(
            """
            SELECT
                d.thread_id AS thread_id,
                t.title AS title,
                COUNT(DISTINCT d.id) AS num_documents,
                COALESCE(SUM(d.num_chunks), 0) AS num_chunks
            FROM rag_documents d
            LEFT JOIN chat_threads t ON t.id = d.thread_id
            WHERE d.thread_id IS NOT NULL
            GROUP BY d.thread_id, t.title
            ORDER BY MAX(d.created_at) DESC
            """
        ).fetchall()
    return ThreadIndexListResponse(
        threads = [
            ThreadIndexSummary(
                thread_id = r["thread_id"],
                title = r["title"],
                num_documents = int(r["num_documents"]),
                num_chunks = int(r["num_chunks"]),
            )
            for r in rows
        ]
    )


@router.delete("/threads/{thread_id}/documents")
def clear_thread_documents(
    thread_id: str,
    current_subject: str = Depends(get_current_subject),
) -> dict:
    """Drop all RAG artifacts for thread_id; chat thread itself untouched."""
    ingestion.purge_thread_documents([thread_id])
    return {"ok": True}


# --- Ingestion job SSE ---


@router.post("/jobs/{job_id}/cancel")
def cancel_job(
    job_id: str,
    current_subject: str = Depends(get_current_subject),
) -> dict:
    """Stop an in-flight ingestion job. The caller deletes the document
    afterwards to reset the index; this only halts the worker."""
    cancelled = ingestion.cancel_ingestion(job_id)
    return {"ok": True, "cancelled": cancelled}


@router.get("/jobs/{job_id}/events")
async def job_events(
    job_id: str,
    request: Request,
    current_subject: str = Depends(_sse_auth),
) -> StreamingResponse:
    state = ingestion.get_job_state(job_id)
    if state is None:
        with get_connection() as conn:
            row = conn.execute(
                "SELECT * FROM rag_ingestion_jobs WHERE id = ?",
                (job_id,),
            ).fetchone()
        if not row:
            raise HTTPException(status_code = 404, detail = "Job not found")
        return StreamingResponse(
            _replay_terminal_state(row),
            media_type = "text/event-stream",
        )

    consumer_queue = state.subscribe()

    async def stream():
        try:
            initial = {
                "type": "status",
                "status": state.status,
                "stage": state.stage,
                "progress": state.progress,
            }
            yield f"data: {json.dumps(initial)}\n\n"
            while True:
                if await request.is_disconnected():
                    break
                try:
                    event = await asyncio.get_event_loop().run_in_executor(
                        None,
                        consumer_queue.get,
                        True,
                        15.0,
                    )
                except queue_module.Empty:
                    yield ": keep-alive\n\n"
                    if state.status in ("completed", "failed"):
                        break
                    continue
                yield f"data: {json.dumps(event)}\n\n"
                if event.get("type") in ("complete", "error"):
                    break
        finally:
            state.unsubscribe(consumer_queue)

    return StreamingResponse(stream(), media_type = "text/event-stream")


async def _replay_terminal_state(row: Any):
    payload = {
        "type": "status",
        "status": row["status"],
        "stage": row["stage"],
        "progress": row["progress"],
        "error": row["error"],
    }
    yield f"data: {json.dumps(payload)}\n\n"


# --- Search ---


# --- Document preview (file + preview-target) ---


PreviewMediaKind = Literal["pdf", "text", "docx", "html", "image", "unknown"]
PreviewChunkKind = Literal["text", "image", "caption"]


class PreviewPdfRegion(BaseModel):
    pageIndex: int
    pageNumber: int | None = None
    x: float
    y: float
    width: float
    height: float
    confidence: Literal["exact"]
    source: str


class PreviewTargetResponse(BaseModel):
    """Per contracts.md §1.2 / §1.3 — single shape covering both
    cited-chunk and document-row preview modes. The §1.3 metadata-only
    mode returns ``None`` for ``chunkId``/``chunkIndex``/``targetPage``/
    ``snippet``/``kind``/``imageUrl`` (Q2: no first-chunk guessing).
    """

    documentId: str
    filename: str
    contentType: str | None
    mediaKind: PreviewMediaKind
    byteSize: int
    status: str
    kbId: str | None
    threadId: str | None
    chunkId: str | None
    chunkIndex: int | None
    targetPage: int | None
    snippet: str | None
    kind: PreviewChunkKind | None
    imageUrl: str | None
    sourcePageIndex: int | None
    pageCharStart: int | None
    pageCharEnd: int | None
    lineStart: int | None
    lineEnd: int | None
    pdfRegions: list[PreviewPdfRegion] = Field(default_factory = list)


class PreviewFileUrlResponse(BaseModel):
    url: str
    expiresAt: int


class LocatorBackfillResponse(BaseModel):
    documentId: str
    totalChunks: int
    matched: int
    alreadyLocated: int
    ambiguous: int
    missing: int
    skipped: int
    regionsMatched: int
    pagesRefreshed: int


# Extension allowlist for inline rendering / disposition. Anything not in
# this map collapses to ("application/octet-stream", attachment, "unknown").
# .html / .htm intentionally serve as text/plain attachment (decisions Q7 +
# Risk #3) so an uploaded HTML cannot execute in the app origin.
_PREVIEW_EXT_MAP: dict[str, tuple[str, str, PreviewMediaKind]] = {
    ".pdf": ("application/pdf", "inline", "pdf"),
    ".txt": ("text/plain; charset=utf-8", "inline", "text"),
    ".md": ("text/markdown; charset=utf-8", "inline", "text"),
    ".markdown": ("text/markdown; charset=utf-8", "inline", "text"),
    ".docx": (
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "attachment",
        "docx",
    ),
    ".html": ("text/plain; charset=utf-8", "attachment", "html"),
    ".htm": ("text/plain; charset=utf-8", "attachment", "html"),
    ".png": ("image/png", "inline", "image"),
    ".jpg": ("image/jpeg", "inline", "image"),
    ".jpeg": ("image/jpeg", "inline", "image"),
    ".gif": ("image/gif", "inline", "image"),
    ".webp": ("image/webp", "inline", "image"),
}


def _ascii_only(value: str) -> bool:
    try:
        value.encode("ascii")
    except UnicodeEncodeError:
        return False
    return True


def _content_disposition_header(filename: str, disposition: str) -> str:
    """Build a Content-Disposition header. Non-ASCII filenames use RFC 5987
    ``filename*=UTF-8''…`` alongside an ASCII-only ``filename=`` fallback so
    older clients still get something readable.
    """
    from urllib.parse import quote as _urlquote

    safe = filename.replace('"', "").replace("\r", "").replace("\n", "")
    if _ascii_only(safe):
        return f'{disposition}; filename="{safe}"'
    ascii_fallback = safe.encode("ascii", "replace").decode("ascii")
    encoded = _urlquote(safe, safe = "")
    return f'{disposition}; filename="{ascii_fallback}"; ' f"filename*=UTF-8''{encoded}"


def _preview_file_metadata(filename: str) -> tuple[str, str, PreviewMediaKind]:
    """Map a stored filename to (content_type, disposition, mediaKind).

    Unknown extensions always force ``application/octet-stream`` +
    ``attachment`` + ``unknown`` so the browser cannot sniff a sensitive
    type and inline it (Risk #3).
    """
    ext = Path(filename).suffix.lower()
    return _PREVIEW_EXT_MAP.get(
        ext, ("application/octet-stream", "attachment", "unknown")
    )


_PREVIEW_FILE_AUDIENCE = "rag-preview-file"
_PREVIEW_FILE_TTL_SECONDS = 5 * 60
_JWT_ALGORITHM = "HS256"


def _parse_pdf_regions(value: str | None) -> list[PreviewPdfRegion]:
    if not value:
        return []
    try:
        raw = json.loads(value)
    except (TypeError, json.JSONDecodeError):
        return []
    if not isinstance(raw, list):
        return []
    out: list[PreviewPdfRegion] = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        try:
            region = PreviewPdfRegion(**item)
        except Exception:
            continue
        if (
            0 <= region.x <= 1
            and 0 <= region.y <= 1
            and region.width > 0
            and region.height > 0
        ):
            out.append(region)
    return out


def _preview_file_token(
    *,
    subject: str,
    document_id: str,
) -> tuple[str, int]:
    secret = get_jwt_secret(subject)
    if secret is None:
        raise HTTPException(status_code = 401, detail = "Invalid or expired token")
    expires = datetime.now(timezone.utc) + timedelta(seconds = _PREVIEW_FILE_TTL_SECONDS)
    payload = {
        "sub": subject,
        "aud": _PREVIEW_FILE_AUDIENCE,
        "document_id": document_id,
        "exp": expires,
    }
    token = jwt.encode(payload, secret, algorithm = _JWT_ALGORITHM)
    return token, int(expires.timestamp())


def _subject_from_preview_file_token(document_id: str, token: str) -> str:
    try:
        unverified = jwt.decode(
            token,
            options = {
                "verify_signature": False,
                "verify_exp": False,
                "verify_aud": False,
            },
        )
    except jwt.InvalidTokenError as exc:
        raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
    subject = unverified.get("sub")
    if not isinstance(subject, str) or not subject:
        raise HTTPException(status_code = 401, detail = "Invalid preview token")
    secret = get_jwt_secret(subject)
    if secret is None:
        raise HTTPException(status_code = 401, detail = "Invalid preview token")
    try:
        payload = jwt.decode(
            token,
            secret,
            algorithms = [_JWT_ALGORITHM],
            audience = _PREVIEW_FILE_AUDIENCE,
        )
    except jwt.InvalidTokenError as exc:
        raise HTTPException(status_code = 401, detail = "Invalid preview token") from exc
    if payload.get("document_id") != document_id:
        raise HTTPException(status_code = 401, detail = "Invalid preview token")
    return subject


def _resolve_document_file_or_404(doc_row: Any, document_id: str) -> Path:
    try:
        resolved = resolve_under_root(
            doc_row["stored_path"],
            root = rag_uploads_root(),
        )
    except ValueError as exc:
        # Symlink escape / ``..`` / absolute outside root — collapse to
        # "file not found" (the auth row exists, the bytes do not).
        logger.warning(
            "RAG preview: stored_path escaped uploads root for doc %s: %s",
            document_id,
            exc,
        )
        raise HTTPException(
            status_code = 404,
            detail = "Document file not found",
        ) from exc

    if not resolved.is_file():
        raise HTTPException(
            status_code = 404,
            detail = "Document file not found",
        )
    return resolved


def _parse_range_header(range_header: str | None, size: int) -> tuple[int, int] | None:
    if not range_header:
        return None
    if not range_header.startswith("bytes="):
        raise ValueError("unsupported range unit")
    spec = range_header[len("bytes=") :].strip()
    if "," in spec or "-" not in spec:
        raise ValueError("multiple or malformed ranges are not supported")
    start_s, end_s = spec.split("-", 1)
    if not start_s and not end_s:
        raise ValueError("empty range")
    if not start_s:
        suffix = int(end_s)
        if suffix <= 0:
            raise ValueError("invalid suffix range")
        start = max(0, size - suffix)
        end = size - 1
    else:
        start = int(start_s)
        end = int(end_s) if end_s else size - 1
    if start < 0 or end < start or start >= size:
        raise ValueError("range outside file")
    return start, min(end, size - 1)


def _iter_file_range(path: Path, start: int, end: int):
    with path.open("rb") as fh:
        fh.seek(start)
        remaining = end - start + 1
        while remaining > 0:
            chunk = fh.read(min(64 * 1024, remaining))
            if not chunk:
                break
            remaining -= len(chunk)
            yield chunk


def _serve_document_file_row(
    doc_row: Any,
    document_id: str,
    range_header: str | None,
) -> FileResponse | Response | StreamingResponse:
    resolved = _resolve_document_file_or_404(doc_row, document_id)
    content_type, disposition, _media_kind = _preview_file_metadata(doc_row["filename"])
    safe_name = _sanitize_filename(doc_row["filename"])
    headers = {
        "Content-Disposition": _content_disposition_header(safe_name, disposition),
        "X-Content-Type-Options": "nosniff",
        "Cache-Control": "private, max-age=0, must-revalidate",
        "Accept-Ranges": "bytes",
    }
    size = resolved.stat().st_size

    try:
        byte_range = _parse_range_header(range_header, size)
    except (TypeError, ValueError):
        range_headers = dict(headers)
        range_headers["Content-Range"] = f"bytes */{size}"
        return Response(status_code = 416, headers = range_headers)

    if byte_range is not None:
        start, end = byte_range
        range_headers = dict(headers)
        range_headers["Content-Range"] = f"bytes {start}-{end}/{size}"
        range_headers["Content-Length"] = str(end - start + 1)
        return StreamingResponse(
            _iter_file_range(resolved, start, end),
            status_code = 206,
            media_type = content_type,
            headers = range_headers,
        )

    # Starlette's FileResponse handles ordinary downloads efficiently. We
    # still advertise Accept-Ranges so PDF.js can switch to explicit range
    # requests via the signed URL path.
    return FileResponse(
        path = str(resolved),
        media_type = content_type,
        headers = headers,
    )


@router.get(
    "/documents/{document_id}/preview-target",
    response_model = PreviewTargetResponse,
)
def get_document_preview_target(
    document_id: str,
    chunk_id: Optional[str] = Query(None),
    current_subject: str = Depends(get_current_subject),
) -> PreviewTargetResponse:
    """Resolve preview metadata for a document, optionally focused on a chunk.

    See contracts.md §1 for the response shape. ``chunk_id`` is the durable
    ``rag_chunks.id`` carried as ``backendChunkId`` on the frontend; when
    supplied it MUST belong to ``document_id`` or we collapse to 404 (so a
    probe cannot enumerate cross-document chunk ids).
    """
    doc_row = document_for_subject_or_404(document_id, current_subject)
    _ct, _disposition, media_kind = _preview_file_metadata(doc_row["filename"])

    base = {
        "documentId": doc_row["id"],
        "filename": doc_row["filename"],
        "contentType": doc_row["content_type"],
        "mediaKind": media_kind,
        "byteSize": int(doc_row["byte_size"]),
        "status": doc_row["status"],
        "kbId": doc_row["kb_id"],
        "threadId": doc_row["thread_id"],
    }

    if not chunk_id:
        # Q2: document-row preview returns metadata only — frontend MUST
        # NOT fall back to "first chunk".
        return PreviewTargetResponse(
            **base,
            chunkId = None,
            chunkIndex = None,
            targetPage = None,
            snippet = None,
            kind = None,
            imageUrl = None,
            sourcePageIndex = None,
            pageCharStart = None,
            pageCharEnd = None,
            lineStart = None,
            lineEnd = None,
            pdfRegions = [],
        )

    # Single connection enforces membership AND fetches the row in one
    # query. Splitting this into a separate `chunk_belongs_to_document`
    # call would open a second SQLite connection and create a TOCTOU
    # window — if the chunk is deleted between the two calls, the data
    # fetch returns None and the route 500s on the next attribute access
    # (devils-advocate D1.1). The cross-document case still collapses to
    # the same 404 the auth helper emits — never 400 (would leak doc
    # existence).
    with get_connection() as conn:
        chunk_row = conn.execute(
            """
            SELECT id, chunk_index, page_number, text, kind, image_path,
                   source_page_index, page_char_start, page_char_end,
                   line_start, line_end, pdf_regions_json
            FROM rag_chunks WHERE id = ? AND document_id = ?
            """,
            (chunk_id, document_id),
        ).fetchone()

    if chunk_row is None:
        raise HTTPException(
            status_code = 404,
            detail = "Document not found",
        )

    chunk_kind: PreviewChunkKind = chunk_row["kind"] or "text"  # type: ignore[assignment]
    image_url: str | None = None
    if chunk_kind == "image" and chunk_row["image_path"]:
        image_url = (
            f"/api/rag/images/{doc_row['id']}/" f"{Path(chunk_row['image_path']).name}"
        )

    return PreviewTargetResponse(
        **base,
        chunkId = chunk_row["id"],
        chunkIndex = int(chunk_row["chunk_index"]),
        targetPage = (
            int(chunk_row["page_number"])
            if chunk_row["page_number"] is not None
            else None
        ),
        snippet = chunk_row["text"] or "",
        kind = chunk_kind,
        imageUrl = image_url,
        sourcePageIndex = chunk_row["source_page_index"],
        pageCharStart = chunk_row["page_char_start"],
        pageCharEnd = chunk_row["page_char_end"],
        lineStart = chunk_row["line_start"],
        lineEnd = chunk_row["line_end"],
        pdfRegions = _parse_pdf_regions(chunk_row["pdf_regions_json"]),
    )


@router.post(
    "/documents/{document_id}/locators/backfill",
    response_model = LocatorBackfillResponse,
)
def backfill_document_locators_route(
    document_id: str,
    current_subject: str = Depends(get_current_subject),
) -> LocatorBackfillResponse:
    """In-place locator backfill for existing citations.

    This preserves ``document_id`` and ``chunk_id``. Chunks are updated only
    when their text has one unambiguous match in the parsed document text;
    duplicate or missing matches remain null.
    """
    doc_row = document_for_subject_or_404(document_id, current_subject)
    resolved = _resolve_document_file_or_404(doc_row, document_id)
    try:
        result = backfill_document_locators(document_id, resolved)
    except Exception as exc:
        logger.warning(
            "RAG locator backfill failed",
            document_id = document_id,
            error = str(exc),
        )
        raise HTTPException(
            status_code = 400,
            detail = "Document locators could not be backfilled",
        ) from exc
    return LocatorBackfillResponse(
        documentId = result.document_id,
        totalChunks = result.total_chunks,
        matched = result.matched,
        alreadyLocated = result.already_located,
        ambiguous = result.ambiguous,
        missing = result.missing,
        skipped = result.skipped,
        regionsMatched = result.regions_matched,
        pagesRefreshed = result.pages_refreshed,
    )


@router.get(
    "/documents/{document_id}/file-url",
    response_model = PreviewFileUrlResponse,
)
def get_document_file_url(
    document_id: str,
    current_subject: str = Depends(get_current_subject),
) -> PreviewFileUrlResponse:
    """Mint a short-lived signed URL for PDF.js range requests.

    The normal bearer-protected `/file` route stays available for blob
    fallback; this route keeps bearer access tokens out of query strings.
    """
    document_for_subject_or_404(document_id, current_subject)
    token, expires_at = _preview_file_token(
        subject = current_subject,
        document_id = document_id,
    )
    url = (
        f"/api/rag/documents/{quote(document_id, safe = '')}/file-signed"
        f"?token={quote(token, safe = '')}"
    )
    return PreviewFileUrlResponse(url = url, expiresAt = expires_at)


@router.get("/documents/{document_id}/file-signed", response_model = None)
def get_signed_document_file(
    document_id: str,
    request: Request,
    token: str = Query(...),
) -> FileResponse | Response | StreamingResponse:
    subject = _subject_from_preview_file_token(document_id, token)
    doc_row = document_for_subject_or_404(document_id, subject)
    return _serve_document_file_row(
        doc_row,
        document_id,
        request.headers.get("range"),
    )


@router.get("/documents/{document_id}/file", response_model = None)
def get_document_file(
    document_id: str,
    request: Request,
    current_subject: str = Depends(get_current_subject),
) -> FileResponse | Response | StreamingResponse:
    """Serve the original uploaded bytes for ``document_id``.

    Path resolution per contracts.md §2.1: the route NEVER accepts a
    client-supplied filename. ``stored_path`` is DB-issued and we run it
    through ``resolve_under_root`` which delegates to ``_assert_contained``
    (realpath + symlink/junction-safe). Any escape (symlink, junction,
    ``..``, absolute outside root) collapses to the second 404.

    Content-Type and disposition come from the extension allowlist
    (``_preview_file_metadata``). HTML/DOCX/unknown serve as
    ``attachment`` with safe content-type so an uploaded ``.html`` can
    never execute in the app origin (Risk #3 / decisions Q7).
    """
    doc_row = document_for_subject_or_404(document_id, current_subject)
    return _serve_document_file_row(
        doc_row,
        document_id,
        request.headers.get("range"),
    )


def _execute_search(
    scope: str,
    *,
    scope_embedder: str | None,
    query: str,
    mode: str,
    top_k: int,
    document_ids: list[str] | None,
    enable_rerank: bool,
    reranker_model: str | None,
    min_score: float,
) -> list[SearchHit]:
    """Run one retrieval against an already-resolved scope. Shared by the
    /search and /prefetch endpoints."""
    candidate_k = max(top_k, RAG_RERANK_CANDIDATE_K) if enable_rerank else top_k

    if mode == "bm25":
        hits = retrieval.retrieve_bm25(scope, query, candidate_k)
    elif mode == "dense":
        hits = retrieval.retrieve_dense(
            scope,
            query,
            candidate_k,
            document_ids = document_ids,
            embedder_model = scope_embedder,
        )
    else:
        hits = retrieval.retrieve_hybrid(
            scope,
            query,
            k = candidate_k,
            document_ids = document_ids,
            embedder_model = scope_embedder,
        )

    if min_score > 0.0:
        hits = retrieval.filter_by_min_score(hits, min_score)

    chunk_ids = [h.chunk_id for h in hits]
    chunk_lookup: dict[str, dict] = {}
    if chunk_ids:
        placeholders = ",".join("?" for _ in chunk_ids)
        with get_connection() as conn:
            rows = conn.execute(
                f"""
                SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text,
                       c.page_number, c.kind, c.image_path, c.linked_chunk_id,
                       c.source_page_index, c.page_char_start, c.page_char_end,
                       c.line_start, c.line_end,
                       d.filename
                FROM rag_chunks c
                JOIN rag_documents d ON d.id = c.document_id
                WHERE c.id IN ({placeholders})
                """,
                chunk_ids,
            ).fetchall()
        for r in rows:
            chunk_lookup[r["chunk_id"]] = dict(r)

    if enable_rerank:
        pairs = [
            (hit, chunk_lookup[hit.chunk_id]["text"])
            for hit in hits
            if hit.chunk_id in chunk_lookup
        ]
        hits = reranker.rerank(
            query,
            pairs,
            model_name = reranker_model,
            top_k = top_k,
        )
    else:
        hits = hits[:top_k]

    out: list[SearchHit] = []
    for hit in hits:
        meta = chunk_lookup.get(hit.chunk_id)
        if not meta:
            continue
        kind = meta.get("kind", "text") or "text"
        image_url: str | None = None
        if kind == "image" and meta.get("image_path"):
            image_url = (
                f"/api/rag/images/{meta['document_id']}/{Path(meta['image_path']).name}"
            )
        out.append(
            SearchHit(
                chunk_id = hit.chunk_id,
                document_id = meta["document_id"],
                chunk_index = meta["chunk_index"],
                text = meta["text"] or "",
                score = hit.score,
                page_number = meta.get("page_number"),
                filename = meta.get("filename"),
                kind = kind,
                image_url = image_url,
                source_page_index = meta.get("source_page_index"),
                page_char_start = meta.get("page_char_start"),
                page_char_end = meta.get("page_char_end"),
                line_start = meta.get("line_start"),
                line_end = meta.get("line_end"),
            )
        )
    return out


@router.post("/search", response_model = SearchResponse)
def search(
    payload: SearchRequest,
    current_subject: str = Depends(get_current_subject),
) -> SearchResponse:
    if bool(payload.kb_id) == bool(payload.thread_id):
        raise HTTPException(
            status_code = 400,
            detail = "exactly one of kb_id or thread_id must be supplied",
        )
    if payload.kb_id:
        _kb_or_404(payload.kb_id)
        scope = kb_scope(payload.kb_id)
    else:
        scope = thread_scope(payload.thread_id)

    # Query must use the same embedder as the scope (dim must match).
    scope_embedder = _resolve_scope_embedder(scope)
    logger.info(
        "RAG search: scope=%s embedder=%s mode=%s top_k=%d min_score=%.3f rerank=%s query=%r",
        scope,
        scope_embedder or "<default>",
        payload.mode,
        payload.top_k,
        payload.min_score,
        payload.enable_rerank,
        payload.query[:120],
    )

    out = _execute_search(
        scope,
        scope_embedder = scope_embedder,
        query = payload.query,
        mode = payload.mode,
        top_k = payload.top_k,
        document_ids = payload.document_ids,
        enable_rerank = payload.enable_rerank,
        reranker_model = payload.reranker_model,
        min_score = payload.min_score,
    )
    logger.info("RAG search: returning %d hits", len(out))
    return SearchResponse(hits = out)


@router.post("/prefetch", response_model = PrefetchResponse)
def prefetch(
    payload: PrefetchRequest,
    current_subject: str = Depends(get_current_subject),
) -> PrefetchResponse:
    """External-provider RAG prefetch: decompose the question via the helper,
    retrieve per sub-query, merge+dedup, return chunks for the frontend to
    inject before calling the provider. The local tool path is unaffected."""
    if bool(payload.kb_id) == bool(payload.thread_id):
        raise HTTPException(
            status_code = 400,
            detail = "exactly one of kb_id or thread_id must be supplied",
        )
    if payload.kb_id:
        _kb_or_404(payload.kb_id)
        scope = kb_scope(payload.kb_id)
    else:
        scope = thread_scope(payload.thread_id)

    scope_embedder = _resolve_scope_embedder(scope)

    # Momentarily load the pre-cached helper to split the question into up to
    # 3 focused queries; falls back to [query] on any failure.
    from core.rag.query_decompose import decompose_query

    queries = decompose_query(payload.query)
    logger.info(
        "RAG prefetch: scope=%s embedder=%s mode=%s n_queries=%d rerank=%s",
        scope,
        scope_embedder or "<default>",
        payload.mode,
        len(queries),
        payload.enable_rerank,
    )

    # Retrieve per query, merge, dedup by chunk_id (keep first/highest-ranked
    # occurrence), then cap at top_k.
    merged: list[SearchHit] = []
    seen: set[str] = set()
    for q in queries:
        hits = _execute_search(
            scope,
            scope_embedder = scope_embedder,
            query = q,
            mode = payload.mode,
            top_k = payload.top_k,
            document_ids = None,
            enable_rerank = payload.enable_rerank,
            reranker_model = payload.reranker_model,
            min_score = payload.min_score,
        )
        for h in hits:
            if h.chunk_id in seen:
                continue
            seen.add(h.chunk_id)
            merged.append(h)

    merged = merged[: payload.top_k]
    logger.info("RAG prefetch: returning %d merged hits", len(merged))
    return PrefetchResponse(queries = queries, hits = merged)