app/services/citation_exploration_store.py

from __future__ import annotations

import json
import os
import tempfile
import threading
import time
from typing import Any, Literal, Mapping

from pydantic import BaseModel, ConfigDict, Field

from app.services.citation_graph import CitationGraphService


CACHE_VERSION = 2
_DEFAULT_ROOT = os.path.expanduser("~/.studybuddy/citation_exploration")
_RICH_METADATA_TTL = 7 * 86400
_SPARSE_METADATA_TTL = 6 * 3600
_TOPOLOGY_TTL = 30 * 86400
_ERROR_TTL = 15 * 60


class CachedReferenceSet(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)

    provider: str
    work_id: str
    referenced_work_ids: list[str] = Field(default_factory=list)
    updated_at: float
    expires_at: float
    error: str = ""

    @property
    def fresh(self) -> bool:
        return self.expires_at > time.time()


class CachedNeighborhood(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)

    key: str
    exploration_kind: Literal["paper", "author"]
    seed_candidate_id: str = ""
    author_key: str = ""
    paper_ids: list[str] = Field(default_factory=list)
    primary_edges: list[dict[str, Any]] = Field(default_factory=list)
    induced_edges: list[dict[str, Any]] = Field(default_factory=list)
    primary_state: Literal["complete", "degraded"] = "complete"
    induced_state: Literal["pending", "complete", "degraded"] = "pending"
    warnings: list[str] = Field(default_factory=list)
    updated_at: float
    expires_at: float

    @property
    def fresh(self) -> bool:
        return self.expires_at > time.time()


class ExplorationCache(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)

    cache_version: int = CACHE_VERSION
    papers: dict[str, dict[str, Any]] = Field(default_factory=dict)
    neighborhoods: dict[str, CachedNeighborhood] = Field(default_factory=dict)
    reference_sets: dict[str, CachedReferenceSet] = Field(default_factory=dict)
    edges: dict[str, dict[str, Any]] = Field(default_factory=dict)
    updated_at: float = 0.0


