import requests
import xml.etree.ElementTree as ET
import logging
import re
from typing import List, Dict, Any
from urllib.parse import quote

logger = logging.getLogger("ArxivTool")

class ArxivTool:
    """
    arXiv 真实文献检索工具 (精准学术关键词解析)
    """
    BASE_URL = "http://export.arxiv.org/api/query"

    @classmethod
    def search_papers(cls, query: str, max_results: int = 5, category: str = "cs.AI") -> List[Dict[str, Any]]:
        """
        检索 arXiv 真实论文
        """
        search_term = cls._extract_english_keywords(query)
        
        # 组装干净且经过转义的 API 查询参数
        search_query = f"cat:{category} AND all:{search_term}" if category else f"all:{search_term}"
        
        params = {
            "search_query": search_query,
            "start": 0,
            "max_results": max_results,
            "sortBy": "relevance",
            "sortOrder": "descending"
        }
        
        try:
            logger.info(f"正在从 arXiv 真实检索: original='{query}', term='{search_term}', cat='{category}'")
            response = requests.get(cls.BASE_URL, params=params, timeout=15)
            response.raise_for_status()
            papers = cls._parse_arxiv_xml(response.text)
            
            # 如果针对提炼的词仍为 0 篇，用该领域的基础关键词兜底检索一次（非 Mock 数据）
            if not papers and category:
                logger.warning(f"关键词 [{search_term}] 未命中论文，使用领域默认核心词 'large language model' 检索...")
                fallback_params = {
                    "search_query": f"cat:{category} AND all:\"large language model\"",
                    "start": 0,
                    "max_results": max_results,
                    "sortBy": "relevance",
                    "sortOrder": "descending"
                }
                fb_resp = requests.get(cls.BASE_URL, params=fallback_params, timeout=15)
                if fb_resp.status_code == 200:
                    papers = cls._parse_arxiv_xml(fb_resp.text)

            return papers
        except Exception as e:
            logger.error(f"arXiv API 检索请求异常: {e}")
            return []

    @staticmethod
    def _extract_english_keywords(text: str) -> str:
        """
        将中文难题提炼为精准的 arXiv 学术搜索项（带双引号短语）
        """
        # 中文/学术核心领域关键词词典映射
        kw_map = [
            ("大语言模型", "\"large language model\""),
            ("大模型", "\"large language model\""),
            ("推理", "reasoning"),
            ("幻觉", "hallucination"),
            ("自一致性", "\"Self-Consistency\""),
            ("多智能体", "\"multi-agent\""),
            ("智能体", "agent"),
            ("量子", "\"quantum computing\""),
            ("脉冲星", "\"pulsar timing\""),
            ("基因", "\"genome editing\"")
        ]
        
        extracted = []
        for zh, en in kw_map:
            if zh in text and en not in extracted:
                extracted.append(en)
        
        # 如果提取到了关键词，用 AND 或空格连接
        if extracted:
            return " AND ".join(extracted[:2]) # 选取前 2 个最核心的做联合精确检索
        
        # 提取原文本中的英文短语（忽略包含括号等干扰字符）
        clean_text = re.sub(r'[（）\(\)\[\]]', ' ', text)
        english_words = re.findall(r'[a-zA-Z0-9\-]+', clean_text)
        if len(english_words) >= 2:
            return f"\"{english_words[0]} {english_words[1]}\""
        elif len(english_words) == 1:
            return english_words[0]
            
        return "\"large language model\""

    @staticmethod
    def _parse_arxiv_xml(xml_content: str) -> List[Dict[str, Any]]:
        root = ET.fromstring(xml_content)
        ns = {'atom': 'http://www.w3.org/2005/Atom'}
        
        papers = []
        for entry in root.findall('atom:entry', ns):
            title = entry.find('atom:title', ns).text.strip().replace('\n', ' ')
            summary = entry.find('atom:summary', ns).text.strip().replace('\n', ' ')
            published = entry.find('atom:published', ns).text[:10] if entry.find('atom:published', ns) is not None else ""
            id_url = entry.find('atom:id', ns).text
            
            authors = []
            for author in entry.findall('atom:author', ns):
                name = author.find('atom:name', ns)
                if name is not None:
                    authors.append(name.text)
                    
            papers.append({
                "title": title,
                "authors": authors,
                "summary": summary,
                "published": published,
                "url": id_url,
                "doi": f"10.48550/arXiv.{id_url.split('/')[-1]}"
            })
            
        return papers

if __name__ == "__main__":
    test_q = "如何突破大语言模型在多步复杂科学推理中的逻辑幻觉问题，并构建具备自一致性（Self-Consistency）校验能力的自进化多智能体科研发现闭环？"
    results = ArxivTool.search_papers(test_q, max_results=5, category="cs.AI")
    print(f"真实检索到的论文数: {len(results)} 篇")
    for p in results:
        print(f"- [{p['published']}] {p['title']}")