"""Core Store class: ingest, search, metadata queries.

Persistence:
  - Relational data (corpora, refs, papers, blocks, links) via SQLAlchemy
    ORM (portable across SQLite, Postgres, MySQL via connection string)
  - Vector search via Chroma (LlamaIndex) or pgvector (SQLAlchemy column)

Identity model:
  - ``Corpus`` = document type registry with write policy.
  - ``Ref`` = one row per known document (paper, note, todo, wiki, etc.).
    Holds slug, DOI, S2 ID, arxiv ID, title, etc.  Auto-int PK.
    Polymorphic on ``corpus_id``.
  - ``Paper`` = ingestion receipt (1:1 optional on Ref). PDF hash,
    bundle path. Only exists when content is ingested.
  - ``Block`` = text chunks from ingested documents (FK to Ref).
  - ``Link`` = slug-based edge between refs or blocks. Replaces Citation.
  - ``LinkType`` = relation registry (cites/cited_by, annotates, etc.).
"""

from __future__ import annotations

import gzip
import json
import logging
import re
import unicodedata
from collections.abc import Callable
from datetime import datetime
from pathlib import Path
from typing import Any

from precis_summary import pick_best_summary
from precis_summary.rake import telegram_precis
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, joinedload, sessionmaker

from acatome_store.config import StoreConfig
from acatome_store.models import (
    Base,
    Block,
    Corpus,
    Link,
    LinkType,
    Note,
    Paper,
    Ref,
    create_blocks_view,
    seed_block_types,
    seed_corpora,
    seed_link_types,
)
from acatome_store.vector import VectorIndex, create_index

log = logging.getLogger(__name__)

# Block types that skip embedding (same as acatome_extract.enrich)
_SKIP_EMBED_TYPES = {"section_header", "title", "author", "equation", "junk"}


def _get_embedder(
    config: StoreConfig,
) -> Callable[[list[str]], list[list[float]]] | None:
    """Build an embedding function from the store's configured profile.

    Returns None if the embedding backend is unavailable.
    """
    if config.embed_provider == "chroma":
        try:
            from chromadb.utils.embedding_functions import (
                DefaultEmbeddingFunction,
            )

            ef = DefaultEmbeddingFunction()

            def _chroma_embed(texts: list[str]) -> list[list[float]]:
                results = ef(texts)
                return [
                    e.tolist() if hasattr(e, "tolist") else list(e) for e in results
                ]

            return _chroma_embed
        except Exception:
            return None

    if config.embed_provider == "sentence-transformers":
        try:
            from sentence_transformers import SentenceTransformer

            model = SentenceTransformer(config.embed_model)
            dim = config.embed_index_dim or config.embed_dim

            def _st_embed(texts: list[str]) -> list[list[float]]:
                embs = model.encode(texts, normalize_embeddings=True)
                return [e[:dim].tolist() for e in embs]

            return _st_embed
        except Exception:
            return None

    return None


def _reembed_blocks(
    blocks: list[dict[str, Any]],
    embedder: Callable[[list[str]], list[list[float]]],
    profile: str = "default",
) -> list[dict[str, Any]]:
    """Re-embed blocks using the given embedder, replacing existing embeddings."""
    texts = []
    indices = []
    for i, b in enumerate(blocks):
        if b.get("type") in _SKIP_EMBED_TYPES:
            continue
        text = b.get("text", "").strip()
        if not text:
            continue
        texts.append(text)
        indices.append(i)

    if not texts:
        return blocks

    embeddings = embedder(texts)
    for idx, emb in zip(indices, embeddings):
        if "embeddings" not in blocks[idx]:
            blocks[idx]["embeddings"] = {}
        blocks[idx]["embeddings"][profile] = emb

    return blocks


