app/services/paper_acquisition.py

"""Explicit, bounded acquisition of verified open-access project papers."""
from __future__ import annotations

import hashlib
import os
import re
import time
from collections.abc import Mapping
from typing import Any, Literal
from urllib.parse import urljoin, urlparse

from pydantic import BaseModel, ConfigDict

from app.schemas.project import ProjectFile
from app.services.paper_metadata import PaperMetadataService
from app.services.project_activity import ProjectActivityService
from app.services.project_files import project_upload_dir
from app.services.project_service import ProjectService


AcquisitionStatus = Literal["added", "already_present", "unavailable", "failed"]
_PDF_TYPES = {"application/pdf", "application/x-pdf", "application/acrobat", "applications/vnd.pdf"}
_DEFAULT_TRUSTED_OA_HOSTS = {
    "arxiv.org", "export.arxiv.org", "aclanthology.org", "openreview.net",
    "proceedings.mlr.press", "proceedings.neurips.cc", "papers.nips.cc",
    "pmc.ncbi.nlm.nih.gov", "www.ncbi.nlm.nih.gov", "europepmc.org",
    "www.ebi.ac.uk", "pdfs.semanticscholar.org",
}


def trusted_oa_hosts() -> set[str]:
    extension = {host.strip().lower() for host in os.getenv("PAPER_OA_ALLOWED_HOSTS", "").split(",") if host.strip()}
    return _DEFAULT_TRUSTED_OA_HOSTS | extension


def is_acquirable_pdf_url(value: str) -> bool:
    parsed = urlparse(str(value or ""))
    host = (parsed.hostname or "").lower()
    return parsed.scheme == "https" and bool(host) and any(
        host == trusted or host.endswith(f".{trusted}") for trusted in trusted_oa_hosts()
    )


def sanitize_acquisition_candidate(candidate: Mapping[str, Any]) -> dict[str, Any]:
    payload = dict(candidate)
    arxiv_id = PaperMetadataService.normalize_arxiv(payload.get("arxiv_id") or "")
    if arxiv_id and not payload.get("pdf_url"):
        payload["arxiv_id"] = arxiv_id
        payload["pdf_url"] = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
        payload["is_open_access"] = True
        payload["pdf_provider"] = payload.get("pdf_provider") or "arXiv"
    can_acquire = bool(payload.get("is_open_access") and is_acquirable_pdf_url(str(payload.get("pdf_url") or "")))
    payload["can_acquire"] = can_acquire
    if not can_acquire:
        payload["is_open_access"] = False
        payload["pdf_url"] = ""
    return payload


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

    ok: bool
    status: AcquisitionStatus
    message: str
    project: dict[str, Any] | None = None
    citation_graph: dict[str, Any] | None = None
    partial: bool = False
    recovery_required: bool = False


