src/storage.py

from __future__ import annotations

import logging
import os
import sys
import tempfile
from pathlib import Path
from typing import Any

from PIL import Image

logger = logging.getLogger(__name__)

_DEFAULT_STORAGE_PATH = "granite-cache"

_SUBDIRS = ("tmp", "sessions", "cache", "uploads", "models")

_MODEL_CACHE_ENV_VARS = {
    "HF_HOME": "models/huggingface",
    "DOCLING_CACHE_DIR": "models/docling",
    "TORCH_HOME": "models/torch",
    "EASYOCR_MODULE_PATH": "models/easyocr",
}


def get_storage_path(sub_path: str = "") -> Path:
    base = os.environ.get("STORAGE_PATH", "").strip()
    root = Path(base).resolve() if base else Path(tempfile.gettempdir())
    return root / sub_path if sub_path else root


def get_temp_dir() -> Path:
    return get_storage_path("tmp")


def init_storage() -> None:
    base = os.environ.get("STORAGE_PATH", "").strip()
    if base:
        root = Path(base).resolve()
        logger.info("STORAGE_PATH set: %s", root)
    else:
        root = Path(tempfile.gettempdir())
        logger.info("STORAGE_PATH not set — using system temp dir: %s", root)

    for sub in _SUBDIRS:
        d = root / sub
        try:
            d.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            logger.error("Cannot create storage dir %s: %s", d, e)
            sys.exit(1)

    if base:
        for env_var, sub_path in _MODEL_CACHE_ENV_VARS.items():
            if env_var not in os.environ:
                cache_dir = root / sub_path
                cache_dir.mkdir(parents=True, exist_ok=True)
                os.environ[env_var] = str(cache_dir)
                logger.info("Set %s=%s", env_var, cache_dir)

    probe = root / "tmp" / ".write_test"
    try:
        probe.touch()
        probe.unlink()
    except OSError as e:
        logger.error("Storage path %s is not writable: %s", root, e)
        sys.exit(1)

    logger.info("Storage ready: %s  (subdirs: %s)", root, ", ".join(_SUBDIRS))


def get_sessions_path() -> Path:
    return get_storage_path("sessions")


def parse_cache_per_session() -> bool:
    """Return True when Docling parse results should be stored per-session.

    Controlled by ``PARSE_CACHE_PER_SESSION`` env var (default ``"true"``).
    Only meaningful in disk mode.
    """
    return os.environ.get("PARSE_CACHE_PER_SESSION", "true").strip().lower() != "false"


def use_disk_images() -> bool:
    """Return True when bulky session images should be externalized to disk.

    Controlled by ``SESSION_IMAGES`` env var: ``"disk"`` (default) saves
    PIL Images to disk and stores LazyImage proxies in session state;
    ``"memory"`` keeps real PIL Images in-memory (original behaviour).
    """
    return os.environ.get("SESSION_IMAGES", "disk").strip().lower() != "memory"


class LazyImage:
    """PIL Image proxy that loads from disk on every access.

    Stored in figure dict ``"image"`` slots after upload processing.
    Any PIL attribute access (e.g. ``.convert()``, ``.save()``, ``.size``)
    reads the PNG from disk each time so no PIL data stays in memory.
    """

    __slots__ = ("_path",)

    def __init__(self, path: Path) -> None:
        object.__setattr__(self, "_path", path)

    def _load(self) -> Image.Image:
        path = object.__getattribute__(self, "_path")
        return Image.open(path).copy()

    @property
    def disk_path(self) -> Path:
        return object.__getattribute__(self, "_path")

    def __getattr__(self, name: str) -> Any:
        return getattr(self._load(), name)

    def __repr__(self) -> str:
        path = object.__getattribute__(self, "_path")
        return f"<LazyImage {path.name}>"


def save_image(session_id: str, category: str, index: int, img: Image.Image) -> LazyImage | Image.Image:
    """Save a PIL Image to disk and return a LazyImage proxy, or pass through.

    When ``SESSION_STORAGE=memory``, returns *img* unchanged (original behaviour).
    Otherwise saves to ``STORAGE_PATH/sessions/<session_id>/<category>/`` and
    returns a :class:`LazyImage`.
    """
    if not use_disk_images():
        return img
    cat_dir = get_sessions_path() / session_id / category
    cat_dir.mkdir(parents=True, exist_ok=True)
    path = cat_dir / f"{index:04d}.png"
    img.save(path, format="PNG")
    return LazyImage(path)


def resolve_for_gradio(img: "LazyImage | Image.Image | None") -> "Path | Image.Image | None":
    """Return a value Gradio's gr.Image component can postprocess.

    LazyImage is an internal proxy unknown to Gradio; hand it the disk path
    instead (Gradio accepts str/Path for image files). PIL Images and None
    pass through unchanged.
    """
    if isinstance(img, LazyImage):
        return img.disk_path
    return img


def _parse_cache_dir(file_hash: str, session_id: str | None) -> Path:
    if session_id and parse_cache_per_session():
        return get_sessions_path() / session_id
    return get_storage_path("cache") / file_hash


def load_parse_cache(file_hash: str, session_id: str | None = None) -> dict | None:
    """Load a previously persisted Docling parse result from disk.

    Returns a dict with ``html``, ``text``, and an empty ``figures`` list,
    or ``None`` if no cached result exists or disk mode is off.
    """
    if not use_disk_images():
        return None
    import json
    path = _parse_cache_dir(file_hash, session_id) / "result.json"
    if not path.exists():
        return None
    data = json.loads(path.read_text())
    data["figures"] = []
    return data


def save_parse_cache(file_hash: str, result: dict, session_id: str | None = None) -> None:
    """Persist the html and text from a Docling parse result to disk.

    No-op when disk mode is off. Figures (PIL Images) are excluded.
    """
    if not use_disk_images():
        return
    import json
    cache_dir = _parse_cache_dir(file_hash, session_id)
    cache_dir.mkdir(parents=True, exist_ok=True)
    payload = {"html": result.get("html", ""), "text": result.get("text", "")}
    (cache_dir / "result.json").write_text(json.dumps(payload))


def save_upload(session_id: str, data: bytes, suffix: str) -> Path:
    """Save the original uploaded file into the session directory.

    Returns the path to the saved file.  When ``SESSION_IMAGES=memory``,
    returns ``None`` (caller should fall back to a temp file).
    """
    if not use_disk_images():
        return None
    session_dir = get_sessions_path() / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    path = session_dir / f"upload{suffix}"
    path.write_bytes(data)
    return path


def v1_mode() -> bool:
    """Return True when browser storage mode is enabled (v1 API)."""
    return os.environ.get("BROWSER_MEMORY_MODE", "").strip().lower() in ("1", "true", "yes")


def debug_keep_files() -> bool:
    """Return True when intermediate files should be preserved for debugging."""
    return os.environ.get("DEBUG_KEEP_FILES", "").strip().lower() in ("1", "true", "yes")


def cleanup_session_files(session_id: str) -> None:
    """Remove all files for a session from disk.

    No-op when DEBUG_KEEP_FILES is set or disk mode is off.
    """
    if debug_keep_files() or not use_disk_images():
        return
    import shutil

    session_dir = get_sessions_path() / session_id
    if session_dir.exists():
        shutil.rmtree(session_dir, ignore_errors=True)