services/paper_fetcher.py

import os
import requests
import fitz  # PyMuPDF
import logging
import io
import trafilatura
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class PaperFetcher:
    def __init__(self):
        self.api_key = os.environ.get("OPENALEX_API_KEY")
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
            "Accept-Language": "en-US,en;q=0.9",
            "Accept-Encoding": "gzip, deflate, br",
            "DNT": "1",
            "Connection": "keep-alive",
            "Upgrade-Insecure-Requests": "1",
            "Sec-Fetch-Dest": "document",
            "Sec-Fetch-Mode": "navigate",
            "Sec-Fetch-Site": "none",
            "Sec-Fetch-User": "?1"
        }
        self.session = requests.Session()
        self.session.headers.update(self.headers)

    def get_work_metadata(self, work_id):
        """Fetch metadata from OpenAlex."""
        if not work_id.startswith("https://openalex.org/"):
            if work_id.startswith("W"):
                work_id = f"https://openalex.org/{work_id}"
        
        url = f"https://api.openalex.org/works/{work_id.split('/')[-1]}"
        params = {}
        if self.api_key:
            params["api_key"] = self.api_key
            
        try:
            response = requests.get(url, params=params, headers=self.headers)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"Failed to fetch metadata from OpenAlex: {e}")
            return None

    def get_semantic_scholar_pdf(self, doi):
        """Fallback: Fetch PDF URL from Semantic Scholar API."""
        if not doi:
            return None
        
        clean_doi = doi.replace("https://doi.org/", "")
        url = f"https://api.semanticscholar.org/graph/v1/paper/DOI:{clean_doi}?fields=openAccessPdf"
        
        try:
            logger.info(f"Checking Semantic Scholar for DOI: {clean_doi}...")
            response = requests.get(url, headers=self.headers, timeout=15)
            response.raise_for_status()
            data = response.json()
            oa_pdf = data.get("openAccessPdf")
            if oa_pdf and oa_pdf.get("url"):
                return oa_pdf["url"]
        except Exception as e:
            logger.error(f"Semantic Scholar fallback failed: {e}")
        return None

    def fetch_pdf_content(self, pdf_url):
        """Fetch PDF content with hardened redirect and session handling."""
        try:
            logger.info(f"Fetching PDF from {pdf_url}...")
            # For PMC and similar, we need to be careful with redirects
            response = self.session.get(pdf_url, timeout=30, allow_redirects=True)
            
            # Handle the "Too many redirects" or meta-refresh redirects manually if needed
            if response.status_code == 200 and 'application/pdf' in response.headers.get('Content-Type', ''):
                return response.content
            
            # If it's HTML, we might have been redirected to a challenge page
            if 'text/html' in response.headers.get('Content-Type', ''):
                logger.warning(f"PDF URL {pdf_url} returned HTML instead of PDF. Possibly a bot challenge.")
            
            response.raise_for_status()
            return response.content
        except Exception as e:
            logger.error(f"Failed to fetch PDF from {pdf_url}: {e}")
            return None

    def extract_text_from_bytes(self, pdf_bytes):
        """Extract all text content from PDF bytes."""
        try:
            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
            text = ""
            for page in doc:
                text += page.get_text()
            doc.close()
            return text.strip()
        except Exception as e:
            logger.error(f"Failed to extract text from PDF bytes: {e}")
            return None

    def extract_from_html(self, url):
        """Extract content from HTML landing page by downloading first to handle cookies/redirects."""
        try:
            logger.info(f"Attempting HTML extraction from {url}...")
            # Use our session to get the HTML
            response = self.session.get(url, timeout=20)
            response.raise_for_status()
            
            downloaded = response.text
            if downloaded:
                # Pass the HTML content directly to trafilatura
                result = trafilatura.extract(downloaded, include_comments=False, include_tables=True)
                if result:
                    logger.info(f"Successfully extracted {len(result)} chars from HTML.")
                    return result.strip()
        except Exception as e:
            logger.error(f"HTML extraction failed for {url}: {e}")
        return None

    def fetch_full_text(self, work_id):
        """Main method: arXiv -> OpenAlex -> Semantic Scholar -> HTML Scrape -> Abstract."""
        metadata = self.get_work_metadata(work_id)
        if not metadata:
            return None

        # 1. Gather all potential PDF candidates
        pdf_candidates = []
        
        # Check arXiv
        ids = metadata.get("ids", {})
        if "arxiv" in ids:
            arxiv_id = ids["arxiv"].split("/")[-1].replace("abs/", "").replace("arxiv:", "")
            pdf_candidates.append(f"https://arxiv.org/pdf/{arxiv_id}.pdf")
            
        # Check OpenAlex locations
        best_loc = metadata.get("best_oa_location")
        if best_loc and best_loc.get("pdf_url"):
            pdf_candidates.append(best_loc["pdf_url"])
            
        for loc in metadata.get("locations", []):
            if loc.get("pdf_url") and loc["pdf_url"] not in pdf_candidates:
                pdf_candidates.append(loc["pdf_url"])

        # 2. Try each PDF candidate
        logger.info(f"Found {len(pdf_candidates)} PDF candidates for {work_id}")
        for url in pdf_candidates:
            pdf_bytes = self.fetch_pdf_content(url)
            if pdf_bytes:
                text = self.extract_text_from_bytes(pdf_bytes)
                if text and len(text) > 200: 
                    logger.info(f"Successfully extracted {len(text)} chars from {url}")
                    return text
                else:
                    logger.warning(f"Extraction from {url} was too short or empty.")

        # 3. If PDF fails, try Semantic Scholar for a new PDF link
        ss_pdf_url = self.get_semantic_scholar_pdf(metadata.get("doi"))
        if ss_pdf_url and ss_pdf_url not in pdf_candidates:
            pdf_bytes = self.fetch_pdf_content(ss_pdf_url)
            if pdf_bytes:
                text = self.extract_text_from_bytes(pdf_bytes)
                if text and len(text) > 200:
                    return text

        # 4. If all PDFs fail, try HTML scraping from landing page
        landing_page = metadata.get("landing_page_url")
        if not landing_page and best_loc:
            landing_page = best_loc.get("landing_page_url")
            
        if landing_page:
            text = self.extract_from_html(landing_page)
            if text and len(text) > 500: # HTML extraction should be substantial
                return text

        # 5. Final Safety Net: Reconstruct Abstract
        logger.warning(f"Could not get full text for {work_id}. Falling back to abstract.")
        return self.get_abstract(metadata)

    def get_abstract(self, metadata):
        """Reconstruct abstract from inverted index."""
        inverted_index = metadata.get("abstract_inverted_index")
        if not inverted_index:
            return ""
        
        max_index = 0
        for indices in inverted_index.values():
            if indices:
                max_index = max(max_index, max(indices))
        
        abstract_list = [""] * (max_index + 1)
        for word, indices in inverted_index.items():
            for index in indices:
                abstract_list[index] = word
                
        return " ".join(abstract_list).strip()

# Test Block
if __name__ == "__main__":
    fetcher = PaperFetcher()
    # Testing the PeerJ paper that has been failing
    test_id = "W2741809807"
    full_text = fetcher.fetch_full_text(test_id)
    
    if full_text:
        print("\n--- Success! First 200 characters ---")
        print(full_text[:200])
        print(f"\nTotal characters fetched: {len(full_text):,}")
    else:
        print("Failed to fetch any text.")