arxiv.py

import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
import os
import re
import time
import tarfile
import shutil
import subprocess
from glob import glob
from typing import List, Dict, Optional
from pathlib import Path
import sys

try:
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass


DIRECT_OPENER = urllib.request.build_opener()


ARXIV_API = "https://export.arxiv.org/api/query"
ATOM_NS = {"atom": "http://www.w3.org/2005/Atom"}

class ArxivAgent:
    def __init__(self, max_results: int = 30, pdf_dir: str = "arxiv_papers", 
                 md_dir: str = "arxiv_papers_md", ddl: int = 180, 
                 download_mode: str = "source_first", download_dir: str | None = None):

        self.max_results = max_results
        self.pdf_dir = pdf_dir
        self.md_dir = md_dir
        self.ddl = ddl
        self.download_mode = download_mode
        self.download_dir = download_dir if download_dir is not None else pdf_dir
            
        self.last_query_start_ts: Optional[float] = None
    
    def search_and_analyze(self, query: str) -> List[Dict]:
        self.last_query_start_ts = time.time()
        
        search_query = f'all:{query}'
        
        raw_results = self._search_arxiv(search_query, max_results=self.max_results * 2) 

        final_papers = raw_results[:self.max_results]
        
        for paper in final_papers:
            paper['relevance_score'] = 1.0 
        
        return final_papers
    
    def download_papers(self, papers: List[Dict]) -> List[str]:
        downloaded_files = []
        
        for i, paper in enumerate(papers, 1):
            title = self._clean_filename(paper['title'])
            arxiv_id = paper.get('arxiv_id', '')
            start_time = time.time()
            
            global_start = self.last_query_start_ts if self.last_query_start_ts else start_time
            deadline = global_start + self.ddl
            
            try:
                base_name = f"{arxiv_id}_{title[:50]}" if arxiv_id else title[:50]
                
                existing_md = self._check_existing_markdown(base_name)
                if existing_md:
                    print(f"[{i}/{len(papers)}] Already exists: {title[:50]}... ({existing_md})")
                    downloaded_files.append(existing_md)
                    continue
                
                if time.time() >= deadline:
                    md_path = self._write_abstract_markdown(paper, base_name)
                    if md_path:
                        print(f"[{i}/{len(papers)}] Timeout, saving abstract: {title[:50]}...")
                        downloaded_files.append(md_path)
                    continue
                
                result = None
                if self.download_mode == "source_first":
                    result = self._download_source_first(paper, base_name, deadline)
                elif self.download_mode == "pdf_first":
                    result = self._download_pdf_first(paper, base_name, deadline)
                
                if result:
                    print(f"[{i}/{len(papers)}] {result['status']}: {title[:50]}...")
                    downloaded_files.append(result['path'])
                else:
                    print(f"[{i}/{len(papers)}] Download failed: {title[:50]}...")
                
                time.sleep(1)
                
            except Exception as e:
                print(f"[{i}/{len(papers)}] Processing failed: {str(e)}")
        
        return downloaded_files
    
    def _search_arxiv(self, query: str, max_results: int = 10) -> List[Dict]:
        params = {
            "search_query": query,
            "start": 0,
            "max_results": max_results,
            "sortBy": "relevance",
            "sortOrder": "descending"
        }
        
        url = f"{ARXIV_API}?{urllib.parse.urlencode(params)}"
        print(f"[DEBUG] Requesting arXiv API: {url[:100]}...")
        req = urllib.request.Request(
            url,
            headers={"User-Agent": "arxiv-agent/1.0 (research@example.com)"}
        )
        
        try:
            print(f"[DEBUG] Sending request, timeout set to 30 seconds...")
            with DIRECT_OPENER.open(req, timeout=30) as resp:
                xml_text = resp.read()
                print(f"[DEBUG] Response received, length: {len(xml_text)} bytes")
            
            root = ET.fromstring(xml_text)
            papers = []
            
            for entry in root.findall("atom:entry", ATOM_NS):
                paper = self._parse_entry(entry)
                if paper:
                    papers.append(paper)
            
            print(f"[DEBUG] Parsing complete, found {len(papers)} papers")
            return papers
            
        except Exception as e:
            print(f"[ERROR] Search failed: {type(e).__name__}: {str(e)}")
            return []
    
    def _parse_entry(self, entry) -> Optional[Dict]:
        title = (entry.findtext("atom:title", default="", namespaces=ATOM_NS) or "").strip()
        summary = (entry.findtext("atom:summary", default="", namespaces=ATOM_NS) or "").strip()
        published = entry.findtext("atom:published", default="", namespaces=ATOM_NS)
        abs_url = entry.findtext("atom:id", default="", namespaces=ATOM_NS)
        
        authors = []
        for a in entry.findall("atom:author", ATOM_NS):
            name = a.findtext("atom:name", default="", namespaces=ATOM_NS)
            if name:
                authors.append(name)
        
        pdf_url = ""
        source_url = ""
        for link in entry.findall("atom:link", ATOM_NS):
            if link.attrib.get("type") == "application/pdf":
                pdf_url = link.attrib.get("href", "")
            elif link.attrib.get("title") == "pdf":
                pdf_url = link.attrib.get("href", "")
        
        arxiv_id = abs_url.rsplit("/", 1)[-1] if abs_url else ""
        
        if arxiv_id:
            source_url = f"https://arxiv.org/e-print/{arxiv_id}"
        
        if not all([title, summary, arxiv_id]):
            return None
        
        return {
            "title": " ".join(title.split()),
            "abstract": " ".join(summary.split()),
            "published": published,
            "authors": authors,
            "abs_url": abs_url,
            "pdf_url": pdf_url,
            "source_url": source_url,
            "arxiv_id": arxiv_id,
        }
    
    def _download_pdf(self, paper: Dict, title: str) -> Optional[str]:
        if not paper['pdf_url']:
            return None
        
        filename = f"{paper['arxiv_id']}_{title[:50]}.pdf"
        filepath = os.path.join(self.download_dir, filename)
        
        req = urllib.request.Request(
            paper['pdf_url'],
            headers={"User-Agent": "arxiv-agent/1.0"}
        )
        
        with DIRECT_OPENER.open(req, timeout=60) as response:
            with open(filepath, 'wb') as f:
                f.write(response.read())
        
        return filepath
    
    def _download_source(self, paper: Dict, title: str, deadline: Optional[float] = None) -> Optional[str]:
        if not paper['source_url']:
            return None

        filename = f"{paper['arxiv_id']}_{title[:50]}_source.tar.gz"
        filepath = os.path.join(self.download_dir, filename)

        try:
            req = urllib.request.Request(
                paper['source_url'],
                headers={"User-Agent": "arxiv-agent/1.0"}
            )

            per_read_timeout = 10
            if deadline is not None:
                remaining = max(0.0, deadline - time.time())
                per_read_timeout = max(1.0, min(10.0, remaining))

            with DIRECT_OPENER.open(req, timeout=per_read_timeout) as response:
                with open(filepath, 'wb') as f:
                    while True:
                        if (deadline is not None) and (time.time() >= deadline):
                            try: f.close()
                            except Exception: pass
                            try:
                                if os.path.exists(filepath): os.remove(filepath)
                            except Exception: pass
                            return None
                        try:
                            chunk = response.read(64 * 1024)
                        except Exception:
                            chunk = b""
                        if not chunk:
                            break
                        f.write(chunk)

            return filepath
        except Exception:
            try:
                if os.path.exists(filepath):
                    os.remove(filepath)
            except Exception:
                pass
            return None
    
    def _clean_filename(self, filename: str) -> str:
        filename = re.sub(r'[<>:"/\\|?*]', '_', filename)
        return filename[:100]
    
    def _check_existing_markdown(self, base_name: str) -> Optional[str]:
        md_paths = glob(os.path.join(self.download_dir, f"{base_name}*.md"))
        if md_paths:
            return md_paths[0]
        md_paths = glob(os.path.join(self.md_dir, f"{base_name}*.md"))
        if md_paths:
            return md_paths[0]
        return None

    def _download_source_first(self, paper: Dict, base_name: str, deadline: Optional[float] = None) -> Optional[Dict]:
        if paper['source_url'] and time.time() < deadline:
            try:
                print(f"[INFO] Attempting to download source: {paper['arxiv_id']}")
                archive_path = self._download_source(paper, base_name, deadline=deadline)
                
                if archive_path and time.time() < deadline:
                    print(f"[INFO] Attempting to convert source to Markdown: {paper['arxiv_id']}")
                    md_path = self.convert_source_archive_to_markdown(archive_path, deadline=deadline)
                    if md_path:
                        target_md_path = os.path.join(self.md_dir, os.path.basename(md_path))
                        shutil.copy(md_path, target_md_path)
                        return {'path': target_md_path, 'status': 'Source converted to Markdown'}
            except Exception as e:
                print(f"[WARNING] Source processing failed: {e}")
        
        return self._download_pdf_first(paper, base_name, deadline)

    def _download_pdf_first(self, paper: Dict, base_name: str, deadline: Optional[float] = None) -> Optional[Dict]:
        pdf_path = None
        if paper['pdf_url'] and time.time() < deadline:
            try:
                print(f"[INFO] Attempting to download PDF: {paper['arxiv_id']}")
                pdf_path = self._download_pdf(paper, base_name)
            except Exception as e:
                print(f"[WARNING] PDF download failed: {e}")
        
        md_path = self._write_abstract_markdown(paper, base_name)
        if md_path:
            return {'path': md_path, 'status': 'Saved abstract Markdown'}
            
        return None
        
    def convert_source_archive_to_markdown(self, archive_path: str, deadline: Optional[float] = None) -> Optional[str]:
        try:
            if not os.path.isfile(archive_path) or not archive_path.endswith('.tar.gz'):
                return None

            if (deadline is not None) and (time.time() >= deadline):
                return None

            base_dir = os.path.dirname(archive_path)
            base_name = os.path.basename(archive_path)[:-7]
            extract_dir = os.path.join(base_dir, base_name)
            if not os.path.exists(extract_dir):
                os.makedirs(extract_dir)

            with tarfile.open(archive_path, 'r:gz') as tar:
                ok = _safe_extract_tar(tar, path=extract_dir, deadline=deadline)
                if not ok:
                    return None

            if (deadline is not None) and (time.time() >= deadline):
                return None

            main_tex = _guess_main_tex_file(extract_dir, deadline=deadline)
            if not main_tex:
                if (deadline is not None) and (time.time() >= deadline):
                    return None
                tex_files = glob(os.path.join(extract_dir, '**', '*.tex'), recursive=True)
                if not tex_files:
                    return None
                main_tex = tex_files[0]

            out_md = os.path.join(base_dir, f"{base_name}_source.md")
            if shutil.which('pandoc'):
                try:
                    remaining = None
                    if deadline is not None:
                        remaining = max(0.0, deadline - time.time())
                        if remaining <= 1.0:
                            raise TimeoutError("deadline too close for pandoc conversion")
                    subprocess.run(
                        ['pandoc', '-f', 'latex', '-t', 'gfm', main_tex, '-o', out_md],
                        cwd=os.path.dirname(main_tex),
                        check=True,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        shell=False,
                        timeout=remaining if remaining and remaining > 0 else None,
                    )
                    return out_md
                except Exception:
                    pass

            if (deadline is not None) and (time.time() >= deadline):
                return None

            try:
                full_text = _read_tex_with_includes(main_tex, deadline=deadline)
                md_text = _latex_to_markdown_basic(full_text)
                with open(out_md, 'w', encoding='utf-8') as f:
                    f.write(md_text)
                return out_md
            except Exception:
                return None
        except Exception:
            return None

    def _write_abstract_markdown(self, paper: Dict, title_clean: Optional[str] = None) -> Optional[str]:
        try:
            title = paper.get('title') or ''
            authors = paper.get('authors') or []
            published = (paper.get('published') or '')[:10]
            abs_url = paper.get('abs_url') or ''
            pdf_url = paper.get('pdf_url') or ''
            source_url = paper.get('source_url') or ''
            arxiv_id = paper.get('arxiv_id') or ''
            abstract = paper.get('abstract') or ''

            safe_title = self._clean_filename(title if not title_clean else title_clean)
            filename = f"{arxiv_id}_{safe_title[:50]}_abstract.md" if arxiv_id else f"{safe_title[:50]}_abstract.md"
            out_path = os.path.join(self.download_dir, filename)

            lines = []
            if title:
                lines.append(f"# {title}")
                lines.append("")
            meta = []
            if authors:
                meta.append(f"**Authors**: {', '.join(authors)}")
            if published:
                meta.append(f"**Published**: {published}")
            if arxiv_id:
                meta.append(f"**arXiv ID**: {arxiv_id}")
            if meta:
                lines.append("\n".join(meta))
                lines.append("")
            links = []
            if abs_url:
                links.append(f"- [Abstract]({abs_url})")
            if pdf_url:
                links.append(f"- [PDF]({pdf_url})")
            if source_url:
                links.append(f"- [Source]({source_url})")
            if links:
                lines.append("**Links**:")
                lines.extend(links)
                lines.append("")
            lines.append("## Abstract")
            lines.append("")
            lines.append(abstract)
            lines.append("")

            with open(out_path, 'w', encoding='utf-8') as f:
                f.write("\n".join(lines))
            return out_path
        except Exception:
            return None
        

