src/data_manager.py

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

This module encapsulates the logic for computing the heavy data
transformations in ``pipeline.py`` and persisting the results to disk.
It adds a small amount of resilience around caching and uses
``logging`` instead of printing directly to stdout.  The cache files
include 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 typing import Dict
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.
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()

# Build cache file paths with version tags.  This allows caches from
# different versions of the pipeline to coexist without overwriting
# each other.  For example, ``daioe_weighted_v1.csv``.
WEIGHTED_CACHE: Path = DATA_DIR / f"daioe_weighted_{CACHE_VERSION}.csv"
SIMPLE_CACHE: Path = DATA_DIR / f"daioe_simple_{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() -> Dict[str, pd.DataFrame]:
    """Runs the heavy pipeline calculation."""
    return pipeline.run_pipeline()


def load_payload(force_recompute: bool = False) -> Dict[str, pd.DataFrame]:
    """
    Load data from disk cache if available, otherwise compute and save.

    Parameters
    ----------
    force_recompute : bool, optional
        If ``True``, recompute the pipeline even if cache files exist.

    Returns
    -------
    Dict[str, pd.DataFrame]
        A dictionary with keys ``"weighted"`` and ``"simple"``
        containing the respective DataFrames.
    """
    # If a cached payload exists and recomputation is not forced, return it
    if not force_recompute and WEIGHTED_CACHE.exists() and SIMPLE_CACHE.exists():
        logger.info("Loading pipeline output from cache directory %s", DATA_DIR)
        try:
            weighted_df = pd.read_csv(WEIGHTED_CACHE)
            simple_df = pd.read_csv(SIMPLE_CACHE)
            return {"weighted": weighted_df, "simple": simple_df}
        except Exception as exc:
            # If reading the cache fails, fall back to recomputing
            logger.warning(
                "Error reading cache files %s and %s: %s; falling back to recompute",
                WEIGHTED_CACHE,
                SIMPLE_CACHE,
                exc,
            )

    if force_recompute:
        # Clear the LRU cache before recomputing
        _compute_pipeline_payload.cache_clear()

    # Cache miss or forced recompute: run the heavy pipeline
    logger.info("Computing pipeline data – this may take a while…")
    payload = _compute_pipeline_payload()

    # Persist to disk atomically
    try:
        _atomic_to_csv(payload["weighted"], WEIGHTED_CACHE)
        _atomic_to_csv(payload["simple"], SIMPLE_CACHE)
        logger.info(
            "Cache updated: weighted=%s, simple=%s", WEIGHTED_CACHE.name, SIMPLE_CACHE.name
        )
    except Exception as exc:
        logger.warning("Could not write cache files: %s", exc)

    return payload