#!/usr/bin/env python3
"""
Tool: ncbi_entrez_batch_fetcher.py
Category: Public Databases & Automated Retrieval
Description: Batch downloads GenBank, nucleotide FASTA, and PubMed abstract metadata
             directly from the NCBI E-utilities REST API with rate limiting and error handling.
Usage:
    python ncbi_entrez_batch_fetcher.py --db nuccore --ids NM_000518.5,NM_000558.5 -o sequences.fasta
    python ncbi_entrez_batch_fetcher.py --db pubmed --term "optogenetics ion channel 2026" --max 5
"""

import argparse
import time
import urllib.request
import urllib.parse
import xml.etree.ElementTree as ET
import sys

BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"

def esearch(db: str, term: str, retmax: int = 10, email: str = "user@example.com") -> list:
    """Queries NCBI Entrez ESearch and returns list of UIDs."""
    params = urllib.parse.urlencode({
        'db': db,
        'term': term,
        'retmax': retmax,
        'retmode': 'xml',
        'email': email
    })
    url = f"{BASE_URL}/esearch.fcgi?{params}"
    req = urllib.request.Request(url, headers={'User-Agent': 'OmicsPortal-Script/1.0'})
    
    with urllib.request.urlopen(req, timeout=20) as resp:
        xml_data = resp.read()
        root = ET.fromstring(xml_data)
        ids = [id_elem.text for id_elem in root.findall(".//IdList/Id")]
        return ids

def efetch(db: str, id_list: list, rettype: str = "fasta", retmode: str = "text", email: str = "user@example.com") -> str:
    """Fetches full records for a list of UIDs."""
    params = urllib.parse.urlencode({
        'db': db,
        'id': ",".join(id_list),
        'rettype': rettype,
        'retmode': retmode,
        'email': email
    })
    url = f"{BASE_URL}/efetch.fcgi?{params}"
    req = urllib.request.Request(url, headers={'User-Agent': 'OmicsPortal-Script/1.0'})
    
    with urllib.request.urlopen(req, timeout=30) as resp:
        return resp.read().decode('utf-8', errors='ignore')

def main():
    parser = argparse.ArgumentParser(description="NCBI Entrez Batch Automated Downloader.")
    parser.add_argument("--db", default="nuccore", choices=["nuccore", "protein", "pubmed"], help="Target NCBI database")
    parser.add_argument("--ids", help="Comma-separated accession IDs")
    parser.add_argument("--term", help="Search query string")
    parser.add_argument("--max", type=int, default=3, help="Maximum records to retrieve")
    parser.add_argument("-o", "--output", help="Optional output file to save results")
    args = parser.parse_args()

    print("=== NCBI Entrez Automated Batch Downloader ===")
    
    id_list = []
    if args.ids:
        id_list = [x.strip() for x in args.ids.split(',') if x.strip()]
    elif args.term:
        print(f"Searching {args.db} for '{args.term}'...")
        id_list = esearch(args.db, args.term, retmax=args.max)
        print(f"Found {len(id_list)} matching IDs: {id_list}")
    else:
        # Default demo accession: Human Beta-Globin (HBB) NM_000518.5
        id_list = ["NM_000518.5"]
        print(f"No query provided. Fetching demo accession: {id_list[0]}\n")

    if not id_list:
        print("No IDs to fetch.", file=sys.stderr)
        return

    rettype = "fasta" if args.db in ("nuccore", "protein") else "abstract"
    print(f"Fetching {len(id_list)} records from {args.db} (type={rettype})...")
    
    try:
        content = efetch(args.db, id_list, rettype=rettype, retmode="text")
        if args.output:
            with open(args.output, "w") as f:
                f.write(content)
            print(f"[Successfully saved records to {args.output}]")
        else:
            print("\nPreview of Fetched Data:\n" + content[:400] + "...")
    except Exception as e:
        print(f"[Network/NCBI Error]: {e}", file=sys.stderr)

if __name__ == "__main__":
    main()