class Store:
    """Acatome paper store.

    - SQLAlchemy ORM for relational CRUD (refs, papers, blocks, citations)
    - LlamaIndex ChromaVectorStore for Chroma vector search
    - pgvector column on blocks table for Postgres vector search
    """

    def __init__(self, config: StoreConfig | None = None):
        self._config = config or StoreConfig.from_global()
        self._index: VectorIndex | None = None
        self._ensure_dirs()
        self._init_db()

    def _ensure_dirs(self) -> None:
        self._config.store_path.mkdir(parents=True, exist_ok=True)

    def _init_db(self) -> None:
        """Create SQLAlchemy engine, session factory, and tables."""
        url = self._config.db_url
        is_pg = url.startswith("postgresql")

        if is_pg:
            try:
                from acatome_store.models import add_pgvector_column

                add_pgvector_column(self._config.embed_dim)
            except ImportError:
                pass

        self._engine = create_engine(url)

        # SQLite needs this for ON DELETE CASCADE to work
        if url.startswith("sqlite"):
            from sqlalchemy import event

            @event.listens_for(self._engine, "connect")
            def _set_sqlite_pragma(dbapi_conn, connection_record):
                cursor = dbapi_conn.cursor()
                cursor.execute("PRAGMA foreign_keys=ON")
                cursor.close()

        Base.metadata.create_all(self._engine)
        self._Session = sessionmaker(bind=self._engine)

        # Auto-migrate: add missing columns if Postgres
        if is_pg:
            self._ensure_embedding_column()
            self._ensure_missing_columns()

        # Seed lookup tables and create convenience view
        with self._Session() as session:
            seed_corpora(session)
            seed_block_types(session)
            seed_link_types(session)
        try:
            create_blocks_view(self._engine)
        except Exception:
            pass  # SQLite doesn't support CREATE OR REPLACE VIEW in all versions

    def _ensure_missing_columns(self) -> None:
        """Add missing columns to existing tables (Postgres only)."""
        migrations = [
            ("refs", "tags", "TEXT"),
            ("refs", "corpus_id", "VARCHAR DEFAULT 'papers' REFERENCES corpora(id)"),
            ("refs", "slug", "VARCHAR UNIQUE"),
            ("refs", "published_date", "DATE"),
            ("refs", "metadata", "TEXT"),
            ("corpora", "write_policy", "VARCHAR DEFAULT 'ingestion'"),
            ("notes", "origin", "VARCHAR"),
        ]
        try:
            from sqlalchemy import text

            with self._engine.connect() as conn:
                for table, column, col_type in migrations:
                    row = conn.execute(
                        text(
                            "SELECT 1 FROM information_schema.columns "
                            f"WHERE table_name='{table}' AND column_name='{column}'"
                        )
                    ).fetchone()
                    if not row:
                        conn.execute(
                            text(f"ALTER TABLE {table} ADD COLUMN {column} {col_type}")
                        )
                conn.commit()
        except Exception:
            pass  # non-fatal

    def _ensure_embedding_column(self) -> None:
        """Add embedding column to blocks table if missing (Postgres only)."""
        try:
            from sqlalchemy import text

            with self._engine.connect() as conn:
                conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
                row = conn.execute(
                    text(
                        "SELECT 1 FROM information_schema.columns "
                        "WHERE table_name='blocks' AND column_name='embedding'"
                    )
                ).fetchone()
                if not row:
                    conn.execute(
                        text("ALTER TABLE blocks ADD COLUMN embedding vector(384)")
                    )
                conn.commit()
        except Exception:
            pass  # pgvector not installed or other issue — non-fatal

    @property
    def index(self) -> VectorIndex:
        """Lazy-init vector index."""
        if self._index is None:
            sf = self._Session if self._config.vector_backend == "postgres" else None
            self._index = create_index(self._config, session_factory=sf)
        return self._index

    # ------------------------------------------------------------------
    # Ref helpers
    # ------------------------------------------------------------------

    def _find_ref(self, session: Session, header: dict) -> Ref | None:
        """Find an existing Ref by DOI, S2 ID, or arxiv ID."""
        doi = header.get("doi")
        if doi:
            ref = session.execute(
                select(Ref).where(Ref.doi == doi)
            ).scalar_one_or_none()
            if ref:
                return ref

        s2_id = header.get("s2_id")
        if s2_id:
            ref = session.execute(
                select(Ref).where(Ref.s2_id == s2_id)
            ).scalar_one_or_none()
            if ref:
                return ref

        arxiv_id = header.get("arxiv_id")
        if arxiv_id:
            ref = session.execute(
                select(Ref).where(Ref.arxiv_id == arxiv_id)
            ).scalar_one_or_none()
            if ref:
                return ref

        return None

    def _upsert_ref(self, session: Session, header: dict) -> Ref:
        """Find or create a Ref from bundle header, updating sparse fields."""
        ref = self._find_ref(session, header)
        if ref is None:
            ref = Ref()
            session.add(ref)

        # Update direct columns (fill in blanks, never overwrite with None)
        # Note: slug is NOT set here — ingest() handles it after collision check
        for field in ("doi", "s2_id", "arxiv_id", "title"):
            val = header.get(field)
            if val and not getattr(ref, field, None):
                setattr(ref, field, val)

        if header.get("authors"):
            if not ref.authors:
                ref.authors = json.dumps(header["authors"])
        if header.get("year") and not ref.year:
            ref.year = header["year"]
        if header.get("keywords"):
            if not ref.keywords:
                ref.keywords = json.dumps(header["keywords"])

        # Metadata JSON: journal, entry_type, source, etc.
        meta = ref._meta.copy() if ref.meta else {}
        for meta_field in ("journal", "entry_type", "source"):
            val = header.get(meta_field)
            if val and not meta.get(meta_field):
                meta[meta_field] = val
        if (meta and not ref.meta) or meta:
            ref.meta = json.dumps(meta)

        session.flush()  # ensure ref.id is assigned
        return ref

    @staticmethod
    def _merge_tags(ref: Ref, new_tags: list[str]) -> None:
        """Additively merge tags into a Ref (deduped, sorted)."""
        existing: list[str] = json.loads(ref.tags) if ref.tags else []
        merged = sorted(set(existing) | set(new_tags))
        ref.tags = json.dumps(merged)

    @staticmethod
    def _make_slug(authors: list[dict], year: int | None, title: str) -> str:
        """Generate human-readable slug: {surname}{year}{keyword}."""
        _STOP = {"a", "an", "the", "new", "on", "of", "in", "for", "to", "and", "with"}
        # First author surname
        surname = "anon"
        if authors:
            name = authors[0].get("name", "") if isinstance(authors[0], dict) else str(authors[0])
            if name and ";" in name:
                name = name.split(";")[0].strip()
            surname = name.split(",")[0].strip().lower() if name else "anon"
        surname = unicodedata.normalize("NFKD", surname).encode("ascii", "ignore").decode()
        surname = re.sub(r"[^a-z]", "", surname)[:30] or "anon"
        yr = str(year) if year else "0000"
        ascii_title = unicodedata.normalize("NFKD", title).encode("ascii", "ignore").decode()
        words = re.findall(r"[a-z]+", ascii_title.lower())
        keyword = next((w for w in words if w not in _STOP), words[0] if words else "")
        if not keyword:
            import hashlib
            keyword = hashlib.sha256(title.encode()).hexdigest()[:6] if title.strip() else "untitled"
        return f"{surname}{yr}{keyword}"

    @staticmethod
    def _disambiguate_slug(session: Session, base_slug: str) -> str:
        """Append a/b/c… suffix to make a slug unique."""
        for suffix in "abcdefghijklmnopqrstuvwxyz":
            candidate = f"{base_slug}{suffix}"
            exists = session.execute(
                select(Ref).where(Ref.slug == candidate)
            ).scalar_one_or_none()
            if not exists:
                return candidate
        # Extremely unlikely: 26 collisions — fall back to numeric
        n = 2
        while True:
            candidate = f"{base_slug}{n}"
            exists = session.execute(
                select(Ref).where(Ref.slug == candidate)
            ).scalar_one_or_none()
            if not exists:
                return candidate
            n += 1

    # ------------------------------------------------------------------
    # Ingest
    # ------------------------------------------------------------------

    def ingest(self, bundle_path: str | Path, *, tags: list[str] | None = None) -> int:
        """Ingest a .acatome bundle into the store.

        Dedup rules:
        - Same pdf_hash → skip (return existing ref_id)
        - Same DOI, different PDF → skip (return existing ref_id)
        - Same slug, different ref → raise ValueError
        - Same ref, re-ingest → atomic replace of Paper + blocks

        Tags are always merged additively, even on dedup skip.

        Returns:
            ref_id (int) of the ingested paper.
        """
        bundle_path = Path(bundle_path)
        data = _read_bundle(bundle_path)
        header = data["header"]
        slug = (
            header.get("slug", "").replace("~", "").replace("\u203a", "")
        )  # separators reserved
        pdf_hash = header["pdf_hash"]

        with self._Session() as session:
            # Dedup: same pdf_hash already ingested
            existing = session.execute(
                select(Paper).where(Paper.pdf_hash == pdf_hash)
            ).scalar_one_or_none()
            if existing:
                if tags:
                    ref = session.get(Ref, existing.ref_id)
                    if ref:
                        self._merge_tags(ref, tags)
                        session.commit()
                return existing.ref_id

            # Find or create the Ref (identity)
            ref = self._upsert_ref(session, header)

            # Dedup: ref has a Paper with a different PDF hash
            if ref.paper and ref.paper.pdf_hash != pdf_hash:
                return ref.id

            # Generate slug from metadata if bundle didn't provide one
            if not slug:
                slug = self._make_slug(
                    header.get("authors", []),
                    header.get("year"),
                    header.get("title", ""),
                )

            # Slug collision: auto-disambiguate with a/b/c… suffix
            existing_ref = session.execute(
                select(Ref).where(Ref.slug == slug)
            ).scalar_one_or_none()
            if existing_ref and existing_ref.id != ref.id:
                slug = self._disambiguate_slug(session, slug)
                log.info("slug collision → %s", slug)
            ref.slug = slug

            # Atomic replace: delete old Paper + blocks if re-ingesting
            if ref.paper:
                session.delete(ref.paper)
            # Blocks FK to ref, so delete them explicitly
            for old_block in list(ref.blocks):
                session.delete(old_block)
            session.flush()

            # Insert Paper (ingestion receipt)
            paper = Paper(
                ref_id=ref.id,
                pdf_hash=pdf_hash,
                bundle_path=str(bundle_path.resolve()),
                verified=bool(header.get("verified")),
            )
            session.add(paper)
            session.flush()

            # Insert abstract as a block (if present)
            abstract_text = header.get("abstract", "")
            if abstract_text:
                session.add(
                    Block(
                        node_id=f"ref:{ref.id}:abstract",
                        profile="default",
                        ref_id=ref.id,
                        page=None,
                        block_index=None,
                        block_type="abstract",
                        text=abstract_text,
                    )
                )

            # Insert document summary as a block (if present)
            enrich_meta = data.get("enrichment_meta") or {}
            # New format: paper_summaries dict; old format: paper_summary string
            paper_summaries = enrich_meta.get("paper_summaries") or {}
            doc_summary = pick_best_summary(paper_summaries) or enrich_meta.get(
                "paper_summary", ""
            )
            if doc_summary:
                session.add(
                    Block(
                        node_id=f"ref:{ref.id}:document_summary",
                        profile="default",
                        ref_id=ref.id,
                        page=None,
                        block_index=None,
                        block_type="document_summary",
                        text=doc_summary,
                    )
                )

            # Insert blocks (block_index is a global sequential counter)
            blocks = data.get("blocks", [])
            for block_index, b in enumerate(blocks):
                node_id = b["node_id"]

                bbox = b.get("bbox")
                block = Block(
                    node_id=node_id,
                    profile="default",
                    ref_id=ref.id,
                    page=b.get("page", 0),
                    block_index=block_index,
                    block_type=b.get("type", "text"),
                    text=b.get("text", ""),
                    summary=pick_best_summary(b.get("summaries")) or b.get("summary"),
                    section_path=json.dumps(b.get("section_path", [])),
                    bbox_x0=bbox[0] if bbox else None,
                    bbox_y0=bbox[1] if bbox else None,
                    bbox_x1=bbox[2] if bbox else None,
                    bbox_y1=bbox[3] if bbox else None,
                )
                session.add(block)

            # Aggregate block-level RAKE keywords into paper-level keywords
            if not ref.keywords:
                rake_parts = []
                for b in blocks:
                    summaries = b.get("summaries") or {}
                    rake = summaries.get("rake", "")
                    if rake:
                        rake_parts.append(rake)
                if rake_parts:
                    combined = "; ".join(rake_parts)
                    paper_kw = telegram_precis(combined, min_n=3, max_n=8)
                    ref.keywords = json.dumps(paper_kw.split("; "))

            if tags:
                self._merge_tags(ref, tags)

            session.commit()
            ref_id = ref.id

        # Check embedding model compatibility and re-embed if needed
        if blocks:
            bundle_models = enrich_meta.get("embedding_models", {})
            bundle_default = bundle_models.get("default", {})
            bundle_model = bundle_default.get("model", "")

            needs_reembed = (
                not bundle_model  # no enrichment metadata → embed from scratch
                or bundle_model != self._config.embed_model
            )

            if needs_reembed:
                log.warning(
                    "Embedding model mismatch: bundle has '%s', system expects '%s' "
                    "— re-embedding %d blocks",
                    bundle_model,
                    self._config.embed_model,
                    len(blocks),
                )
                embedder = _get_embedder(self._config)
                if embedder:
                    try:
                        blocks = _reembed_blocks(blocks, embedder)
                        # Write embeddings back to bundle so future reingests skip this
                        _update_bundle_embeddings(
                            bundle_path,
                            data,
                            blocks,
                            self._config.embed_model,
                            self._config.embed_dim,
                        )
                    except Exception:
                        log.warning("Re-embedding failed, skipping vector index")
                        blocks = []  # skip indexing
                else:
                    log.warning(
                        "No embedder available for '%s', skipping vector index",
                        self._config.embed_provider,
                    )
                    blocks = []  # skip indexing

            # Index embeddings in vector store (best-effort)
            try:
                self.index.add_blocks(str(ref_id), blocks)
            except Exception:
                pass

        return ref_id

    # ------------------------------------------------------------------
    # Direct ref creation (for direct-write corpora: todos, wiki, notes)
    # ------------------------------------------------------------------

    def create_ref(
        self,
        slug: str,
        *,
        corpus_id: str = "papers",
        title: str = "",
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        blocks: list[dict[str, Any]] | None = None,
    ) -> int:
        """Create a ref directly (no .acatome bundle needed).

        For use by direct-write corpora (todos, wiki, notes, journal).
        Raises ValueError if slug already exists or corpus not found.

        Args:
            slug: Unique slug for the ref.
            corpus_id: Corpus this ref belongs to.
            title: Human-readable title.
            metadata: JSON-serialisable dict stored in ref.metadata.
            tags: Optional list of tags.
            blocks: Optional list of block dicts with keys:
                text, block_type (default "text"), section_path (default []).

        Returns:
            ref_id (int) of the created ref.
        """
        if not slug:
            raise ValueError("slug is required for create_ref()")
        if "~" in slug or "\u203a" in slug:
            raise ValueError(
                f"slug must not contain '~' or '\u203a' (reserved as URI selector separator): {slug}"
            )

        with self._Session() as session:
            # Validate corpus exists
            corpus = session.get(Corpus, corpus_id)
            if not corpus:
                raise ValueError(f"Unknown corpus: {corpus_id}")

            # Check slug uniqueness
            existing = session.execute(
                select(Ref).where(Ref.slug == slug)
            ).scalar_one_or_none()
            if existing:
                raise ValueError(f"Slug already exists: {slug}")

            ref = Ref(
                corpus_id=corpus_id,
                slug=slug,
                title=title or None,
            )
            if metadata:
                ref.meta = json.dumps(metadata)
            session.add(ref)
            session.flush()  # get ref.id

            if tags:
                self._merge_tags(ref, tags)

            # Insert blocks
            for i, b in enumerate(blocks or []):
                node_id = f"{slug}-b{i:04d}"
                block = Block(
                    node_id=node_id,
                    profile="default",
                    ref_id=ref.id,
                    page=0,
                    block_index=i,
                    block_type=b.get("block_type", "text"),
                    text=b.get("text", ""),
                    section_path=json.dumps(b.get("section_path", [])),
                )
                session.add(block)

            session.commit()
            return ref.id

    def update_ref_metadata(
        self,
        slug: str,
        metadata: dict[str, Any],
        *,
        merge: bool = True,
    ) -> None:
        """Update the metadata JSON for a ref.

        Args:
            slug: Ref slug to update.
            metadata: Dict of keys to set.
            merge: If True (default), merge with existing metadata.
                   If False, replace entirely.

        Raises ValueError if ref not found.
        """
        with self._Session() as session:
            ref = session.execute(
                select(Ref).where(Ref.slug == slug)
            ).scalar_one_or_none()
            if not ref:
                raise ValueError(f"Ref not found: {slug}")

            if merge and ref.meta:
                try:
                    existing = json.loads(ref.meta)
                except (ValueError, TypeError):
                    existing = {}
                existing.update(metadata)
                ref.meta = json.dumps(existing)
            else:
                ref.meta = json.dumps(metadata)
            session.commit()

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

    def search_text(
        self,
        query: str,
        top_k: int = 5,
        where: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """Semantic text search via vector index.

        Returns list of result dicts with text, metadata, distance,
        and enriched paper info.
        """
        hits = self.index.search_text(query, top_k=top_k, where=where)

        for hit in hits:
            pid = hit.get("metadata", {}).get("paper_id")
            if pid:
                paper = self.get(pid)
                if paper:
                    hit["paper"] = {
                        "slug": paper.get("slug"),
                        "title": paper.get("title"),
                        "year": paper.get("year"),
                        "doi": paper.get("doi"),
                    }

        # Batch-fetch block summaries
        lookups = []
        for hit in hits:
            meta = hit.get("metadata", {})
            pid = meta.get("paper_id")
            bi = meta.get("block_index")
            if pid is not None and bi is not None:
                lookups.append((int(pid), int(bi)))
        if lookups:
            summaries = self._batch_block_summaries(lookups)
            for hit in hits:
                meta = hit.get("metadata", {})
                pid = meta.get("paper_id")
                bi = meta.get("block_index")
                if pid is not None and bi is not None:
                    s = summaries.get((int(pid), int(bi)))
                    if s:
                        hit["summary"] = s

        return hits

    def _batch_block_summaries(
        self, keys: list[tuple[int, int]]
    ) -> dict[tuple[int, int], str]:
        """Fetch block summaries for a list of (ref_id, block_index) pairs."""
        from sqlalchemy import tuple_

        with self._Session() as session:
            stmt = select(Block.ref_id, Block.block_index, Block.summary).where(
                tuple_(Block.ref_id, Block.block_index).in_(keys),
                Block.summary.isnot(None),
                Block.profile == "default",
            )
            rows = session.execute(stmt).all()
            return {(r[0], r[1]): r[2] for r in rows if r[2]}

    # ------------------------------------------------------------------
    # CRUD
    # ------------------------------------------------------------------

    def get(self, identifier) -> dict[str, Any] | None:
        """Get document by ref_id (int), slug, or DOI.

        Returns merged ref + paper dict, or None.
        """
        with self._Session() as session:
            ref = None

            # Try int ref_id
            if isinstance(identifier, int):
                ref = session.get(Ref, identifier)
            elif isinstance(identifier, str):
                # Try slug (on Ref)
                ref = session.execute(
                    select(Ref)
                    .options(joinedload(Ref.paper))
                    .where(Ref.slug == identifier)
                ).scalar_one_or_none()

                # Try DOI
                if not ref:
                    ref = session.execute(
                        select(Ref).where(Ref.doi == identifier)
                    ).scalar_one_or_none()

                # Try as int string
                if not ref:
                    try:
                        ref = session.get(Ref, int(identifier))
                    except (ValueError, TypeError):
                        pass

            if ref:
                return ref.to_dict()
        return None

    def get_blocks(
        self,
        identifier,
        block_type: str | None = None,
        supplement: str | None = None,
    ) -> list[dict[str, Any]]:
        """Get blocks for a paper by ref_id, slug, or DOI.

        Args:
            identifier: ref_id (int), slug, or DOI.
            block_type: Optional filter (e.g. "text", "abstract").
            supplement: If set, only blocks for that supplement.
                        If None (default), only main-paper blocks.
                        Use "*" for all blocks (main + supplements).
        """
        paper = self.get(identifier)
        if not paper or "ref_id" not in paper:
            return []
        ref_id = paper["ref_id"]
        with self._Session() as session:
            stmt = select(Block).where(Block.ref_id == ref_id)
            if block_type:
                stmt = stmt.where(Block.block_type == block_type)
            if supplement == "*":
                pass  # no filter — return all
            elif supplement is not None:
                stmt = stmt.where(Block.supplement == supplement)
            else:
                stmt = stmt.where(Block.supplement.is_(None))
            stmt = stmt.order_by(Block.page, Block.block_index)
            rows = session.execute(stmt).scalars().all()
            return [r.to_dict() for r in rows]

    def update_block_text(
        self,
        identifier,
        node_id: str,
        text: str,
    ) -> None:
        """Update the text content of a specific block.

        Args:
            identifier: ref_id (int), slug, or DOI.
            node_id: The block's node_id.
            text: New text content.

        Raises ValueError if ref or block not found.
        """
        paper = self.get(identifier)
        if not paper or "ref_id" not in paper:
            raise ValueError(f"Ref not found: {identifier}")
        ref_id = paper["ref_id"]
        with self._Session() as session:
            block = session.execute(
                select(Block).where(
                    Block.ref_id == ref_id,
                    Block.node_id == node_id,
                )
            ).scalar_one_or_none()
            if not block:
                raise ValueError(f"Block not found: {node_id} in ref {identifier}")
            block.text = text
            session.commit()

    # ------------------------------------------------------------------
    # Figures
    # ------------------------------------------------------------------

    _FIG_NUM_RE = re.compile(r"(?:Fig(?:ure)?\.?\s*|Scheme\s*)(\d+)", re.IGNORECASE)

    def get_figures(self, identifier) -> list[dict[str, Any]]:
        """Get figure metadata for a paper, with parsed figure numbers.

        Returns list of dicts: {fig_num, caption, page, block_index, node_id}.
        Figures are numbered by caption label (Fig 1, Figure 2, Scheme 3).
        Figures without a parseable number are assigned sequentially.
        """
        blocks = self.get_blocks(identifier, block_type="figure")
        figures: list[dict[str, Any]] = []
        next_auto = 1
        used_nums: set[int] = set()

        # First pass: extract explicit numbers
        for b in blocks:
            caption = b.get("text", "")
            m = self._FIG_NUM_RE.search(caption)
            num = int(m.group(1)) if m else None
            if num is not None:
                used_nums.add(num)
            figures.append(
                {
                    "fig_num": num,
                    "caption": caption,
                    "page": b.get("page"),
                    "block_index": b.get("block_index"),
                    "node_id": b.get("node_id", ""),
                }
            )

        # Second pass: assign auto numbers to figures without explicit labels
        for fig in figures:
            if fig["fig_num"] is None:
                while next_auto in used_nums:
                    next_auto += 1
                fig["fig_num"] = next_auto
                used_nums.add(next_auto)
                next_auto += 1

        return figures

    def get_figure_image(self, identifier, fig_num: int) -> dict[str, Any] | None:
        """Get figure image data from the bundle.

        Args:
            identifier: slug, ref_id, or DOI.
            fig_num: Figure number (from caption label).

        Returns:
            Dict with keys: fig_num, caption, page, image_bytes, image_ext,
            or None if not found.
        """
        # Find the figure metadata
        figures = self.get_figures(identifier)
        fig_meta = next((f for f in figures if f["fig_num"] == fig_num), None)
        if fig_meta is None:
            return None

        # Read the bundle to get image bytes
        paper = self.get(identifier)
        if not paper:
            return None

        with self._Session() as session:
            paper_row = session.execute(
                select(Paper).where(Paper.ref_id == paper["ref_id"])
            ).scalar_one_or_none()
            if not paper_row or not paper_row.bundle_path:
                return None
            bundle_path = Path(paper_row.bundle_path)

        if not bundle_path.exists():
            return None

        data = _read_bundle(bundle_path)
        # Match by node_id (reliable across bundle and store)
        target_node_id = fig_meta["node_id"]
        for block in data.get("blocks", []):
            if block.get("node_id") == target_node_id and block.get("type") == "figure":
                b64 = block.get("image_base64", "")
                mime = block.get("image_mime", "image/png")
                if not b64:
                    return None
                import base64

                ext = ".png" if "png" in mime else ".jpg"
                return {
                    "fig_num": fig_num,
                    "caption": fig_meta["caption"],
                    "page": fig_meta["page"],
                    "image_bytes": base64.b64decode(b64),
                    "image_ext": ext,
                }

        return None

    def get_toc(
        self, identifier, supplement: str | None = None
    ) -> list[dict[str, Any]]:
        """Get a lightweight table-of-contents for a paper.

        Returns one entry per block: node_id, page, block_type,
        summary (if available, else truncated text), section_path.
        """
        paper = self.get(identifier)
        if not paper or "ref_id" not in paper:
            return []
        ref_id = paper["ref_id"]
        with self._Session() as session:
            stmt = (
                select(Block)
                .where(Block.ref_id == ref_id)
                .order_by(Block.page, Block.block_index)
            )
            if supplement == "*":
                pass
            elif supplement is not None:
                stmt = stmt.where(Block.supplement == supplement)
            else:
                stmt = stmt.where(Block.supplement.is_(None))
            rows = session.execute(stmt).scalars().all()
            toc = []
            for r in rows:
                preview = r.summary or (
                    r.text[:120] + "…" if len(r.text) > 120 else r.text
                )
                entry = {
                    "node_id": r.node_id,
                    "block_index": r.block_index,
                    "page": r.page,
                    "block_type": r.block_type,
                    "section_path": r.section_path,
                    "preview": preview,
                }
                if r.summary:
                    entry["has_summary"] = True
                toc.append(entry)
            return toc

    def reindex_blocks(self, identifier=None) -> int:
        """Reassign block_index as a global sequential counter per paper.

        Fixes papers ingested with the old per-page block_index scheme.
        If *identifier* is given, reindex that paper only; otherwise all.
        Returns the number of papers reindexed.
        """
        papers = (
            [self.get(identifier)] if identifier else self.list_papers(limit=10_000)
        )
        count = 0
        with self._Session() as session:
            for paper in papers:
                if not paper or "ref_id" not in paper:
                    continue
                ref_id = paper["ref_id"]
                rows = (
                    session.execute(
                        select(Block)
                        .where(
                            Block.ref_id == ref_id,
                            Block.supplement.is_(None),
                        )
                        .order_by(Block.page, Block.block_index, Block.node_id)
                    )
                    .scalars()
                    .all()
                )
                for new_idx, row in enumerate(rows):
                    if row.block_index != new_idx:
                        row.block_index = new_idx
                count += 1
            session.commit()
        return count

    def delete(self, identifier) -> bool:
        """Delete the ingested Paper + blocks (keeps the Ref stub)."""
        paper_dict = self.get(identifier)
        if not paper_dict or "ref_id" not in paper_dict:
            return False
        ref_id = paper_dict["ref_id"]
        with self._Session() as session:
            ref = session.get(Ref, ref_id)
            if not ref:
                return False
            # Blocks FK to ref — delete explicitly
            for block in list(ref.blocks):
                session.delete(block)
            if ref.paper:
                session.delete(ref.paper)
            session.commit()
        try:
            self.index.delete_paper(str(ref_id))
        except Exception:
            pass
        return True

    def list_papers(
        self,
        limit: int = 100,
        offset: int = 0,
        since: datetime | None = None,
    ) -> list[dict[str, Any]]:
        """List ingested papers (joined with ref metadata).

        Args:
            limit: Max papers to return.
            offset: Pagination offset.
            since: Only return papers ingested on or after this datetime (UTC, naive).
        """
        with self._Session() as session:
            stmt = (
                select(Ref)
                .join(Paper)
                .options(joinedload(Ref.paper))
                .order_by(Paper.ingested_at.desc())
            )
            if since is not None:
                stmt = stmt.where(Paper.ingested_at >= since)
            rows = (
                session.execute(stmt.limit(limit).offset(offset))
                .unique()
                .scalars()
                .all()
            )
            result = []
            for r in rows:
                block_count = (
                    session.query(Block)
                    .filter(Block.ref_id == r.id, Block.profile == "default")
                    .count()
                )
                kw = None
                if r.keywords:
                    try:
                        kw = json.loads(r.keywords)
                    except (json.JSONDecodeError, TypeError):
                        kw = r.keywords
                tags: list[str] = []
                if r.tags:
                    try:
                        tags = json.loads(r.tags)
                    except (json.JSONDecodeError, TypeError):
                        pass
                result.append(
                    {
                        "ref_id": r.id,
                        "slug": r.slug,
                        "title": r.title,
                        "authors": r.authors,
                        "year": r.year,
                        "doi": r.doi,
                        "corpus_id": r.corpus_id,
                        "block_count": block_count,
                        "keywords": kw,
                        "tags": tags,
                        "ingested_at": r.paper.ingested_at if r.paper else None,
                    }
                )
            return result

    def stats(self) -> dict[str, Any]:
        """Return store statistics and connection info."""
        cfg = self._config
        with self._Session() as session:
            total_refs = session.query(Ref).count()
            total_ingested = session.query(Paper).count()
            total_blocks = session.query(Block).count()
            verified = session.query(Paper).filter(Paper.verified.is_(True)).count()
        indexed = 0
        try:
            indexed = self.index.count()
        except Exception:
            pass

        # Connection info
        info: dict[str, Any] = {
            "metadata_backend": cfg.metadata_backend,
        }
        if cfg.metadata_backend == "postgres":
            info["pg_host"] = cfg.pg_host
            info["pg_port"] = cfg.pg_port
            info["pg_database"] = cfg.pg_database
            info["pg_schema"] = cfg.pg_schema
            info["pg_user"] = cfg.pg_user
        else:
            info["db_path"] = str(cfg.store_path / "acatome.db")

        info.update(
            {
                "vector_backend": cfg.vector_backend,
                "embed_model": cfg.embed_model,
                "embed_dim": cfg.embed_dim,
                "store_path": str(cfg.store_path),
                "total_refs": total_refs,
                "total_papers": total_ingested,
                "total_blocks": total_blocks,
                "verified": verified,
                "indexed_blocks": indexed,
            }
        )
        return info

    # ------------------------------------------------------------------
    # Links
    # ------------------------------------------------------------------

    def create_link(
        self,
        src_slug: str,
        dst_slug: str,
        relation: str = "cites",
        *,
        src_node_id: str | None = None,
        dst_node_id: str | None = None,
    ) -> Link:
        """Create a link between two refs or blocks.

        Args:
            src_slug: Source ref slug.
            dst_slug: Target ref slug.
            relation: Link type name (must exist in link_types table).
            src_node_id: Optional source block node_id.
            dst_node_id: Optional target block node_id.

        Returns:
            The created Link object.

        Raises:
            ValueError: If relation is not a valid link type, or if
                src_slug/dst_slug don't resolve to existing refs.
        """
        with self._Session() as session:
            # Validate relation
            lt = session.get(LinkType, relation)
            if not lt:
                valid = [r.name for r in session.execute(select(LinkType)).scalars()]
                raise ValueError(
                    f"Unknown link relation '{relation}'. "
                    f"Valid types: {', '.join(sorted(valid))}"
                )
            # Validate source ref exists
            src_ref = session.execute(
                select(Ref).where(Ref.slug == src_slug)
            ).scalar_one_or_none()
            if not src_ref:
                raise ValueError(f"Source ref not found: '{src_slug}'")
            # Validate target ref exists
            dst_ref = session.execute(
                select(Ref).where(Ref.slug == dst_slug)
            ).scalar_one_or_none()
            if not dst_ref:
                raise ValueError(f"Target ref not found: '{dst_slug}'")

            link = Link(
                src_slug=src_slug,
                src_node_id=src_node_id,
                dst_slug=dst_slug,
                dst_node_id=dst_node_id,
                relation=relation,
            )
            session.add(link)
            session.commit()
            session.refresh(link)
            return link

    def get_links(
        self,
        slug: str,
        *,
        node_id: str | None = None,
        relation: str | None = None,
        direction: str = "both",
    ) -> list[dict[str, Any]]:
        """Get links for a ref or block.

        Args:
            slug: Ref slug.
            node_id: Optional block node_id (filters to block-level links).
            relation: Optional filter by link type name.
            direction: 'outbound', 'inbound', or 'both'.

        Returns:
            List of link dicts, each with an extra 'display_relation' key
            showing the natural-language label from the viewing side.

        Raises:
            ValueError: If direction is not one of 'outbound', 'inbound', 'both'.
        """
        if direction not in ("outbound", "inbound", "both"):
            raise ValueError(
                f"Invalid direction '{direction}', must be 'outbound', 'inbound', or 'both'"
            )
        results = []
        with self._Session() as session:
            # Build outbound query (this slug is src)
            if direction in ("outbound", "both"):
                q = select(Link).where(Link.src_slug == slug)
                if node_id is not None:
                    q = q.where(Link.src_node_id == node_id)
                if relation:
                    q = q.where(Link.relation == relation)
                q = q.order_by(Link.created_at)
                for link in session.execute(q).scalars():
                    d = link.to_dict()
                    d["display_relation"] = link.relation
                    d["direction"] = "outbound"
                    results.append(d)

            # Build inbound query (this slug is dst)
            if direction in ("inbound", "both"):
                q = select(Link).where(Link.dst_slug == slug)
                if node_id is not None:
                    q = q.where(Link.dst_node_id == node_id)
                if relation:
                    q = q.where(Link.relation == relation)
                q = q.order_by(Link.created_at)
                for link in session.execute(q).scalars():
                    d = link.to_dict()
                    # Use inverse label for inbound links
                    lt = session.get(LinkType, link.relation)
                    d["display_relation"] = lt.inverse if lt else link.relation
                    d["direction"] = "inbound"
                    results.append(d)

        return results

    def get_link_count(self, slug: str) -> dict[str, int]:
        """Get link counts for a ref, grouped by display relation.

        Returns:
            Dict mapping display_relation to count.
            E.g. {"cites": 3, "cited_by": 5, "annotated_by": 2}
        """
        links = self.get_links(slug)
        counts: dict[str, int] = {}
        for link in links:
            rel = link["display_relation"]
            counts[rel] = counts.get(rel, 0) + 1
        return counts

    def delete_link(self, link_id: int) -> bool:
        """Delete a link by ID.

        Returns:
            True if deleted, False if link not found.
        """
        with self._Session() as session:
            link = session.get(Link, link_id)
            if not link:
                return False
            session.delete(link)
            session.commit()
            return True

    def delete_links_for_slug(self, slug: str) -> int:
        """Delete all links where slug is src or dst.

        Returns:
            Number of links deleted.
        """
        from sqlalchemy import delete as sa_delete
        from sqlalchemy import or_

        with self._Session() as session:
            stmt = sa_delete(Link).where(
                or_(Link.src_slug == slug, Link.dst_slug == slug)
            )
            result = session.execute(stmt)
            session.commit()
            return result.rowcount

    # ------------------------------------------------------------------
    # Notes (DEPRECATED — use links with relation='annotates')
    # ------------------------------------------------------------------

    def add_note(
        self,
        content: str,
        *,
        ref_id: int | None = None,
        block_node_id: str | None = None,
        block_profile: str = "default",
        title: str | None = None,
        tags: list[str] | None = None,
        origin: str | None = None,
    ) -> int:
        """Add a note on a paper (ref_id) or block (block_node_id).

        Returns note id.
        """
        with self._Session() as session:
            note = Note(
                ref_id=ref_id,
                block_node_id=block_node_id,
                block_profile=block_profile if block_node_id else None,
                title=title,
                content=content,
                tags=json.dumps(tags) if tags else None,
                origin=origin,
            )
            session.add(note)
            session.commit()
            return note.id

    def get_notes(
        self,
        *,
        ref_id: int | None = None,
        block_node_id: str | None = None,
    ) -> list[dict[str, Any]]:
        """Get notes for a paper or block."""
        with self._Session() as session:
            q = select(Note)
            if ref_id is not None:
                q = q.where(Note.ref_id == ref_id)
            if block_node_id is not None:
                q = q.where(Note.block_node_id == block_node_id)
            q = q.order_by(Note.created_at.desc())
            rows = session.execute(q).scalars().all()
            return [r.to_dict() for r in rows]

    def update_note(self, note_id: int, **kwargs) -> bool:
        """Update a note. Pass content=, title=, tags= as kwargs."""
        with self._Session() as session:
            note = session.get(Note, note_id)
            if not note:
                return False
            for key in ("content", "title"):
                if key in kwargs:
                    setattr(note, key, kwargs[key])
            if "tags" in kwargs:
                note.tags = json.dumps(kwargs["tags"]) if kwargs["tags"] else None
            session.commit()
            return True

    def delete_note(self, note_id: int) -> bool:
        """Delete a note by id."""
        with self._Session() as session:
            note = session.get(Note, note_id)
            if not note:
                return False
            session.delete(note)
            session.commit()
            return True

    # ------------------------------------------------------------------
    # Supplements
    # ------------------------------------------------------------------

    def ingest_supplement(
        self,
        parent_identifier,
        bundle_path: str | Path,
        supplement_name: str,
    ) -> int:
        """Ingest a supplement PDF bundle into an existing paper.

        The supplement blocks are stored with ``supplement=name`` on the
        parent paper's ref_id.  The paper's ``supplements`` JSON list is
        updated to include the new name.

        Args:
            parent_identifier: ref_id (int), slug, or DOI of the parent.
            bundle_path: Path to the .acatome bundle for the supplement.
            supplement_name: Lowercase string id (e.g. "s1", "methods").

        Returns:
            ref_id (int) of the parent paper.
        """
        parent = self.get(parent_identifier)
        if not parent or "ref_id" not in parent:
            raise ValueError(f"Parent paper not found: {parent_identifier}")

        ref_id = parent["ref_id"]
        supplement_name = supplement_name.lower()
        bundle_path = Path(bundle_path)
        data = _read_bundle(bundle_path)

        with self._Session() as session:
            paper = session.execute(
                select(Paper).where(Paper.ref_id == ref_id)
            ).scalar_one_or_none()
            if not paper:
                raise ValueError(f"No ingested paper for ref_id={ref_id}")

            # Delete old supplement blocks with same name (re-ingest)
            old = (
                session.execute(
                    select(Block).where(
                        Block.ref_id == ref_id,
                        Block.supplement == supplement_name,
                    )
                )
                .scalars()
                .all()
            )
            for b in old:
                session.delete(b)
            session.flush()

            # Insert supplement blocks (block_index is a global sequential counter)
            blocks = data.get("blocks", [])
            for block_index, b in enumerate(blocks):
                node_id = b["node_id"]

                bbox = b.get("bbox")
                block = Block(
                    node_id=f"{node_id}:supp:{supplement_name}",
                    profile="default",
                    ref_id=ref_id,
                    page=b.get("page", 0),
                    block_index=block_index,
                    block_type=b.get("type", "text"),
                    text=b.get("text", ""),
                    summary=pick_best_summary(b.get("summaries")) or b.get("summary"),
                    supplement=supplement_name,
                    section_path=json.dumps(b.get("section_path", [])),
                    bbox_x0=bbox[0] if bbox else None,
                    bbox_y0=bbox[1] if bbox else None,
                    bbox_x1=bbox[2] if bbox else None,
                    bbox_y1=bbox[3] if bbox else None,
                )
                session.add(block)

            # Update supplements list on paper
            existing = json.loads(paper.supplements) if paper.supplements else []
            if supplement_name not in existing:
                existing.append(supplement_name)
                existing.sort()
            paper.supplements = json.dumps(existing)

            session.commit()

        # Index embeddings (best-effort)
        if blocks:
            try:
                self.index.add_blocks(str(ref_id), blocks)
            except Exception:
                pass

        return ref_id

    def get_supplements(self, identifier) -> list[str]:
        """Return list of supplement names for a paper, or empty list."""
        paper = self.get(identifier)
        if not paper:
            return []
        supps = paper.get("supplements")
        if not supps:
            return []
        try:
            return json.loads(supps)
        except (json.JSONDecodeError, TypeError):
            return []

    # ------------------------------------------------------------------
    # Tags
    # ------------------------------------------------------------------

    def add_tags(self, identifier, tags: list[str]) -> bool:
        """Add tags to a paper (additive, deduped).

        Args:
            identifier: ref_id (int), slug, or DOI.
            tags: Tag strings to add.

        Returns:
            True if the ref was found and updated.
        """
        paper_dict = self.get(identifier)
        if not paper_dict or "id" not in paper_dict:
            return False
        with self._Session() as session:
            ref = session.get(Ref, paper_dict["id"])
            if not ref:
                return False
            self._merge_tags(ref, tags)
            session.commit()
        return True

    def remove_tags(self, identifier, tags: list[str]) -> bool:
        """Remove tags from a paper.

        Args:
            identifier: ref_id (int), slug, or DOI.
            tags: Tag strings to remove.

        Returns:
            True if the ref was found and updated.
        """
        paper_dict = self.get(identifier)
        if not paper_dict or "id" not in paper_dict:
            return False
        with self._Session() as session:
            ref = session.get(Ref, paper_dict["id"])
            if not ref:
                return False
            existing: list[str] = json.loads(ref.tags) if ref.tags else []
            remaining = sorted(set(existing) - set(tags))
            ref.tags = json.dumps(remaining) if remaining else None
            session.commit()
        return True

    def get_tags(self, identifier) -> list[str]:
        """Return tags for a paper, or empty list."""
        paper_dict = self.get(identifier)
        if not paper_dict:
            return []
        raw = paper_dict.get("tags")
        if not raw:
            return []
        try:
            return json.loads(raw) if isinstance(raw, str) else raw
        except (json.JSONDecodeError, TypeError):
            return []

    def find_by_tag(self, tag: str) -> list[dict[str, Any]]:
        """Find all papers with a given tag.

        Returns list of dicts with ref_id, slug, title, year, tags.
        """
        with self._Session() as session:
            refs = session.execute(select(Ref)).scalars().all()
            results = []
            for r in refs:
                tags: list[str] = json.loads(r.tags) if r.tags else []
                if tag in tags:
                    results.append(
                        {
                            "ref_id": r.id,
                            "slug": r.slug,
                            "title": r.title,
                            "year": r.year,
                            "tags": tags,
                        }
                    )
            return results

    def list_tags(self) -> dict[str, int]:
        """Return all tags with counts, sorted by count descending."""
        with self._Session() as session:
            refs = session.execute(select(Ref)).scalars().all()
            counts: dict[str, int] = {}
            for r in refs:
                tags: list[str] = json.loads(r.tags) if r.tags else []
                for t in tags:
                    counts[t] = counts.get(t, 0) + 1
            return dict(sorted(counts.items(), key=lambda x: -x[1]))

    # ------------------------------------------------------------------
    # Retractions
    # ------------------------------------------------------------------

    def retract(self, identifier, note: str = "") -> bool:
        """Mark a document as retracted.

        Args:
            identifier: ref_id (int), slug, or DOI.
            note: Optional retraction note.

        Returns:
            True if the ref was found and updated.
        """
        paper_dict = self.get(identifier)
        if not paper_dict or "id" not in paper_dict:
            return False
        with self._Session() as session:
            ref = session.get(Ref, paper_dict["id"])
            if not ref:
                return False
            ref._set_meta_field("retracted", True)
            if note:
                ref._set_meta_field("retraction_note", note)
            session.commit()
        return True

    def unretract(self, identifier) -> bool:
        """Remove retraction flag from a document."""
        paper_dict = self.get(identifier)
        if not paper_dict or "id" not in paper_dict:
            return False
        with self._Session() as session:
            ref = session.get(Ref, paper_dict["id"])
            if not ref:
                return False
            ref._set_meta_field("retracted", False)
            ref._set_meta_field("retraction_note", None)
            session.commit()
        return True

    def reset_schema(self) -> None:
        """Drop all tables and recreate from the current model.

        WARNING: destroys all data. Use for dev/reingest workflows only.
        """
        # Drop dependent views first (Postgres blocks DROP TABLE otherwise)
        if self._config.db_url.startswith("postgresql"):
            from sqlalchemy import text

            with self._engine.connect() as conn:
                conn.execute(text("DROP VIEW IF EXISTS blocks_v CASCADE"))
                conn.commit()
        Base.metadata.drop_all(self._engine)
        self._init_db()

    def close(self) -> None:
        """Dispose of the engine."""
        if hasattr(self, "_engine"):
            self._engine.dispose()


def _update_bundle_embeddings(
    bundle_path: Path,
    data: dict[str, Any],
    blocks: list[dict[str, Any]],
    embed_model: str,
    embed_dim: int,
) -> None:
    """Write re-embedded blocks back into the .acatome bundle.

    Updates the blocks' embeddings and enrichment_meta so future
    reingests don't need to re-embed.
    """
    data["blocks"] = blocks
    em = data.get("enrichment_meta") or {}
    models = em.get("embedding_models") or {}
    models["default"] = {"model": embed_model, "dim": embed_dim}
    em["embedding_models"] = models
    em.setdefault("profiles", [])
    if "default" not in em["profiles"]:
        em["profiles"].append("default")
    data["enrichment_meta"] = em
    try:
        with gzip.open(bundle_path, "wt", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, separators=(",", ":"))
    except Exception:
        log.warning("Failed to write embeddings back to %s", bundle_path)


def _read_bundle(path: str | Path) -> dict[str, Any]:
    """Read a .acatome bundle (gzipped JSON)."""
    path = Path(path)
    with gzip.open(path, "rt", encoding="utf-8") as f:
        return json.load(f)
