src/data_manager.py

"""Data manager for loading and caching SCB employment pipeline results.

This module encapsulates the logic for computing the SCB-only
transformations in ``pipeline.py`` and persisting the result to disk.
It adds a small amount of resilience around caching and uses
``logging`` instead of printing directly to stdout.  The cache file
includes a version tag to make it easy to invalidate caches when
fundamental changes are made to the pipeline logic.
"""

import os
import tempfile
import logging
from pathlib import Path
from functools import lru_cache

import pandas as pd

from . import pipeline

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Cache setup
# ---------------------------------------------------------------------------
# A version tag to embed into the cache filenames.  Bump this value
# whenever the underlying ``pipeline`` logic changes in a way that
# invalidates existing caches (or when you want to force a fresh recompute).
CACHE_VERSION: str = "v1"


def _resolve_cache_dir() -> Path:
    """Select a writable directory for caching.

    The lookup order is:

    1. The ``DATA_CACHE_DIR`` environment variable, if set.
    2. A ``data`` folder at the repository root.
    3. A temporary directory in ``/tmp``.

    Each candidate path is tested for writability by attempting to
    create and delete a sentinel file.  The first path that succeeds
    is returned.  If none succeed, a final fallback directory in ``/tmp``
    is created and returned.
    """
    candidates: list[Path] = []
    env = os.getenv("DATA_CACHE_DIR")
    if env:
        # Expand relative or user paths to absolute
        candidates.append(Path(env).expanduser().resolve())

    # Repo root /data (two levels up from this file)
    candidates.append(Path(__file__).resolve().parent.parent / "data")
    # Temp fallback
    candidates.append(Path(tempfile.gettempdir()) / "employment_ai_cache")

    for path in candidates:
        try:
            path.mkdir(parents=True, exist_ok=True)
            test_file = path / ".write_test"
            test_file.write_text("ok", encoding="utf-8")
            test_file.unlink()
            return path
        except Exception:
            continue

    # Final fallback: ensure the last candidate exists
    fallback = Path(tempfile.gettempdir()) / "employment_ai_cache"
    fallback.mkdir(parents=True, exist_ok=True)
    return fallback


# Resolve the directory once at import time
DATA_DIR: Path = _resolve_cache_dir()

# Single cache file for the SCB-only output DataFrame.
SCB_CACHE: Path = DATA_DIR / f"scb_employment_{CACHE_VERSION}.csv"


def _atomic_to_csv(df: pd.DataFrame, path: Path) -> None:
    """Write a DataFrame to CSV atomically.

    The CSV is first written to a temporary file in the same directory
    and then renamed to the final location.  This avoids leaving a
    partially written file if the process is interrupted mid‑write.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    df.to_csv(tmp_path, index=False)
    tmp_path.replace(path)


@lru_cache(maxsize=1)
def _compute_pipeline_payload() -> pd.DataFrame:
    """Run the SCB-only pipeline calculation (cached in-process)."""
    return pipeline.run_pipeline()


def load_payload() -> pd.DataFrame:
    """Load employment data from disk cache if available, otherwise compute and save.

    To force a rebuild, delete the cache file or bump `CACHE_VERSION` and restart.
    """
    if SCB_CACHE.exists():
        logger.info("Loading pipeline output from cache directory %s", DATA_DIR)
        try:
            return pd.read_csv(SCB_CACHE)
        except Exception as exc:
            # If reading the cache fails, fall back to recomputing
            logger.warning(
                "Error reading cache file %s: %s; falling back to recompute",
                SCB_CACHE,
                exc,
            )

    logger.info("Computing SCB employment data – this may take a while…")
    payload = _compute_pipeline_payload()

    # Persist to disk atomically
    try:
        _atomic_to_csv(payload, SCB_CACHE)
        logger.info("Cache updated: %s", SCB_CACHE.name)
    except Exception as exc:
        logger.warning("Could not write cache file: %s", exc)

    return payload