class CitationExplorationStore:
    """Atomic, project-scoped cache for external paper graph exploration."""

    _locks: dict[str, threading.RLock] = {}
    _locks_guard = threading.Lock()

    def __init__(self, project_id: str, root: str | None = None) -> None:
        if not CitationGraphService.is_valid_project_id(project_id):
            raise ValueError("Invalid project id.")
        self.project_id = project_id
        self.root = root or _DEFAULT_ROOT
        os.makedirs(self.root, exist_ok=True)
        with self._locks_guard:
            self._lock = self._locks.setdefault(self._path(), threading.RLock())

    def _path(self) -> str:
        return os.path.join(self.root, f"{self.project_id}.json")

    def load(self) -> ExplorationCache:
        with self._lock:
            try:
                with open(self._path(), encoding="utf-8") as file:
                    payload = json.load(file)
                if int(payload.get("cache_version") or 0) != CACHE_VERSION:
                    return ExplorationCache()
                return ExplorationCache(**payload)
            except (FileNotFoundError, OSError, ValueError, TypeError, json.JSONDecodeError):
                return ExplorationCache()

    def save(self, cache: ExplorationCache) -> None:
        cache.updated_at = time.time()
        with self._lock:
            descriptor, temporary = tempfile.mkstemp(
                prefix=f".{self.project_id}.", suffix=".tmp", dir=self.root, text=True
            )
            try:
                with os.fdopen(descriptor, "w", encoding="utf-8") as file:
                    json.dump(cache.model_dump(), file, indent=2)
                    file.flush()
                    os.fsync(file.fileno())
                os.replace(temporary, self._path())
            finally:
                if os.path.exists(temporary):
                    os.remove(temporary)

    @staticmethod
    def paper_is_rich(paper: Mapping[str, Any]) -> bool:
        authors = paper.get("author_details") or []
        destinations = paper.get("destinations") or []
        abstract_known = bool(paper.get("abstract")) or paper.get("abstract_status") == "unavailable"
        return bool(paper.get("title") and authors and destinations and abstract_known)

    @classmethod
    def paper_fresh(cls, paper: Mapping[str, Any]) -> bool:
        try:
            return float(paper.get("cache_expires_at") or 0) > time.time()
        except (TypeError, ValueError):
            return False

    @classmethod
    def prepare_paper(cls, paper: Mapping[str, Any], *, error: str = "") -> dict[str, Any]:
        payload = dict(paper)
        now = time.time()
        ttl = _ERROR_TTL if error else _RICH_METADATA_TTL if cls.paper_is_rich(payload) else _SPARSE_METADATA_TTL
        payload["cache_updated_at"] = now
        payload["cache_expires_at"] = now + ttl
        payload["metadata_quality"] = "rich" if cls.paper_is_rich(payload) else "sparse"
        if error:
            payload["last_provider_error"] = error
        return payload

    def candidate(self, candidate_id: str) -> dict[str, Any] | None:
        paper = self.load().papers.get(candidate_id)
        return dict(paper) if paper else None

    def upsert_papers(self, papers: list[Mapping[str, Any]]) -> ExplorationCache:
        cache = self.load()
        for incoming in papers:
            candidate_id = str(incoming.get("candidate_id") or "")
            if not candidate_id:
                continue
            prepared = self.prepare_paper(incoming, error=str(incoming.get("error") or ""))
            existing = cache.papers.get(candidate_id)
            if existing and self.paper_is_rich(existing) and not self.paper_is_rich(prepared):
                merged = dict(prepared)
                for key, value in existing.items():
                    if value not in (None, "", [], {}):
                        merged[key] = value
                prepared = merged
            cache.papers[candidate_id] = prepared
        self.save(cache)
        return cache

    def neighborhood(self, key: str, *, fresh_only: bool = True) -> CachedNeighborhood | None:
        row = self.load().neighborhoods.get(key)
        return row if row and (row.fresh or not fresh_only) else None

    def save_neighborhood(self, neighborhood: CachedNeighborhood) -> ExplorationCache:
        cache = self.load()
        cache.neighborhoods[neighborhood.key] = neighborhood
        for edge in [*neighborhood.primary_edges, *neighborhood.induced_edges]:
            edge_id = str(edge.get("id") or "")
            if edge_id:
                cache.edges[edge_id] = dict(edge)
        self.save(cache)
        return cache

    @staticmethod
    def new_neighborhood(
        *,
        key: str,
        exploration_kind: Literal["paper", "author"],
        paper_ids: list[str],
        primary_edges: list[dict[str, Any]],
        seed_candidate_id: str = "",
        author_key: str = "",
        warnings: list[str] | None = None,
        degraded: bool = False,
    ) -> CachedNeighborhood:
        now = time.time()
        return CachedNeighborhood(
            key=key,
            exploration_kind=exploration_kind,
            seed_candidate_id=seed_candidate_id,
            author_key=author_key,
            paper_ids=paper_ids,
            primary_edges=primary_edges,
            primary_state="degraded" if degraded else "complete",
            warnings=warnings or [],
            updated_at=now,
            expires_at=now + (_ERROR_TTL if degraded else _TOPOLOGY_TTL),
        )

    def reference_set(self, candidate_id: str, *, fresh_only: bool = True) -> CachedReferenceSet | None:
        row = self.load().reference_sets.get(candidate_id)
        return row if row and (row.fresh or not fresh_only) else None

    def save_reference_sets(self, rows: Mapping[str, CachedReferenceSet]) -> ExplorationCache:
        cache = self.load()
        cache.reference_sets.update(rows)
        self.save(cache)
        return cache

    @staticmethod
    def new_reference_set(
        *, provider: str, work_id: str, referenced_work_ids: list[str], error: str = ""
    ) -> CachedReferenceSet:
        now = time.time()
        return CachedReferenceSet(
            provider=provider,
            work_id=work_id,
            referenced_work_ids=sorted(set(referenced_work_ids)),
            updated_at=now,
            expires_at=now + (_ERROR_TTL if error else _TOPOLOGY_TTL),
            error=error,
        )

    def merge_induced_edges(
        self, neighborhood_key: str, edges: list[dict[str, Any]], *, warnings: list[str] | None = None
    ) -> CachedNeighborhood:
        cache = self.load()
        neighborhood = cache.neighborhoods.get(neighborhood_key)
        if neighborhood is None:
            raise LookupError("Exploration neighborhood not found.")
        by_id = {str(edge.get("id") or ""): dict(edge) for edge in neighborhood.induced_edges if edge.get("id")}
        for edge in edges:
            edge_id = str(edge.get("id") or "")
            if edge_id:
                by_id[edge_id] = dict(edge)
                cache.edges[edge_id] = dict(edge)
        neighborhood.induced_edges = list(by_id.values())
        neighborhood.induced_state = "degraded" if warnings else "complete"
        neighborhood.warnings = list(dict.fromkeys([*neighborhood.warnings, *(warnings or [])]))
        neighborhood.updated_at = time.time()
        neighborhood.expires_at = time.time() + (_ERROR_TTL if warnings else _TOPOLOGY_TTL)
        cache.neighborhoods[neighborhood_key] = neighborhood
        self.save(cache)
        return neighborhood