app/api/deps.py

"""Shared dependencies for API routes."""

import json
import logging
import pickle
import time
from functools import lru_cache
from pathlib import Path
from typing import Any
from uuid import uuid4

from app.config import settings
from app.core.ai_agent import MappingAgent
from app.models.template import TemplateSchema

logger = logging.getLogger(__name__)

# In-memory session store (swap for Redis/DynamoDB in production)
_sessions: dict[str, dict[str, Any]] = {}

SESSION_TTL_HOURS = settings.session_ttl_hours


def get_session_store() -> dict[str, dict[str, Any]]:
    return _sessions


# Keys whose values are heavy and should be serialized to disk
_HEAVY_KEYS = {"all_metadata", "all_sample_values"}


def _session_cache_path(session_id: str, key: str) -> Path:
    return settings.upload_dir / f".session_{session_id}_{key}.pkl"


def create_session() -> str:
    session_id = str(uuid4())
    _sessions[session_id] = {"status": "created", "created_at": time.time()}
    return session_id


def persist_session_data(session_id: str, key: str, value: Any) -> None:
    """Store a heavy value on disk and keep only a sentinel in memory."""
    settings.ensure_dirs()
    cache_path = _session_cache_path(session_id, key)
    try:
        with open(cache_path, "wb") as f:
            pickle.dump(value, f, protocol=pickle.HIGHEST_PROTOCOL)
        # Store a lightweight sentinel so get_session knows to reload
        session = _sessions.get(session_id)
        if session is not None:
            session[key] = None  # sentinel
            session[f"_{key}_on_disk"] = True
    except Exception as exc:
        logger.warning("Failed to persist session data %s/%s: %s", session_id, key, exc)
        # Fall back to keeping it in memory
        session = _sessions.get(session_id)
        if session is not None:
            session[key] = value


def _load_heavy(session_id: str, key: str) -> Any | None:
    """Load a heavy value back from disk."""
    cache_path = _session_cache_path(session_id, key)
    if cache_path.exists():
        try:
            with open(cache_path, "rb") as f:
                return pickle.load(f)  # noqa: S301
        except Exception as exc:
            logger.warning("Failed to load session data %s/%s: %s", session_id, key, exc)
    return None


def get_session(session_id: str) -> dict[str, Any] | None:
    session = _sessions.get(session_id)
    if session:
        session["last_accessed"] = time.time()
        # Lazily reload heavy keys from disk
        for key in _HEAVY_KEYS:
            if session.get(f"_{key}_on_disk") and session.get(key) is None:
                loaded = _load_heavy(session_id, key)
                if loaded is not None:
                    session[key] = loaded
    return session


def cleanup_expired_sessions() -> int:
    """Remove sessions older than SESSION_TTL_HOURS and delete their files.
    Returns number of sessions cleaned up."""
    now = time.time()
    cutoff = now - SESSION_TTL_HOURS * 3600
    expired = [
        sid for sid, s in _sessions.items()
        if s.get("last_accessed", s.get("created_at", 0)) < cutoff
    ]
    for sid in expired:
        session = _sessions.pop(sid, {})
        # Clean up uploaded files
        for path in session.get("source_paths", {}).values():
            try:
                Path(path).unlink(missing_ok=True)
            except Exception:
                pass
        # Clean up reference files
        for path in session.get("reference_paths", {}).values():
            try:
                Path(path).unlink(missing_ok=True)
            except Exception:
                pass
        # Clean up output files
        for path in session.get("output_files", []):
            try:
                Path(path).unlink(missing_ok=True)
            except Exception:
                pass
        # Clean up serialized session data
        for key in _HEAVY_KEYS:
            try:
                _session_cache_path(sid, key).unlink(missing_ok=True)
            except Exception:
                pass
    if expired:
        logger.info("Cleaned up %d expired sessions", len(expired))
    return len(expired)


@lru_cache
def get_mapping_agent() -> MappingAgent:
    """Default agent — for metadata-only requests (no sensitive content)."""
    return MappingAgent(
        model=settings.llm_model,
        api_key=settings.llm_api_key or None,
        base_url=settings.llm_base_url,
    )


@lru_cache
def get_secure_mapping_agent() -> MappingAgent:
    """Secure agent — for requests with reference files (may contain sensitive content).
    Falls back to default agent if not configured."""
    if settings.has_secure_llm:
        return MappingAgent(
            model=settings.secure_llm_model,
            api_key=settings.secure_llm_api_key or None,
            base_url=settings.secure_llm_base_url,
        )
    return get_mapping_agent()


@lru_cache
def load_template(template_name: str = "stars_v1") -> TemplateSchema:
    template_dir = Path(__file__).parent.parent.parent / "templates"
    template_files = {
        "stars_v1": "stars_v1.json",
        "qvey_v1": "qvey_v1.json",
        # backward compat
        "aseesa_standard_v1": "stars_v1.json",
    }
    filename = template_files.get(template_name)
    if not filename:
        raise ValueError(f"Unknown template: {template_name}")

    template_path = template_dir / filename
    with open(template_path) as f:
        data = json.load(f)
    return TemplateSchema(**data)