import ads
import json
import argparse
import sys

# 如果你没有在环境变量里设置 ADS_DEV_KEY，将使用以下的硬编码 Token 
ads.config.token = "dpJWki7eHJ48TwlKz2AUyhXAxBgZrKo6AjE8hZwp" 

def main():
    parser = argparse.ArgumentParser(description="Search ADS and return metadata")
    parser.add_argument("--query", required=True, help="ADS Search Query")
    parser.add_argument("--output", required=True, help="Output JSON file path")
    parser.add_argument("--rows", type=int, default=10, help="Number of rows to return")
    parser.add_argument("--year_range", help="Year range to filter, e.g. 2018-2023 or 2020")
    
    args = parser.parse_args()
    
    print(f"Searching ADS for query: {args.query}")
    query_params = {
        "q": args.query,
        "rows": args.rows,
        "fl": ["bibcode", "title", "author", "year", "abstract", "citation_count", "reference_count", "pub", "doi"]
    }
    
    if args.year_range:
        if '-' in args.year_range:
            start_year, end_year = args.year_range.split('-')
            query_params["fq"] = f"year:[{start_year} TO {end_year}]"
        else:
            query_params["fq"] = f"year:{args.year_range}"
            
    try:
        papers = list(ads.SearchQuery(**query_params))
        results = []
        for p in papers:
            record = {
                "bibcode": getattr(p, "bibcode", "") or "",
                "title": getattr(p, "title", [""])[0] if getattr(p, "title", None) else "",
                "author": getattr(p, "author", []),
                "year": getattr(p, "year", "") or "",
                "abstract": getattr(p, "abstract", "") or "",
                "citation_count": getattr(p, "citation_count", 0) or 0,
                "reference_count": getattr(p, "reference_count", 0) or 0,
                "pub": getattr(p, "pub", "") or "",
                "doi": getattr(p, "doi", [""])[0] if getattr(p, "doi", None) else ""
            }
            results.append(record)
            
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
            
        print(f"Found {len(results)} papers. Saved metadata to {args.output}.")
        
        # 打印简单摘要到终端
        for i, r in enumerate(results[:5]):
            print(f"\n[{i+1}] {r['title']} ({r['year']})")
            print(f"    Bibcode: {r['bibcode']} | Citations: {r['citation_count']}")
            authors = ", ".join(r['author'][:3]) + (" et al." if len(r['author']) > 3 else "")
            print(f"    Authors: {authors}")
            
    except Exception as e:
        print(f"Query Failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()