backend/hf_sync.py

"""Hugging Face Dataset acts as the persistent 'cloud drive' for uploaded PDFs.

Layout inside the dataset repo:
    pdfs/<filename>.pdf          <- raw PDFs (source of truth)
    vector_cache/manifest.json   <- vector cache metadata
    vector_cache/index/**        <- persisted LanceDB files

Local disk on HF Spaces is ephemeral, so we try to restore the vector cache
from the dataset first and only rebuild from PDFs when the cache is missing or
stale.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path
import shutil
from tempfile import TemporaryDirectory
from typing import List

from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

from config import DATASET_ID, HF_TOKEN

logger = logging.getLogger(__name__)

PDF_PREFIX = "pdfs/"
VECTOR_PREFIX = "vector_cache/"
VECTOR_INDEX_PREFIX = f"{VECTOR_PREFIX}index/"
VECTOR_MANIFEST_PATH = f"{VECTOR_PREFIX}manifest.json"
VECTOR_CACHE_SCHEMA_VERSION = 1


def is_configured() -> bool:
    return bool(HF_TOKEN and DATASET_ID)


def _vector_manifest(pdf_filenames: List[str], embed_model: str) -> dict:
    return {
        "schema_version": VECTOR_CACHE_SCHEMA_VERSION,
        "embed_model": embed_model,
        "pdf_files": sorted(pdf_filenames),
    }


def _api() -> HfApi:
    if not HF_TOKEN or not DATASET_ID:
        raise RuntimeError(
            "HF_TOKEN and DATASET_ID must be set as environment variables / Space secrets."
        )
    return HfApi(token=HF_TOKEN)


def _ensure_dataset_exists(api: HfApi) -> None:
    try:
        api.repo_info(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        try:
            api.create_repo(
                repo_id=DATASET_ID,
                repo_type="dataset",
                private=True,
                exist_ok=True,
            )
            logger.info("Created dataset repo %s", DATASET_ID)
        except Exception as exc:  # noqa: BLE001
            raise RuntimeError(
                f"Dataset {DATASET_ID!r} does not exist and could not be created automatically. "
                "Create it on Hugging Face or update DATASET_ID / HF_TOKEN permissions."
            ) from exc


def list_remote_pdfs() -> List[str]:
    """Return PDF filenames (basename only) currently stored in the dataset."""
    api = _api()
    try:
        files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        return []
    return [
        f[len(PDF_PREFIX):]
        for f in files
        if f.startswith(PDF_PREFIX) and f.lower().endswith(".pdf")
    ]


def download_vector_store(
    target_dir: Path,
    expected_pdf_filenames: List[str],
    embed_model: str,
) -> bool:
    """Restore the persisted LanceDB folder when its manifest matches the dataset PDFs."""
    if not is_configured():
        return False

    api = _api()
    try:
        manifest_path = hf_hub_download(
            repo_id=DATASET_ID,
            repo_type="dataset",
            filename=VECTOR_MANIFEST_PATH,
            token=HF_TOKEN,
        )
    except (EntryNotFoundError, RepositoryNotFoundError):
        logger.info("No vector cache manifest in dataset yet")
        return False
    except Exception as exc:  # noqa: BLE001
        logger.warning("Could not download vector cache manifest: %s", exc)
        return False

    try:
        manifest = json.loads(Path(manifest_path).read_text("utf-8"))
    except Exception as exc:  # noqa: BLE001
        logger.warning("Could not parse vector cache manifest: %s", exc)
        return False

    expected_manifest = _vector_manifest(expected_pdf_filenames, embed_model)
    if manifest != expected_manifest:
        logger.info("Vector cache manifest is stale; rebuilding from PDFs")
        return False

    try:
        files = api.list_repo_files(repo_id=DATASET_ID, repo_type="dataset")
    except RepositoryNotFoundError:
        return False

    vector_files = [f for f in files if f.startswith(VECTOR_INDEX_PREFIX)]
    if not vector_files:
        logger.info("Vector cache manifest exists but no cache files were found")
        return False

    shutil.rmtree(target_dir, ignore_errors=True)
    target_dir.mkdir(parents=True, exist_ok=True)

    for repo_path in vector_files:
        relative_path = repo_path[len(VECTOR_INDEX_PREFIX):]
        if not relative_path:
            continue
        downloaded = hf_hub_download(
            repo_id=DATASET_ID,
            repo_type="dataset",
            filename=repo_path,
            token=HF_TOKEN,
        )
        destination = target_dir / relative_path
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(downloaded, destination)

    logger.info("Vector cache restored from dataset (%d files)", len(vector_files))
    return True


def delete_vector_store() -> bool:
    """Remove the persisted vector cache folder from the dataset."""
    if not is_configured():
        return False

    api = _api()
    try:
        api.delete_folder(
            path_in_repo=VECTOR_PREFIX.rstrip("/"),
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message="Delete vector cache",
        )
        logger.info("Vector cache deleted from HF Dataset")
        return True
    except (EntryNotFoundError, RepositoryNotFoundError):
        return False


def upload_vector_store(local_dir: Path, pdf_filenames: List[str], embed_model: str) -> bool:
    """Persist the local LanceDB folder into the dataset along with a manifest."""
    if not is_configured():
        return False

    api = _api()
    _ensure_dataset_exists(api)

    local_files = [p for p in local_dir.rglob("*") if p.is_file()]
    if not pdf_filenames or not local_files:
        delete_vector_store()
        return bool(not pdf_filenames)

    with TemporaryDirectory(prefix="iam-earth-vector-cache-") as tmp_dir:
        stage_dir = Path(tmp_dir)
        index_dir = stage_dir / "index"
        shutil.copytree(local_dir, index_dir, dirs_exist_ok=True)
        (stage_dir / "manifest.json").write_text(
            json.dumps(_vector_manifest(pdf_filenames, embed_model), indent=2, ensure_ascii=False),
            "utf-8",
        )

        api.upload_folder(
            repo_id=DATASET_ID,
            repo_type="dataset",
            folder_path=stage_dir,
            path_in_repo=VECTOR_PREFIX.rstrip("/"),
            commit_message="Sync vector cache",
            delete_patterns="**",
            ignore_patterns=["**/.DS_Store"],
        )

    logger.info("Vector cache synced to HF Dataset")
    return True


def download_all_pdfs(target_dir: Path) -> List[Path]:
    """Download every PDF in the dataset into target_dir. Returns local paths."""
    target_dir.mkdir(parents=True, exist_ok=True)
    paths: List[Path] = []
    for name in list_remote_pdfs():
        try:
            local = hf_hub_download(
                repo_id=DATASET_ID,
                repo_type="dataset",
                filename=f"{PDF_PREFIX}{name}",
                local_dir=str(target_dir),
                token=HF_TOKEN,
            )
            paths.append(Path(local))
        except EntryNotFoundError:
            continue
    return paths


def upload_pdf(local_path: Path, filename: str) -> None:
    """Push a single PDF into the dataset under pdfs/<filename>."""
    api = _api()
    try:
        _ensure_dataset_exists(api)
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=f"{PDF_PREFIX}{filename}",
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message=f"Add {filename}",
        )
    except RuntimeError:
        raise
    except Exception as exc:  # noqa: BLE001
        raise RuntimeError(
            f"Failed to upload {filename} to dataset {DATASET_ID!r}: {exc}"
        ) from exc


def delete_pdf(filename: str) -> bool:
    """Delete pdfs/<filename> from the dataset. Returns True if removed."""
    api = _api()
    try:
        api.delete_file(
            path_in_repo=f"{PDF_PREFIX}{filename}",
            repo_id=DATASET_ID,
            repo_type="dataset",
            commit_message=f"Delete {filename}",
        )
        return True
    except EntryNotFoundError:
        return False