#!/usr/bin/env python3
"""Search academic databases for research papers."""

import argparse
import json
import urllib.request
import urllib.parse
import xml.etree.ElementTree as ET
from typing import Optional

def search_semantic_scholar(query: str, limit: int = 10) -> list:
    """Search Semantic Scholar API."""
    base_url = "https://api.semanticscholar.org/graph/v1/paper/search"
    params = {
        "query": query,
        "limit": limit,
        "fields": "title,authors,year,abstract,citationCount,url,openAccessPdf,externalIds"
    }
    url = f"{base_url}?{urllib.parse.urlencode(params)}"
    
    req = urllib.request.Request(url, headers={"User-Agent": "ResearchPaperSearch/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode())
    except Exception as e:
        return [{"error": f"Semantic Scholar API error: {e}"}]
    
    results = []
    for paper in data.get("data", []):
        authors = [a.get("name", "") for a in paper.get("authors", [])]
        result = {
            "source": "semantic_scholar",
            "id": f"semantic:{paper.get('paperId', '')}",
            "title": paper.get("title", ""),
            "authors": authors[:5],  # Limit authors shown
            "year": paper.get("year"),
            "abstract": (paper.get("abstract") or "")[:500],
            "citations": paper.get("citationCount", 0),
            "url": paper.get("url", ""),
            "isOpenAccess": paper.get("openAccessPdf") is not None,
            "doi": paper.get("externalIds", {}).get("DOI")
        }
        results.append(result)
    return results

def search_arxiv(query: str, limit: int = 10) -> list:
    """Search arXiv API."""
    base_url = "http://export.arxiv.org/api/query"
    params = {
        "search_query": f"all:{query}",
        "start": 0,
        "max_results": limit,
        "sortBy": "relevance"
    }
    url = f"{base_url}?{urllib.parse.urlencode(params)}"
    
    req = urllib.request.Request(url, headers={"User-Agent": "ResearchPaperSearch/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            xml_data = resp.read().decode()
    except Exception as e:
        return [{"error": f"arXiv API error: {e}"}]
    
    # Parse Atom XML
    ns = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}
    root = ET.fromstring(xml_data)
    
    results = []
    for entry in root.findall("atom:entry", ns):
        arxiv_id = entry.find("atom:id", ns).text.split("/abs/")[-1]
        authors = [a.find("atom:name", ns).text for a in entry.findall("atom:author", ns)]
        published = entry.find("atom:published", ns).text[:4] if entry.find("atom:published", ns) is not None else None
        
        result = {
            "source": "arxiv",
            "id": f"arxiv:{arxiv_id}",
            "title": entry.find("atom:title", ns).text.strip().replace("\n", " "),
            "authors": authors[:5],
            "year": int(published) if published else None,
            "abstract": (entry.find("atom:summary", ns).text or "").strip()[:500],
            "citations": None,  # arXiv doesn't provide citations
            "url": f"https://arxiv.org/abs/{arxiv_id}",
            "isOpenAccess": True,  # arXiv is always open access
            "pdf": f"https://arxiv.org/pdf/{arxiv_id}.pdf"
        }
        results.append(result)
    return results

def search_pubmed(query: str, limit: int = 10) -> list:
    """Search PubMed via NCBI E-utilities."""
    # Step 1: Search for IDs
    search_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
    search_params = {
        "db": "pubmed",
        "term": query,
        "retmax": limit,
        "retmode": "json",
        "sort": "relevance"
    }
    url = f"{search_url}?{urllib.parse.urlencode(search_params)}"
    
    req = urllib.request.Request(url, headers={"User-Agent": "ResearchPaperSearch/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode())
    except Exception as e:
        return [{"error": f"PubMed search error: {e}"}]
    
    ids = data.get("esearchresult", {}).get("idlist", [])
    if not ids:
        return []
    
    # Step 2: Fetch paper details
    fetch_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
    fetch_params = {
        "db": "pubmed",
        "id": ",".join(ids),
        "retmode": "xml"
    }
    url = f"{fetch_url}?{urllib.parse.urlencode(fetch_params)}"
    
    req = urllib.request.Request(url, headers={"User-Agent": "ResearchPaperSearch/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            xml_data = resp.read().decode()
    except Exception as e:
        return [{"error": f"PubMed fetch error: {e}"}]
    
    root = ET.fromstring(xml_data)
    results = []
    
    for article in root.findall(".//PubmedArticle"):
        pmid = article.find(".//PMID").text if article.find(".//PMID") is not None else ""
        title_elem = article.find(".//ArticleTitle")
        title = "".join(title_elem.itertext()) if title_elem is not None else ""
        
        # Get authors
        authors = []
        for author in article.findall(".//Author"):
            last = author.find("LastName")
            first = author.find("ForeName")
            if last is not None:
                name = last.text
                if first is not None:
                    name = f"{first.text} {last.text}"
                authors.append(name)
        
        # Get year
        year = None
        pub_date = article.find(".//PubDate/Year")
        if pub_date is not None:
            year = int(pub_date.text)
        
        # Get abstract
        abstract_parts = article.findall(".//AbstractText")
        abstract = " ".join("".join(p.itertext()) for p in abstract_parts)[:500]
        
        # Get DOI
        doi = None
        for id_elem in article.findall(".//ArticleId"):
            if id_elem.get("IdType") == "doi":
                doi = id_elem.text
        
        result = {
            "source": "pubmed",
            "id": f"pmid:{pmid}",
            "title": title,
            "authors": authors[:5],
            "year": year,
            "abstract": abstract,
            "citations": None,
            "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/",
            "isOpenAccess": None,  # Would need PMC lookup
            "doi": doi
        }
        results.append(result)
    
    return results

def main():
    parser = argparse.ArgumentParser(description="Search academic databases")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--source", choices=["all", "semantic", "arxiv", "pubmed"], 
                        default="all", help="Database to search")
    parser.add_argument("--limit", type=int, default=10, help="Max results per source")
    parser.add_argument("--json", action="store_true", help="Output raw JSON")
    args = parser.parse_args()
    
    results = []
    
    if args.source in ["all", "semantic"]:
        results.extend(search_semantic_scholar(args.query, args.limit))
    if args.source in ["all", "arxiv"]:
        results.extend(search_arxiv(args.query, args.limit))
    if args.source in ["all", "pubmed"]:
        results.extend(search_pubmed(args.query, args.limit))
    
    if args.json:
        print(json.dumps(results, indent=2))
        return
    
    # Pretty print
    for i, paper in enumerate(results, 1):
        if "error" in paper:
            print(f"\n⚠️  {paper['error']}")
            continue
        
        print(f"\n{'='*60}")
        print(f"[{i}] {paper['title']}")
        print(f"    📚 {paper['source'].upper()} | {paper['id']}")
        if paper.get('authors'):
            print(f"    👤 {', '.join(paper['authors'][:3])}" + (" et al." if len(paper.get('authors', [])) > 3 else ""))
        if paper.get('year'):
            print(f"    📅 {paper['year']}")
        if paper.get('citations'):
            print(f"    📊 {paper['citations']} citations")
        if paper.get('isOpenAccess'):
            print(f"    🔓 Open Access")
        if paper.get('abstract'):
            print(f"    📝 {paper['abstract'][:200]}...")
        print(f"    🔗 {paper['url']}")
    
    print(f"\n{'='*60}")
    print(f"Found {len([r for r in results if 'error' not in r])} papers")

if __name__ == "__main__":
    main()