class PaperAcquisitionService:
    def __init__(
        self,
        metadata_service: PaperMetadataService | None = None,
        project_service: ProjectService | None = None,
        activity_service: ProjectActivityService | None = None,
    ) -> None:
        self.metadata_service = metadata_service or PaperMetadataService()
        self.project_service = project_service or ProjectService()
        self.activity_service = activity_service or ProjectActivityService()
        self.trusted_hosts = trusted_oa_hosts()
        self.max_bytes = max(1_048_576, int(os.getenv("PAPER_ACQUISITION_MAX_BYTES", str(50 * 1024 * 1024))))

    async def acquire(self, project_id: str, candidate: Mapping[str, Any]) -> PaperAcquisitionResult:
        from app.services.citation_graph import CitationGraphService

        if not CitationGraphService.is_valid_project_id(project_id):
            return self._result(False, "failed", "Invalid project id.")
        async with self.project_service.mutation(project_id):
            return await self._acquire_locked(project_id, candidate)

    async def _acquire_locked(self, project_id: str, candidate: Mapping[str, Any]) -> PaperAcquisitionResult:
        from app.services.citation_graph import CitationGraphService

        # Resolution and download are inside the project mutation boundary so a
        # second add cannot race this operation's eventual reload/dedupe/save.
        metadata = await self.metadata_service.resolve(candidate)
        if not metadata.is_open_access or not metadata.pdf_url:
            return self._result(False, "unavailable", "No downloadable open-access PDF found.")
        try:
            content = await self._download_pdf(metadata.pdf_url)
        except ValueError as exc:
            return self._result(False, "unavailable", str(exc))
        except Exception as exc:  # Network failures are useful but must not crash the UI.
            self.activity_service.record(project_id, "paper_acquisition", "Open-access paper download failed", status="warning", metadata={"title": metadata.title, "error": str(exc)})
            return self._result(False, "failed", "The open-access PDF could not be downloaded.")

        file_id = hashlib.sha256(content).hexdigest()
        # Reload immediately before dedupe/update to preserve changes made while
        # metadata was cached or before this operation acquired the mutation lock.
        project = self.project_service.load(project_id)
        if project is None:
            return self._result(False, "failed", "Project not found.")
        project_payload = self._project_payload(project)
        if any(file.file_id == file_id for file in project.files):
            graph, graph_error = await self._graph_after_mutation(project_id)
            if graph is None:
                return self._result(False, "failed", f"Paper is present, but citation refresh failed: {graph_error}", project_payload, partial=True, recovery_required=True)
            return self._result(True, "already_present", "This paper is already in the project.", project_payload, graph)

        upload_dir = project_upload_dir(project_id)
        filename = self._unique_filename(upload_dir, self._safe_filename(metadata.title), file_id)
        destination = os.path.join(upload_dir, filename)
        old_files = list(project.files)
        db = None
        try:
            # Content stays in memory until both signature and content type have passed.
            with open(destination, "xb") as file:
                file.write(content)
            project.files.append(ProjectFile(filename=filename, file_id=file_id))
            project.updated_at = time.time()
            self.project_service.save(project)
            from app.rag.ingestion import LIBRARY_COLLECTION, ingest_file
            from app.websockets.handlers import get_db

            db = get_db()
            ingest_file(content, filename, LIBRARY_COLLECTION, db=db, skip_if_indexed=False, project_id=project_id)
            # Keep the provider-verified identity attached to this exact PDF.
            # The canonical evidence parser remains source-of-truth for paper
            # content, but its layout-derived title must never replace this.
            self.metadata_service.remember_file_metadata(file_id, metadata)
        except Exception as exc:
            cleanup_errors = self._compensate_ingest(db, project_id, file_id)
            self.metadata_service.forget_file_metadata(file_id)
            try:
                if os.path.exists(destination):
                    os.remove(destination)
            except Exception as cleanup_exc:
                cleanup_errors.append(f"file cleanup: {cleanup_exc}")
            project.files = old_files
            project.updated_at = time.time()
            try:
                self.project_service.save(project)
            except Exception as cleanup_exc:
                cleanup_errors.append(f"project registry rollback: {cleanup_exc}")
            self.activity_service.record(project_id, "paper_acquisition", "Adding open-access paper failed", status="error", metadata={"title": metadata.title, "error": str(exc), "cleanup_errors": cleanup_errors})
            message = "The validated PDF could not be added to this project."
            if cleanup_errors:
                message += " Storage cleanup needs attention."
            return self._result(False, "failed", message, partial=bool(cleanup_errors), recovery_required=bool(cleanup_errors))

        self.activity_service.record(project_id, "paper_added", "Added open-access paper", metadata={"title": metadata.title, "filename": filename, "provider": metadata.provider})
        graph, graph_error = await self._graph_after_mutation(project_id, force_refresh=True)
        if graph is not None:
            self.activity_service.record(project_id, "citation_graph_refresh", "Citation map refreshed after paper add")
            return self._result(True, "added", "Paper added to the project.", self._project_payload(project), graph)
        self.activity_service.record(project_id, "citation_graph_refresh", "Citation map refresh failed after paper add", status="warning", metadata={"error": graph_error})
        return self._result(True, "added", "Paper added; citation refresh needs recovery.", self._project_payload(project), None, partial=True, recovery_required=True)

    @staticmethod
    def _compensate_ingest(db: Any, project_id: str, file_id: str) -> list[str]:
        if db is None:
            return []
        errors: list[str] = []
        try:
            from app.rag.evidence_ingestion import EvidenceIngestionService

            EvidenceIngestionService(db=db).remove_document(project_id, file_id)
        except Exception as exc:
            errors.append(f"evidence cleanup: {exc}")
        return errors

    @staticmethod
    async def _graph_after_mutation(project_id: str, force_refresh: bool = False) -> tuple[dict[str, Any] | None, str]:
        from app.services.citation_graph import CitationGraphService
        service = CitationGraphService()
        errors: list[str] = []
        for attempt in range(2):
            try:
                graph = await service.refresh(project_id) if force_refresh or attempt else service.load(project_id)
                if not graph.nodes or graph.file_fingerprint != service.file_fingerprint(project_id):
                    graph = await service.refresh(project_id)
                return graph.model_dump(), ""
            except Exception as exc:
                errors.append(str(exc))
                force_refresh = True
        return None, "; ".join(errors)

    async def _download_pdf(self, source_url: str) -> bytes:
        import httpx

        url = self._validated_url(source_url)
        timeout = httpx.Timeout(connect=8.0, read=20.0, write=20.0, pool=8.0)
        async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client:
            for _ in range(6):
                async with client.stream("GET", url) as response:
                    if response.status_code in {301, 302, 303, 307, 308}:
                        location = response.headers.get("location")
                        if not location:
                            raise ValueError("The open-access source returned an invalid redirect.")
                        url = self._validated_url(urljoin(str(response.url), location))
                        continue
                    if response.status_code != 200:
                        raise ValueError("The open-access source did not provide a downloadable PDF.")
                    content_type = response.headers.get("content-type", "").split(";", 1)[0].strip().lower()
                    if content_type not in _PDF_TYPES:
                        raise ValueError("The open-access source did not provide a PDF content type.")
                    length = response.headers.get("content-length")
                    if length and int(length) > self.max_bytes:
                        raise ValueError("The open-access PDF exceeds the download size limit.")
                    body = bytearray()
                    async for chunk in response.aiter_bytes():
                        body.extend(chunk)
                        if len(body) > self.max_bytes:
                            raise ValueError("The open-access PDF exceeds the download size limit.")
                    content = bytes(body)
                    if not content.startswith(b"%PDF"):
                        raise ValueError("The open-access source did not provide a valid PDF.")
                    return content
        raise ValueError("The open-access source redirected too many times.")

    def _validated_url(self, value: str) -> str:
        if not is_acquirable_pdf_url(value):
            raise ValueError("No downloadable open-access PDF found.")
        return value

    def _trusted_host(self, host: str) -> bool:
        return any(host == trusted or host.endswith(f".{trusted}") for trusted in self.trusted_hosts)

    @staticmethod
    def _safe_filename(title: str) -> str:
        clean = re.sub(r"[^A-Za-z0-9 ._-]+", "_", title or "open-access-paper").strip(" ._")[:140]
        return f"{clean or 'open-access-paper'}.pdf"

    @staticmethod
    def _unique_filename(upload_dir: str, filename: str, file_id: str) -> str:
        stem, ext = os.path.splitext(os.path.basename(filename))
        candidate = f"{stem}{ext or '.pdf'}"
        if not os.path.exists(os.path.join(upload_dir, candidate)):
            return candidate
        return f"{stem}-{file_id[:8]}{ext or '.pdf'}"

    def _project_payload(self, project: Any) -> dict[str, Any]:
        return {**project.model_dump(), "document_id": self.project_service.derive_document_id(project)}

    @staticmethod
    def _result(ok: bool, status: AcquisitionStatus, message: str, project: dict[str, Any] | None = None, citation_graph: dict[str, Any] | None = None, partial: bool = False, recovery_required: bool = False) -> PaperAcquisitionResult:
        return PaperAcquisitionResult(ok=ok, status=status, message=message, project=project, citation_graph=citation_graph, partial=partial, recovery_required=recovery_required)