app/services/citation_exploration.py

from __future__ import annotations

import hashlib
import time
from collections.abc import Mapping
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from app.services.citation_exploration_store import CachedNeighborhood, CitationExplorationStore
from app.services.citation_graph import CitationGraph, CitationGraphEdge, CitationGraphNode, CitationGraphService
from app.services.citation_topology import CitationTopologyService, VerifiedTopologyEdge
from app.services.paper_acquisition import PaperAcquisitionResult, PaperAcquisitionService, sanitize_acquisition_candidate
from app.services.paper_metadata import PaperMetadataService


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

    graph: CitationGraph
    seed_id: str
    seed_candidate_id: str
    neighborhood_key: str
    exploration_kind: Literal["paper", "author"] = "paper"
    primary_state: Literal["complete", "degraded"] = "complete"
    induced_state: Literal["pending", "complete", "degraded"] = "pending"


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

    neighborhood_key: str
    nodes: list[CitationGraphNode] = Field(default_factory=list)
    edges: list[CitationGraphEdge] = Field(default_factory=list)
    induced_state: Literal["complete", "degraded"]
    warnings: list[str] = Field(default_factory=list)


class CitationExplorationService:
    """Cached transient paper and author graph orchestration."""

    def __init__(
        self,
        root: str | None = None,
        graph_service: CitationGraphService | None = None,
        topology_service: CitationTopologyService | None = None,
        metadata_service: PaperMetadataService | None = None,
        acquisition_service: PaperAcquisitionService | None = None,
    ) -> None:
        self.root = root
        self.graph_service = graph_service or CitationGraphService()
        self.metadata_service = metadata_service or PaperMetadataService()
        self.topology = topology_service or CitationTopologyService(self.metadata_service)
        self.acquisition_service = acquisition_service or PaperAcquisitionService()

    def _store(self, project_id: str) -> CitationExplorationStore:
        return CitationExplorationStore(project_id, self.root)

    def _project_candidate(self, project_id: str, candidate_id: str) -> dict[str, Any] | None:
        graph = self.graph_service.load(project_id)
        for node in graph.nodes:
            if node.status == "owned" and node.metadata.get("external_candidate_id") == candidate_id:
                return self.metadata_service.from_candidate({
                    **node.metadata,
                    "candidate_id": candidate_id,
                    "title": node.title,
                    "authors": node.authors,
                    "year": node.year,
                    "abstract": node.abstract,
                    "url": node.url,
                    "doi": node.doi,
                    "arxiv_id": node.arxiv_id,
                    "semantic_scholar_id": node.semantic_scholar_id,
                    "cited_by": node.cited_by,
                }).candidate_dict()
            references = node.metadata.get("verified_unmatched_references") or []
            for reference in references:
                if isinstance(reference, dict) and reference.get("candidate_id") == candidate_id:
                    return sanitize_acquisition_candidate(reference)
        return None

    def get_candidate(self, project_id: str, candidate_id: str) -> dict[str, Any] | None:
        store = self._store(project_id)
        cached = store.candidate(candidate_id)
        if cached:
            return cached
        project_candidate = self._project_candidate(project_id, candidate_id)
        if project_candidate:
            store.upsert_papers([project_candidate])
        return project_candidate

    @staticmethod
    def _node(candidate: Mapping[str, Any], *, seed: bool = False) -> CitationGraphNode:
        candidate_id = str(candidate.get("candidate_id") or "")
        title = str(candidate.get("title") or "Untitled paper")
        author_details = candidate.get("author_details") if isinstance(candidate.get("author_details"), list) else []
        authors = str(candidate.get("authors") or ", ".join(
            str(row.get("name") or "") for row in author_details if isinstance(row, Mapping)
        ))
        year_value = candidate.get("year")
        try:
            year = int(year_value) if year_value is not None else None
        except (TypeError, ValueError):
            year = None
        first_author = authors.split(",", 1)[0].split(" and ", 1)[0].strip()
        label = f"{first_author or 'Paper'}, {year}" if year else first_author or "Paper"
        metadata = dict(candidate)
        metadata["external_candidate_id"] = candidate_id
        metadata["exploration_seed"] = seed
        return CitationGraphNode(
            id=f"external:{candidate_id}",
            type="paper",
            label=label,
            title=title,
            abstract=str(candidate.get("abstract") or ""),
            url=str(candidate.get("url") or candidate.get("canonical_url") or ""),
            year=year,
            authors=authors,
            cited_by=max(0, int(candidate.get("cited_by") or 0)),
            status="connector",
            source_file_ids=[],
            confidence=float(candidate.get("resolution_confidence") or 0.8),
            arxiv_id=str(candidate.get("arxiv_id") or ""),
            doi=str(candidate.get("doi") or ""),
            semantic_scholar_id=str(candidate.get("semantic_scholar_id") or ""),
            metadata=metadata,
        )

    @staticmethod
    def _edge(edge: VerifiedTopologyEdge, *, induced: bool) -> CitationGraphEdge:
        evidence = f"Verified by {edge.provider} work identifiers"
        digest = hashlib.sha256(
            f"{edge.provider}|{edge.source_candidate_id}|{edge.target_candidate_id}".encode("utf-8")
        ).hexdigest()[:20]
        return CitationGraphEdge(
            id=f"explore:{digest}",
            source=f"external:{edge.source_candidate_id}",
            target=f"external:{edge.target_candidate_id}",
            relation="cites",
            evidence=f"{evidence}{' (induced)' if induced else ''}",
            confidence=1.0,
        )

    def _graph_from_neighborhood(
        self, project_id: str, store: CitationExplorationStore, neighborhood: CachedNeighborhood
    ) -> CitationGraph:
        cache = store.load()
        nodes = [
            self._node(cache.papers[candidate_id], seed=candidate_id == neighborhood.seed_candidate_id)
            for candidate_id in neighborhood.paper_ids
            if candidate_id in cache.papers
        ]
        edges = [CitationGraphEdge(**row) for row in [*neighborhood.primary_edges, *neighborhood.induced_edges]]
        return CitationGraph(
            project_id=project_id,
            nodes=nodes,
            edges=edges,
            updated_at=neighborhood.updated_at,
            warnings=neighborhood.warnings,
        )

    def _response(
        self, project_id: str, store: CitationExplorationStore, neighborhood: CachedNeighborhood
    ) -> CitationExplorationResponse:
        graph = self._graph_from_neighborhood(project_id, store, neighborhood)
        if neighborhood.seed_candidate_id:
            seed_candidate_id = neighborhood.seed_candidate_id
        elif graph.nodes:
            seed_candidate_id = str(graph.nodes[0].metadata.get("external_candidate_id") or "")
        else:
            seed_candidate_id = ""
        return CitationExplorationResponse(
            graph=graph,
            seed_id=f"external:{seed_candidate_id}" if seed_candidate_id else "",
            seed_candidate_id=seed_candidate_id,
            neighborhood_key=neighborhood.key,
            exploration_kind=neighborhood.exploration_kind,
            primary_state=neighborhood.primary_state,
            induced_state=neighborhood.induced_state,
        )

    async def explore(self, project_id: str, candidate_id: str) -> CitationExplorationResponse:
        seed = self.get_candidate(project_id, candidate_id)
        if seed is None:
            raise LookupError("Verified citation candidate not found.")
        store = self._store(project_id)
        key = f"paper:{candidate_id}"
        cached = store.neighborhood(key)
        if cached:
            return self._response(project_id, store, cached)

        topology = await self.topology.paper_neighborhood(seed)
        store.upsert_papers(topology.papers)
        store.save_reference_sets(topology.reference_sets)
        primary_edges = [self._edge(edge, induced=False).model_dump() for edge in topology.primary_edges]
        resolved_seed_id = str(topology.seed.get("candidate_id") or candidate_id)
        neighborhood = store.new_neighborhood(
            key=key,
            exploration_kind="paper",
            seed_candidate_id=resolved_seed_id,
            paper_ids=[str(paper.get("candidate_id") or "") for paper in topology.papers if paper.get("candidate_id")],
            primary_edges=primary_edges,
            warnings=topology.warnings,
            degraded=bool(topology.warnings),
        )
        store.save_neighborhood(neighborhood)
        return self._response(project_id, store, neighborhood)

    async def enrich_edges(self, project_id: str, neighborhood_key: str) -> CitationExplorationDelta:
        store = self._store(project_id)
        neighborhood = store.neighborhood(neighborhood_key, fresh_only=False)
        if neighborhood is None:
            raise LookupError("Exploration neighborhood not found.")
        if neighborhood.induced_state in {"complete", "degraded"} and neighborhood.fresh:
            return CitationExplorationDelta(
                neighborhood_key=neighborhood_key,
                edges=[CitationGraphEdge(**row) for row in neighborhood.induced_edges],
                induced_state=neighborhood.induced_state,
                warnings=neighborhood.warnings,
            )

        cache = store.load()
        papers = [cache.papers[candidate_id] for candidate_id in neighborhood.paper_ids if candidate_id in cache.papers]
        reference_sets = {
            candidate_id: cache.reference_sets[candidate_id]
            for candidate_id in neighborhood.paper_ids
            if candidate_id in cache.reference_sets and cache.reference_sets[candidate_id].fresh
        }
        warnings = [] if len(reference_sets) == len(papers) else ["Some paper reference sets were unavailable."]
        induced = self.topology.induced_edges(papers, reference_sets)
        primary_pairs = {
            (str(row.get("source") or ""), str(row.get("target") or ""))
            for row in neighborhood.primary_edges
        }
        edge_rows = [
            self._edge(edge, induced=True).model_dump()
            for edge in induced
            if (f"external:{edge.source_candidate_id}", f"external:{edge.target_candidate_id}") not in primary_pairs
        ]
        merged = store.merge_induced_edges(neighborhood_key, edge_rows, warnings=warnings)
        return CitationExplorationDelta(
            neighborhood_key=neighborhood_key,
            edges=[CitationGraphEdge(**row) for row in merged.induced_edges],
            induced_state=merged.induced_state,
            warnings=merged.warnings,
        )

    async def details(self, project_id: str, candidate_id: str) -> CitationGraphNode:
        candidate = self.get_candidate(project_id, candidate_id)
        if candidate is None:
            raise LookupError("Verified citation candidate not found.")
        resolved = await self.metadata_service.resolve(candidate)
        payload = resolved.candidate_dict()
        self._store(project_id).upsert_papers([payload])
        return self._node(payload, seed=bool(candidate.get("exploration_seed")))

    async def explore_author(
        self, project_id: str, candidate_id: str, author_index: int
    ) -> CitationExplorationResponse:
        candidate = self.get_candidate(project_id, candidate_id)
        if candidate is None:
            raise LookupError("Verified citation candidate not found.")
        # Detail resolution supplies provider author identifiers before accepting the index.
        resolved = await self.metadata_service.resolve(candidate)
        resolved_candidate = resolved.candidate_dict()
        store = self._store(project_id)
        store.upsert_papers([resolved_candidate])
        author_details = resolved_candidate.get("author_details") or []
        if author_index < 0 or author_index >= len(author_details):
            raise LookupError("Author is not available on this paper.")
        author = author_details[author_index]
        openalex_id = str(author.get("openalex_id") or "") if isinstance(author, Mapping) else ""
        cached_key = f"author:openalex:{self.topology._work_id(openalex_id)}:top10" if openalex_id else ""
        if cached_key:
            cached = store.neighborhood(cached_key)
            if cached:
                return self._response(project_id, store, cached)

        key, papers, reference_sets = await self.topology.author_neighborhood(resolved_candidate, author_index)
        store.upsert_papers(papers)
        store.save_reference_sets(reference_sets)
        seed_candidate_id = str(papers[0].get("candidate_id") or "") if papers else ""
        neighborhood = store.new_neighborhood(
            key=key,
            exploration_kind="author",
            author_key=key,
            seed_candidate_id=seed_candidate_id,
            paper_ids=[str(paper.get("candidate_id") or "") for paper in papers if paper.get("candidate_id")],
            primary_edges=[],
            warnings=[] if papers else ["No papers were returned for this author."],
            degraded=not papers,
        )
        store.save_neighborhood(neighborhood)
        return self._response(project_id, store, neighborhood)

    async def download(self, project_id: str, candidate_id: str) -> PaperAcquisitionResult:
        candidate = self.get_candidate(project_id, candidate_id)
        if candidate is None:
            raise LookupError("Verified citation candidate not found.")
        return await self.acquisition_service.acquire(project_id, candidate)