document_store.py

import json
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Literal, Optional

from docling_core.types.doc.document import DoclingDocument

from models import DocumentRecord, FigureRecord


class DocumentStore:
    """Stores PDF copies and manifest metadata, keyed by binary_hash.

    Directory layout (collection root = base_path):
        <base_path>/manifest.json
        <base_path>/documents/<hash>/source.pdf
        <base_path>/documents/<hash>/meta.json        (figure metadata — heavy data)
        <base_path>/documents/<hash>/docling.json     (serialised DoclingDocument)
        <base_path>/documents/<hash>/figures/pic_000.png, tbl_000.png

    The binary_hash is used as the document key and is stored on every
    Chunk as ``file_hash``, enabling round-trip PDF retrieval for visual
    grounding via ``get_pdf_path``.

    The manifest.json is kept lightweight (no figures array, no pdf_path)
    so it loads instantly even for large collections.
    """

    def __init__(self, base_path: str = "./vector_data/vector"):
        """Initialise the store rooted at the collection directory.

        Args:
            base_path: Path to the collection root directory
                       (e.g. ``./vector_data/my_collection``).
                       ``manifest.json`` lives directly under this path.
        """
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @property
    def _manifest_path(self) -> Path:
        return self.base_path / "manifest.json"

    def _pdf_path(self, file_hash: str) -> Path:
        return self.base_path / "documents" / file_hash / "source.pdf"

    def _figures_dir(self, file_hash: str) -> Path:
        return self.base_path / "documents" / file_hash / "figures"

    def _meta_path(self, file_hash: str) -> Path:
        """Path to the per-document meta.json (heavy figure data)."""
        return self.base_path / "documents" / file_hash / "meta.json"

    def get_docling_doc_path(self, file_hash: str) -> Path:
        """Return the path to the stored DoclingDocument JSON for a document."""
        return self.base_path / "documents" / file_hash / "docling.json"

    def _read_manifest(self) -> dict:
        if not self._manifest_path.exists():
            return {}
        with self._manifest_path.open("r", encoding="utf-8") as f:
            return json.load(f)

    def _write_manifest(self, data: dict) -> None:
        tmp = self._manifest_path.with_suffix(".tmp")
        with tmp.open("w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, default=str)
        tmp.replace(self._manifest_path)

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

    def create(self, dl_doc: DoclingDocument, source_pdf_path: str, page_count: Optional[int] = None) -> str:
        """Copy the source PDF into the store and write a pending manifest entry.

        Args:
            dl_doc: Converted DoclingDocument (must have a valid ``origin``).
            source_pdf_path: Path to the original PDF to copy into the store.
            page_count: Total page count to record. When omitted the page count
                is derived from ``dl_doc.pages``, which may be a subset when
                the document was converted in chunks.

        Returns:
            The ``binary_hash`` used as the storage key.
        """
        file_hash = str(dl_doc.origin.binary_hash)

        # Ensure per-document directory exists
        doc_dir = self.base_path / "documents" / file_hash
        doc_dir.mkdir(parents=True, exist_ok=True)

        dest = self._pdf_path(file_hash)
        if not dest.exists():
            shutil.copy2(source_pdf_path, dest)

        filename = dl_doc.origin.filename or ""
        record = DocumentRecord(
            file_hash=file_hash,
            document_name=dl_doc.name,
            filename=filename,
            file_extension=Path(filename).suffix if filename else "",
            page_count=page_count if page_count is not None else len(dl_doc.pages),
            ingested_at=datetime.now(timezone.utc),
            status="pending",
        )
        manifest = self._read_manifest()
        manifest[file_hash] = json.loads(record.model_dump_json())
        self._write_manifest(manifest)

        return file_hash

    def set_status(self, file_hash: str, status: Literal["pending", "complete", "failed"]) -> None:
        """Update the ingestion status of a document in the manifest."""
        manifest = self._read_manifest()
        if file_hash in manifest:
            manifest[file_hash]["status"] = status
            self._write_manifest(manifest)

    def update_chunk_count(self, file_hash: str, count: int) -> None:
        """Cache the number of embedded chunks for a document in the manifest.

        Args:
            file_hash: Target document hash.
            count: Total number of chunks embedded in the vector store.
        """
        manifest = self._read_manifest()
        if file_hash in manifest:
            manifest[file_hash]["chunk_count"] = count
            self._write_manifest(manifest)

    def get_pdf_path(self, file_hash: str) -> Path:
        """Return the path to the stored PDF.

        Raises:
            FileNotFoundError: If no PDF is stored for this hash.
        """
        path = self._pdf_path(file_hash)
        if not path.exists():
            raise FileNotFoundError(f"No PDF found for hash: {file_hash}")
        return path

    def exists(self, file_hash: str) -> bool:
        """Return True only if the document has been fully ingested (status=complete).

        Old manifest entries without a status field are treated as complete for
        backward compatibility.
        """
        entry = self._read_manifest().get(file_hash)
        return entry is not None and entry.get("status", "complete") == "complete"

    def delete(self, file_hash: str) -> None:
        """Delete the document directory (PDF + figures + meta) and manifest entry. No-op if not found."""
        doc_dir = self.base_path / "documents" / file_hash
        if doc_dir.exists():
            shutil.rmtree(doc_dir)
        manifest = self._read_manifest()
        manifest.pop(file_hash, None)
        self._write_manifest(manifest)

    def list(self) -> List[DocumentRecord]:
        """Return a DocumentRecord for every stored document."""
        manifest = self._read_manifest()
        records = []
        for v in manifest.values():
            # Backward compatibility: strip heavy fields that may exist in old manifests
            v.pop("figures", None)
            v.pop("pdf_path", None)
            # Ensure new fields have defaults
            v.setdefault("tags", {})
            v.setdefault("figure_count", 0)
            v.setdefault("chunk_count", 0)
            records.append(DocumentRecord(**v))
        return records

    def list_incomplete(self) -> List[DocumentRecord]:
        """Return all manifest entries that are not status=complete."""
        manifest = self._read_manifest()
        result = []
        for v in manifest.values():
            if v.get("status", "complete") != "complete":
                v.pop("figures", None)
                v.pop("pdf_path", None)
                v.setdefault("tags", {})
                v.setdefault("figure_count", 0)
                v.setdefault("chunk_count", 0)
                result.append(DocumentRecord(**v))
        return result

    def rename(self, file_hash: str, new_name: str) -> None:
        """Update the document_name field of a manifest entry."""
        manifest = self._read_manifest()
        if file_hash not in manifest:
            raise KeyError(f"No manifest entry for hash: {file_hash}")
        manifest[file_hash]["document_name"] = new_name
        self._write_manifest(manifest)

    def update_tags(self, file_hash: str, tags: dict) -> None:
        """Merge ``tags`` into the document's existing tag set and persist.

        Args:
            file_hash: Target document hash.
            tags: Dict of ``key -> value`` tags to apply (merged, not replaced).

        Raises:
            KeyError: If no manifest entry exists for ``file_hash``.
        """
        manifest = self._read_manifest()
        if file_hash not in manifest:
            raise KeyError(f"No manifest entry for hash: {file_hash}")
        existing = manifest[file_hash].get("tags", {})
        existing.update(tags)
        manifest[file_hash]["tags"] = existing
        self._write_manifest(manifest)

    def remove_tags(self, file_hash: str, keys: List[str]) -> None:
        """Remove specified tag keys from a document's tag set.

        Args:
            file_hash: Target document hash.
            keys: List of tag keys to remove. Missing keys are silently ignored.

        Raises:
            KeyError: If no manifest entry exists for ``file_hash``.
        """
        manifest = self._read_manifest()
        if file_hash not in manifest:
            raise KeyError(f"No manifest entry for hash: {file_hash}")
        existing = manifest[file_hash].get("tags", {})
        for k in keys:
            existing.pop(k, None)
        manifest[file_hash]["tags"] = existing
        self._write_manifest(manifest)

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

    def save_figures(self, file_hash: str, figures: List[FigureRecord], images: dict = None) -> None:
        """Save figure metadata to per-doc meta.json and optional PNG images.

        Figure PNGs are named by kind and sequence index: ``pic_000.png``,
        ``tbl_000.png``, etc.  The figure metadata list is stored in
        ``documents/<hash>/meta.json`` (not embedded in the manifest) to keep
        the manifest lightweight.

        The manifest entry is updated with ``has_figures=True`` and
        ``figure_count=N`` only.

        Args:
            file_hash: Document hash.
            figures: List of FigureRecord objects.
            images: Optional dict mapping figure ID to PIL.Image. If provided,
                    images are saved as PNGs and image_path is updated on each record.
        """
        fig_dir = self._figures_dir(file_hash)
        fig_dir.mkdir(parents=True, exist_ok=True)

        # Count by kind for short sequential names
        kind_counters: dict[str, int] = {}

        if images:
            for fig in figures:
                if fig.id in images:
                    img = images[fig.id]
                    kind_prefix = "pic" if fig.kind == "picture" else "tbl"
                    idx = kind_counters.get(kind_prefix, 0)
                    filename = f"{kind_prefix}_{idx:03d}.png"
                    kind_counters[kind_prefix] = idx + 1
                    img_path = fig_dir / filename
                    img.save(str(img_path), "PNG")
                    fig.image_path = filename

        # Write figure metadata to per-doc meta.json (not in manifest)
        meta_path = self._meta_path(file_hash)
        meta_path.parent.mkdir(parents=True, exist_ok=True)
        meta = {
            "file_hash": file_hash,
            "figures": [json.loads(fig.model_dump_json()) for fig in figures],
        }
        with meta_path.open("w", encoding="utf-8") as f:
            json.dump(meta, f, indent=2)

        # Update manifest with lightweight summary only
        doc_manifest = self._read_manifest()
        if file_hash in doc_manifest:
            doc_manifest[file_hash]["has_figures"] = True
            doc_manifest[file_hash]["figure_count"] = len(figures)
            self._write_manifest(doc_manifest)

    def get_figures(self, file_hash: str) -> List[FigureRecord]:
        """Load all figure records for a document from per-doc meta.json.

        Falls back to manifest (for backward compatibility with old data that
        embedded figures there) if meta.json doesn't exist.

        Returns:
            List of FigureRecord, or empty list if no figures were extracted.
        """
        meta_path = self._meta_path(file_hash)
        if meta_path.exists():
            with meta_path.open("r", encoding="utf-8") as f:
                meta = json.load(f)
            figures_data = meta.get("figures", [])
            return [FigureRecord(**d) for d in figures_data]

        # Backward compatibility: read from manifest if meta.json not present
        manifest = self._read_manifest()
        entry = manifest.get(file_hash)
        if not entry:
            return []
        figures_data = entry.get("figures", [])
        return [FigureRecord(**d) for d in figures_data]

    def get_figure_image_path(self, file_hash: str, figure_id: str) -> Path:
        """Return absolute path to a specific figure PNG.

        The figure_id is matched by looking up the record and using its
        ``image_path`` (short filename like ``pic_000.png``).

        Raises:
            FileNotFoundError: If the image doesn't exist.
        """
        figures = self.get_figures(file_hash)
        for fig in figures:
            if fig.id == figure_id:
                if fig.image_path:
                    img_path = self._figures_dir(file_hash) / fig.image_path
                    if img_path.exists():
                        return img_path
                raise FileNotFoundError(f"No figure image for: {figure_id} (doc {file_hash})")
        raise FileNotFoundError(f"No figure record found: {figure_id} for doc {file_hash}")

    def delete_figures(self, file_hash: str) -> None:
        """Remove all figure data for a document (images only; manifest entry deleted by delete())."""
        fig_dir = self._figures_dir(file_hash)
        if fig_dir.exists():
            shutil.rmtree(fig_dir)