def _safe_extract_tar(tar: tarfile.TarFile, path: str, deadline: Optional[float] = None) -> bool:
    for member in tar.getmembers():
        if (deadline is not None) and (time.time() >= deadline):
            return False
        member_path = os.path.join(path, member.name)
        abs_path = os.path.abspath(member_path)
        abs_base = os.path.abspath(path)
        if not abs_path.startswith(abs_base):
            continue
        try:
            tar.extract(member, path)
        except Exception:
            continue
    return True


def _guess_main_tex_file(root_dir: str, deadline: Optional[float] = None) -> Optional[str]:
    candidates = []
    for tex in glob(os.path.join(root_dir, '**', '*.tex'), recursive=True):
        if (deadline is not None) and (time.time() >= deadline):
            break
        try:
            with open(tex, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read(5000)
            if ('\\documentclass' in content) and ('\\begin{document}' in content):
                candidates.append(tex)
        except Exception:
            continue
    if candidates:
        candidates.sort(key=lambda p: (p.count(os.sep), len(p)))
        return candidates[0]
    return None


def _read_tex_with_includes(main_tex_path: str, deadline: Optional[float] = None) -> str:
    visited = set()

    def _read(path: str, base_dir: str) -> str:
        if (deadline is not None) and (time.time() >= deadline):
            return ""
        norm = os.path.normpath(os.path.abspath(path))
        if norm in visited:
            return ""
        visited.add(norm)
        try:
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                text = f.read()
        except Exception:
            return ""

        def _replace_include(m):
            inner = m.group(1).strip()
            if not os.path.splitext(inner)[1]:
                inner = inner + '.tex'
            sub_path = os.path.join(base_dir, inner)
            return _read(sub_path, os.path.dirname(sub_path))

        pattern = re.compile(r"\\(?:input|include)\{([^}]+)\}")
        try:
            text = re.sub(pattern, _replace_include, text)
        except Exception:
            pass
        return text

    return _read(main_tex_path, os.path.dirname(main_tex_path))


def _latex_to_markdown_basic(tex: str) -> str:
    s = tex
    lines = []
    for line in s.splitlines():
        if '%' in line:
            idx = line.find('%')
            if idx >= 0:
                line = line[:idx]
        lines.append(line)
    s = "\n".join(lines)

    def remove_env(name: str, text: str) -> str:
        return re.sub(rf"\\begin\{{{name}\}}[\s\S]*?\\end\{{{name}\}}", "", text, flags=re.DOTALL)
    for env in ["figure", "table", "algorithm", "lstlisting"]:
        s = remove_env(env, s)

    s = re.sub(r"\\section\*?\{([^}]*)\}", r"\n\n## \1\n\n", s)
    s = re.sub(r"\\subsection\*?\{([^}]*)\}", r"\n\n### \1\n\n", s)
    s = re.sub(r"\\subsubsection\*?\{([^}]*)\}", r"\n\n#### \1\n\n", s)

    s = re.sub(r"\\textbf\{([^}]*)\}", r"**\1**", s)
    s = re.sub(r"\\emph\{([^}]*)\}", r"*\1*", s)
    s = re.sub(r"\\textit\{([^}]*)\}", r"*\1*", s)

    s = re.sub(r"\\cite[t|p]?\{[^}]*\}", "", s)
    s = re.sub(r"\\ref\{[^}]*\}", "", s)
    s = re.sub(r"\\label\{[^}]*\}", "", s)

    s = re.sub(r"\\[a-zA-Z]+(?:\[[^\]]*\])?\{([^}]*)\}", r"\1", s)
    s = re.sub(r"\\[a-zA-Z]+(?:\[[^\]]*\])?", "", s)

    s = re.sub(r"\n{3,}", "\n\n", s)
    s = re.sub(r"[ \t]{2,}", " ", s)
    return s.strip() + "\n"




def search_relevant_papers(query: str, max_results: int = 30) -> List[Dict]:
    print(f"[DEBUG] search_relevant_papers called, query: '{query}', max_results={max_results}")
    agent = ArxivAgent(max_results=max_results)
    ranked = agent.search_and_analyze(query)
    print(f"[DEBUG] search_relevant_papers returning {len(ranked)} papers")
    return ranked[:max_results]


def _extract_arxiv_id_from_url(url: str) -> Optional[str]:
    try:
        parsed = urllib.parse.urlparse(url)
        path = parsed.path.strip('/')
        parts = path.split('/')
        if not parts:
            return None

        if parts[0] in {"abs", "pdf", "e-print"} and len(parts) >= 2:
            arxiv_part = parts[1]
            if arxiv_part.endswith('.pdf'):
                arxiv_part = arxiv_part[:-4]
            arxiv_part = re.sub(r"v\d+$", "", arxiv_part)
            if re.match(r"^\d{4}\.\d{4,5}$", arxiv_part):
                return arxiv_part
            if arxiv_part:
                return arxiv_part
        candidate = re.sub(r"v\d+$", "", parts[-1])
        candidate = candidate.replace('.pdf', '')
        if candidate:
            return candidate
        return None
    except Exception:
        return None


def _fetch_metadata_by_id(arxiv_id: str) -> Optional[Dict]:
    try:
        query = f"id_list={arxiv_id}"
        url = f"{ARXIV_API}?{query}"
        req = urllib.request.Request(url, headers={"User-Agent": "arxiv-agent/1.0 (metadata)"})
        with DIRECT_OPENER.open(req, timeout=20) as resp:
            xml_text = resp.read()
        root = ET.fromstring(xml_text)
        entry = root.find("atom:entry", ATOM_NS)
        if entry is None:
            return None
        title = (entry.findtext("atom:title", default="", namespaces=ATOM_NS) or "").strip()
        summary = (entry.findtext("atom:summary", default="", namespaces=ATOM_NS) or "").strip()
        published = entry.findtext("atom:published", default="", namespaces=ATOM_NS)
        abs_url = entry.findtext("atom:id", default=f"https://arxiv.org/abs/{arxiv_id}", namespaces=ATOM_NS)
        
        authors = []
        for a in entry.findall("atom:author", ATOM_NS):
            name = a.findtext("atom:name", default="", namespaces=ATOM_NS)
            if name:
                authors.append(name)
        
        pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
        for link in entry.findall("atom:link", ATOM_NS):
            if link.attrib.get("type") == "application/pdf":
                pdf_url = link.attrib.get("href", pdf_url)
                break
        source_url = f"https://arxiv.org/e-print/{arxiv_id}"
        return {
            "title": " ".join(title.split()),
            "abstract": " ".join(summary.split()),
            "published": published,
            "authors": authors,
            "abs_url": abs_url,
            "pdf_url": pdf_url,
            "source_url": source_url,
            "arxiv_id": arxiv_id,
        }
    except Exception:
        return None