Ncbi Sequence Fetch

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""NCBI E-utilities wrapper for protein and nucleotide sequence retrieval.

Provides subcommands for the core NCBI E-utilities operations needed for
protein sequence retrieval: efetch, esearch, elink, CDS translation,
patent sequence search, and gene-to-protein lookup.

Rate-limited to 3 requests/second (10/s with NCBI_API_KEY).
"""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
#   "python-dotenv",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from typing import Any
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

import dotenv
from science_skills.skills.scienceskillscommon import http_client

_CODON_TABLE = {
    'ATA': 'I',
    'ATC': 'I',
    'ATT': 'I',
    'ATG': 'M',
    'ACA': 'T',
    'ACC': 'T',
    'ACG': 'T',
    'ACT': 'T',
    'AAC': 'N',
    'AAT': 'N',
    'AAA': 'K',
    'AAG': 'K',
    'AGC': 'S',
    'AGT': 'S',
    'AGA': 'R',
    'AGG': 'R',
    'CTA': 'L',
    'CTC': 'L',
    'CTG': 'L',
    'CTT': 'L',
    'CCA': 'P',
    'CCC': 'P',
    'CCG': 'P',
    'CCT': 'P',
    'CAC': 'H',
    'CAT': 'H',
    'CAA': 'Q',
    'CAG': 'Q',
    'CGA': 'R',
    'CGC': 'R',
    'CGG': 'R',
    'CGT': 'R',
    'GTA': 'V',
    'GTC': 'V',
    'GTG': 'V',
    'GTT': 'V',
    'GCA': 'A',
    'GCC': 'A',
    'GCG': 'A',
    'GCT': 'A',
    'GAC': 'D',
    'GAT': 'D',
    'GAA': 'E',
    'GAG': 'E',
    'GGA': 'G',
    'GGC': 'G',
    'GGG': 'G',
    'GGT': 'G',
    'TCA': 'S',
    'TCC': 'S',
    'TCG': 'S',
    'TCT': 'S',
    'TTC': 'F',
    'TTT': 'F',
    'TTA': 'L',
    'TTG': 'L',
    'TAC': 'Y',
    'TAT': 'Y',
    'TAA': '*',
    'TAG': '*',
    'TGC': 'C',
    'TGT': 'C',
    'TGA': '*',
    'TGG': 'W',
}


EUTILS_BASE = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils'

_CLIENT = None


def get_api_client():
  """Returns the lazily initialized HttpClient."""
  global _CLIENT
  if _CLIENT is None:
    api_key = os.environ.get('NCBI_API_KEY', '')
    qps = 10.0 if api_key else 3.0
    _CLIENT = http_client.HttpClient(EUTILS_BASE, qps=qps)
  return _CLIENT


def _eutils_get(endpoint: str, params: dict[str, str | int]) -> str | None:
  """Sends a GET request to an NCBI E-utilities endpoint.

  Handles rate limiting, retries with backoff, and API key
  injection.

  Args:
    endpoint: E-utilities endpoint name (e.g. 'efetch.fcgi').
    params: Query parameters as a dict.

  Returns:
    Response text or None on failure.
  """
  api_key = os.environ.get('NCBI_API_KEY', '')
  if api_key:
    params['api_key'] = api_key

  query_string = urllib.parse.urlencode(params)
  full_url = f'{EUTILS_BASE}/{endpoint}'
  if query_string:
    full_url += f'?{query_string}'

  try:
    return get_api_client().fetch_text(full_url)
  except http_client.HttpError as e:
    print(f'{endpoint} error after all retires: {e}', file=sys.stderr)
    return None


def efetch(
    db: str,
    db_id: str | int,
    retmode: str = 'text',
    rettype: str = 'fasta',
) -> str | None:
  """Fetches data from NCBI efetch endpoint.

  Args:
    db: Database name (protein, nuccore, gene, pubmed, etc.)
    db_id: One or more comma-separated IDs.
    retmode: Return mode (text, xml, json).
    rettype: Return type (fasta, gb, gp, fasta_cds_aa, etc.)

  Returns:
    Response text or None on failure.
  """
  return _eutils_get(
      'efetch.fcgi',
      {
          'db': db,
          'id': str(db_id),
          'rettype': rettype,
          'retmode': retmode,
      },
  )


def esearch(
    db: str,
    term: str,
    retmax: int = 20,
) -> tuple[list[str], int]:
  """Searches NCBI esearch endpoint.

  Args:
    db: Database name (protein, nuccore, gene, pubmed, etc.)
    term: Search query using Entrez syntax.
    retmax: Maximum results to return.

  Returns:
    Tuple of (list of ID strings, total count).
  """
  text = _eutils_get(
      'esearch.fcgi',
      {
          'db': db,
          'term': term,
          'retmode': 'json',
          'retmax': retmax,
      },
  )
  if text:
    try:
      data = json.loads(text)
      result = data.get('esearchresult', {})
      return result.get('idlist', []), int(result.get('count', 0))
    except (json.JSONDecodeError, ValueError) as e:
      print(f'esearch parse error: {e}', file=sys.stderr)
  return [], 0


def elink(dbfrom: str, db: str, link_id: str) -> list[str]:
  """Follows cross-database links via NCBI elink.

  Args:
    dbfrom: Source database (e.g., pubmed).
    db: Target database (e.g., protein).
    link_id: ID in the source database.

  Returns:
    List of linked IDs in the target database.
  """
  text = _eutils_get(
      'elink.fcgi',
      {
          'dbfrom': dbfrom,
          'db': db,
          'id': link_id,
          'retmode': 'json',
      },
  )
  if text:
    try:
      data = json.loads(text)
      linksets = data.get('linksets', [])
      if linksets:
        linksetdbs = linksets[0].get('linksetdbs', [])
        if linksetdbs:
          return linksetdbs[0].get('links', [])
    except (json.JSONDecodeError, ValueError) as e:
      print(f'elink parse error: {e}', file=sys.stderr)
  return []


def esummary(db: str, ids: str | list[str]) -> dict[str, Any]:
  """Fetches summary info for one or more IDs.

  Args:
    db: Database name.
    ids: Comma-separated ID string or list.

  Returns:
    Dict of id -> summary info.
  """
  if isinstance(ids, list):
    ids = ','.join(str(i) for i in ids)
  text = _eutils_get(
      'esummary.fcgi',
      {
          'db': db,
          'id': ids,
          'retmode': 'json',
      },
  )
  if text:
    try:
      data = json.loads(text)
      return data.get('result', {})
    except (json.JSONDecodeError, ValueError) as e:
      print(f'esummary parse error: {e}', file=sys.stderr)
  return {}


def parse_fasta(text: str | None) -> list[tuple[str, str]]:
  """Parses FASTA text into list of (header, sequence) tuples."""
  if not text:
    return []
  entries = []
  current_header = None
  current_seq = []
  for line in text.strip().splitlines():
    line = line.strip()
    if line.startswith('>'):
      if current_header is not None:
        entries.append((current_header, ''.join(current_seq)))
      current_header = line
      current_seq = []
    else:
      current_seq.append(line)
  if current_header is not None:
    entries.append((current_header, ''.join(current_seq)))
  return entries


def translate_dna(seq: str) -> str:
  """Translates a DNA sequence to protein using the standard codon table."""
  protein = []
  seq = seq.upper().replace('\n', '').replace(' ', '')
  for i in range(0, len(seq) - len(seq) % 3, 3):
    codon = seq[i : i + 3]
    if 'N' in codon:
      protein.append('X')
    else:
      aa = _CODON_TABLE.get(codon, 'X')
      protein.append(aa)
  return ''.join(protein).rstrip('*')


def get_longest_orf(seq: str, min_len: int = 75) -> str | None:
  """Finds the longest open reading frame in a DNA sequence.

  Searches all 6 reading frames (3 forward + 3 reverse complement).

  Args:
    seq: DNA sequence string.
    min_len: Minimum ORF length in amino acids.

  Returns:
    Longest ORF protein sequence, or None.
  """
  rev_comp = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', 'N': 'N'}
  rev_seq = ''.join(rev_comp.get(b, 'N') for b in reversed(seq.upper()))

  longest = ''
  for s in [seq.upper(), rev_seq]:
    for frame in range(3):
      translated = translate_dna(s[frame:])
      # Split on stop codons
      peptides = translated.split('*')
      for pep in peptides:
        idx = pep.find('M')
        if idx != -1:
          orf = pep[idx:]
          if len(orf) >= min_len and len(orf) > len(longest):
            longest = orf
  return longest if longest else None


def extract_cds_translations(
    xml_text: str | None,
    identifier: str | None = None,
    target_len: int = 0,
) -> list[dict[str, str]]:
  """Extracts CDS translations from GenBank XML.

  Args:
    xml_text: GenBank XML string from efetch.
    identifier: Optional gene name to filter by.
    target_len: Optional target protein length for best-match selection.

  Returns:
    List of dicts with keys: translation, product,
    gene, locus_tag, protein_id
  """
  del identifier, target_len  # Reserved for future filtering.
  if not xml_text:
    return []
  results = []
  try:
    root = ET.fromstring(xml_text)
    for seq in root.findall('.//GBSeq'):
      for feat in seq.findall('.//GBFeature'):
        key = feat.find('GBFeature_key')
        if key is None or key.text != 'CDS':
          continue
        entry = {
            'translation': '',
            'product': '',
            'gene': '',
            'locus_tag': '',
            'protein_id': '',
            'note': '',
        }
        for qual in feat.findall('.//GBQualifier'):
          name = qual.find('GBQualifier_name')
          val = qual.find('GBQualifier_value')
          if name is not None and val is not None and val.text:
            if name.text == 'translation':
              entry['translation'] = val.text.upper()
            elif name.text in [
                'product',
                'gene',
                'locus_tag',
                'protein_id',
                'note',
            ]:
              entry[name.text] = val.text
        if entry['translation']:
          results.append(entry)
  except ET.ParseError as e:
    print(f'XML parse error: {e}', file=sys.stderr)
  return results


def is_genomic_record(xml_text: str) -> bool:
  """Checks if a GenBank XML record is genomic DNA (not mRNA/CDS)."""
  try:
    root = ET.fromstring(xml_text)
    moltype = root.find('.//GBSeq_moltype')
    if moltype is not None and moltype.text:
      if 'genomic' in moltype.text.lower():
        return True
    for feat in root.iter('GBFeature'):
      k = feat.find('GBFeature_key')
      if k is not None and k.text == 'source':
        for qual in feat.iter('GBQualifier'):
          n = qual.find('GBQualifier_name')
          v = qual.find('GBQualifier_value')
          if (
              n is not None
              and n.text == 'mol_type'
              and v is not None
              and v.text
              and 'genomic' in v.text.lower()
          ):
            return True
        break
  except ET.ParseError:
    pass
  return False


def cmd_fetch_protein(args: argparse.Namespace) -> None:
  """Fetches protein sequence by accession ID."""
  results = []
  for acc in args.accession:
    fasta = efetch('protein', acc)
    entries = parse_fasta(fasta) if fasta else []
    if entries:
      for header, seq in entries:
        results.append({
            'accession': acc,
            'header': header,
            'sequence': seq,
            'length': len(seq),
        })
    else:
      # Try with dot variant
      if '_' in acc:
        fasta = efetch('protein', acc.replace('_', '.'))
        entries = parse_fasta(fasta) if fasta else []
        for header, seq in entries:
          results.append({
              'accession': acc,
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
      if not any(r['accession'] == acc for r in results):
        results.append({
            'accession': acc,
            'error': 'Not found',
        })
  _write_output(results, args.output)


def cmd_fetch_nucleotide(args: argparse.Namespace) -> None:
  """Fetches nucleotide sequence by accession ID."""
  results = []
  for acc in args.accession:
    fasta = efetch('nuccore', acc)
    entries = parse_fasta(fasta) if fasta else []
    if entries:
      for header, seq in entries:
        results.append({
            'accession': acc,
            'header': header,
            'sequence': seq,
            'length': len(seq),
        })
    else:
      results.append({'accession': acc, 'error': 'Not found'})
  _write_output(results, args.output)


def cmd_cds_translate(args: argparse.Namespace) -> None:
  """Fetches CDS and translates to protein sequence.

  Tries multiple approaches:
  1. NCBI fasta_cds_aa (pre-translated CDS protein)
  2. GenBank XML CDS translation annotations
  3. Raw nucleotide to 6-frame ORF finding (if not genomic)

  Args:
    args: Parsed CLI args with accession, target_length, and output attributes.
  """
  results = []
  for acc in args.accession:
    result = {'accession': acc}

    # Approach 1: Pre-translated CDS protein
    fasta = efetch('nuccore', acc, rettype='fasta_cds_aa')
    if fasta:
      entries = parse_fasta(fasta)
      if entries:
        best = _pick_best_by_length(entries, args.target_length)
        if best:
          header, seq = best
          seq = seq.replace('*', '')
          # Check for dominant isoform if multiple entries
          if len(entries) > 1 and not args.target_length:
            dom = _get_dominant_cds(acc)
            if dom:
              seq = dom
              header = f'>dominant_isoform_{acc}'
          result['header'] = header
          result['sequence'] = seq
          result['length'] = len(seq)
          result['method'] = 'fasta_cds_aa'
          results.append(result)
          continue

    # Approach 2: GenBank XML CDS translations
    xml_data = efetch('nuccore', acc, retmode='xml', rettype='gb')
    if xml_data:
      genomic = is_genomic_record(xml_data)
      cds_list = extract_cds_translations(xml_data)
      if cds_list:
        if args.target_length:
          best_cds = min(
              cds_list,
              key=lambda c: abs(len(c['translation']) - args.target_length),
          )
        else:
          best_cds = max(
              cds_list,
              key=lambda c: len(c['translation']),
          )
        result['header'] = (
            f'>{best_cds["gene"]}'
            f'|{best_cds["product"]}'
            f'|{best_cds["protein_id"]}'
        )
        result['sequence'] = best_cds['translation']
        result['length'] = len(best_cds['translation'])
        result['method'] = 'genbank_xml_cds'
        result['cds_info'] = best_cds
        if genomic:
          result['is_genomic'] = True
        results.append(result)
        continue
      if genomic:
        result['is_genomic'] = True

    # Approach 3: Raw nucleotide ORF finding
    fasta = efetch('nuccore', acc)
    if fasta:
      entries = parse_fasta(fasta)
      if entries:
        _, dna_seq = entries[0]
        orf = get_longest_orf(dna_seq)
        if orf:
          tl = args.target_length
          if tl == 0 or abs(len(orf) - tl) < 50:
            result['header'] = f'>ORF_{acc}'
            result['sequence'] = orf
            result['length'] = len(orf)
            result['method'] = 'orf_translation'
            results.append(result)
            continue

    result['error'] = 'No CDS translation found'
    results.append(result)

  _write_output(results, args.output)


def cmd_search(args: argparse.Namespace) -> None:
  """Searches an NCBI database by query terms."""
  ids, count = esearch(args.database, args.query, retmax=args.retmax)
  result = {
      'database': args.database,
      'query': args.query,
      'total_count': count,
      'returned_count': len(ids),
      'ids': ids,
  }

  # Optionally fetch FASTA for returned IDs
  if args.fetch_sequences and ids:
    sequences = []
    batch_size = 20
    for i in range(0, len(ids), batch_size):
      batch = ids[i : i + batch_size]
      fasta = efetch(args.database, ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          sequences.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
    result['sequences'] = sequences
  _write_output(result, args.output)


def cmd_elink(args: argparse.Namespace) -> None:
  """Follows cross-database links from one NCBI database to another."""
  linked_ids = elink(args.dbfrom, args.db, args.id)
  result = {
      'source_db': args.dbfrom,
      'target_db': args.db,
      'source_id': args.id,
      'linked_ids': linked_ids,
      'count': len(linked_ids),
  }

  if args.fetch_sequences and linked_ids:
    sequences = []
    batch_size = 20
    for i in range(0, len(linked_ids), batch_size):
      batch = linked_ids[i : i + batch_size]
      fasta = efetch(args.db, ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          sequences.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })
    result['sequences'] = sequences

  _write_output(result, args.output)


def cmd_gene_protein(args: argparse.Namespace) -> None:
  """Searches for protein sequence by gene name and organism.

  Tries NCBI Protein database with gene name + organism
  filter. Optionally filters by sequence length.

  Args:
    args: Parsed CLI args with gene, organism, target_length, retmax, and output
      attributes.
  """
  gene = args.gene
  organism = args.organism
  results = []

  # Search NCBI Protein
  query = f'{gene}[Gene Name] AND {organism}[Organism]'
  if args.target_length:
    lo = args.target_length - 25
    hi = args.target_length + 25
    query += f' AND {lo}:{hi}[Sequence Length]'

  ids, count = esearch('protein', query, retmax=args.retmax)
  if ids:
    for uid in ids[: args.retmax]:
      fasta = efetch('protein', uid)
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          results.append({
              'uid': uid,
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })

  output = {
      'gene': gene,
      'organism': organism,
      'query': query,
      'total_count': count,
      'results': results,
  }
  _write_output(output, args.output)


def cmd_locus_protein(args: argparse.Namespace) -> None:
  """Searches for protein sequence by locus tag.

  Tries:
  1. NCBI Protein search by gene name + organism
  2. NCBI Nuccore search for GenBank records with CDS
     translations

  Args:
    args: Parsed CLI args with locus, organism, retmax, and output attributes.
  """
  locus = args.locus
  organism = args.organism
  results = []

  # Try NCBI Protein
  query = f'{locus}[Gene Name]'
  if organism:
    query += f' AND {organism}[Organism]'
  ids, _ = esearch('protein', query, retmax=args.retmax)
  if ids:
    for uid in ids[: args.retmax]:
      fasta = efetch('protein', uid)
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          if len(seq) >= 50:
            results.append({
                'uid': uid,
                'header': header,
                'sequence': seq,
                'length': len(seq),
                'source': 'protein_db',
            })

  # Try NCBI Nuccore for CDS translations
  nuc_query = f'{locus}[Accession] OR {locus}[All Fields]'
  nuc_ids, _ = esearch('nuccore', nuc_query, retmax=10)
  if nuc_ids:
    xml_data = efetch(
        'nuccore',
        ','.join(nuc_ids),
        retmode='xml',
        rettype='gb',
    )
    if xml_data:
      cds_list = extract_cds_translations(xml_data, identifier=locus)
      for cds in cds_list:
        results.append({
            'header': f'>{cds["gene"]}|{cds["product"]}|{cds["protein_id"]}',
            'sequence': cds['translation'],
            'length': len(cds['translation']),
            'source': 'nuccore_cds',
            'gene': cds['gene'],
            'locus_tag': cds['locus_tag'],
        })

  output = {
      'locus': locus,
      'organism': organism or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_pubmed_proteins(args: argparse.Namespace) -> None:
  """Finds protein sequences linked to a PubMed article.

  Searches:
  1. NCBI Protein records with PMID filter
  2. NCBI elink PubMed to Protein
  3. NCBI Nuccore records with PMID for CDS
     translations

  Args:
    args: Parsed CLI args with pmid, identifier, and output attributes.
  """
  pmid = args.pmid
  identifier = args.identifier
  results = []

  # Step 1: Protein db search by PMID
  prot_ids, _ = esearch('protein', f'{pmid}[PMID]', retmax=100)

  # Step 2: elink PubMed→Protein if no direct hits
  if not prot_ids:
    prot_ids = elink('pubmed', 'protein', pmid)

  if prot_ids:
    batch_size = 20
    for i in range(0, len(prot_ids), batch_size):
      batch = prot_ids[i : i + batch_size]
      xml_data = efetch('protein', ','.join(batch), retmode='xml', rettype='gp')
      if xml_data:
        try:
          root = ET.fromstring(xml_data)
          # NCBI returns GBSeq or INSDSeq depending on format
          seq_els = root.findall('.//GBSeq')
          pfx = 'GBSeq'
          if not seq_els:
            seq_els = root.findall('.//INSDSeq')
            pfx = 'INSDSeq'
          for gb_seq in seq_els:
            defn = gb_seq.find(f'{pfx}_definition')
            acc = gb_seq.find(f'{pfx}_primary-accession')
            seq_el = gb_seq.find(f'{pfx}_sequence')
            if seq_el is not None and seq_el.text:
              defn_text = (defn.text or '') if defn is not None else ''
              entry = {
                  'accession': acc.text if acc is not None else '',
                  'definition': defn_text,
                  'sequence': seq_el.text.upper(),
                  'length': len(seq_el.text),
                  'source': 'protein_db',
              }
              if identifier:
                entry['matches_identifier'] = (
                    identifier.lower() in defn_text.lower()
                )
              results.append(entry)
        except (ET.ParseError, AttributeError):
          pass

  # Step 3: Nuccore records for CDS translations
  nuc_ids, _ = esearch('nuccore', f'{pmid}[PMID]', retmax=100)
  if nuc_ids:
    for i in range(0, len(nuc_ids), 10):
      batch = nuc_ids[i : i + 10]
      xml_data = efetch('nuccore', ','.join(batch), retmode='xml', rettype='gb')
      if xml_data:
        cds_list = extract_cds_translations(xml_data, identifier=identifier)
        for cds in cds_list:
          entry = {
              'gene': cds['gene'],
              'product': cds['product'],
              'protein_id': cds['protein_id'],
              'sequence': cds['translation'],
              'length': len(cds['translation']),
              'source': 'nuccore_cds',
          }
          if identifier:
            entry['matches_identifier'] = (
                identifier.lower() in cds['gene'].lower()
                or identifier.lower() in cds['product'].lower()
                or identifier.lower() in cds['protein_id'].lower()
            )
          results.append(entry)

  output = {
      'pmid': pmid,
      'identifier': identifier or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_patent_search(args: argparse.Namespace) -> None:
  """Searches for protein sequences from patents.

  Two modes:
  1. By patent number: fetches all protein sequences
     from a specific patent
  2. By keywords: searches NCBI Protein with
     patent[Properties] filter

  Args:
    args: Parsed CLI args with patent_number, keywords, organism, retmax, and
      output attributes.
  """
  results = []

  if args.patent_number:
    # Mode 1: Search by patent number
    patent = args.patent_number.strip()

    # Extract the core 7-10 digit number, ignoring prefixes like US/WO/EP
    # and suffixes like B2/A1
    digit_match = re.search(r'(\d{7,10})', patent)
    digits = digit_match.group(1) if digit_match else ''

    # Build queries in priority order: most specific first.
    # Avoid broad text searches like 'patent {patent}' which match millions
    # of unrelated records and pollute the ID list.
    queries = []
    if digits:
      queries.append(f'{digits}[Patent Publication Number]')

    # Extract country prefix (e.g. US, WO, EP) for exact phrase search
    prefix_match = re.match(r'^([A-Za-z]{2})\s*', patent)
    if prefix_match and digits:
      country = prefix_match.group(1).upper()
      queries.append(f'"{country} {digits}"')

    # Handle international patents
    if re.search(r'^(WO|EP)', patent, re.IGNORECASE):
      stripped = re.sub(r'^(WO|EP)', '', patent)
      queries.append(f'patent {stripped}')

    # Fallback: broad text search only if no specific queries matched
    if not queries:
      queries.append(f'patent {patent}')

    all_ids = []
    for q in queries:
      ids, _ = esearch('protein', q, retmax=args.retmax)
      all_ids.extend(ids)
      # If we already have enough IDs from a specific query, stop
      if len(set(all_ids)) >= args.retmax:
        break
    all_ids = list(dict.fromkeys(all_ids))
    all_ids = all_ids[: args.retmax]  # dedupe preserving order

    if all_ids:
      batch_size = 20
      for i in range(0, len(all_ids), batch_size):
        batch = all_ids[i : i + batch_size]
        fasta = efetch('protein', ','.join(batch))
        if fasta:
          entries = parse_fasta(fasta)
          for header, seq in entries:
            # Filter results that don't mention our patent's digits
            if digits and digits not in header:
              continue

            seq_id_match = re.search(r'Sequence (\d+)', header)
            pat_match = re.search(
                r'patent\s+(?:[a-zA-Z]{2}\s*)?(\d{7,10})',
                header,
                re.IGNORECASE,
            )
            results.append({
                'header': header,
                'sequence': seq,
                'length': len(seq),
                'seq_id_no': (
                    int(seq_id_match.group(1)) if seq_id_match else None
                ),
                'patent_number': pat_match.group(1) if pat_match else patent,
            })

  elif args.keywords:
    # Mode 2: Search by keywords with patent filter
    keywords = args.keywords
    query_parts = []
    for kw in keywords:
      query_parts.append(kw)
    query = ' AND '.join(query_parts)

    if args.organism:
      query += f' AND {args.organism}[Organism]'
    query += ' AND patent[Properties]'

    ids, _ = esearch('protein', query, retmax=args.retmax)
    if ids:
      batch_size = 20
      for i in range(0, len(ids), batch_size):
        batch = ids[i : i + batch_size]
        fasta = efetch('protein', ','.join(batch))
        if fasta:
          entries = parse_fasta(fasta)
          for header, seq in entries:
            pat_match = re.search(
                r'patent\s+US\s*(\d{7,8})',
                header,
                re.IGNORECASE,
            )
            if pat_match:
              results.append({
                  'header': header,
                  'sequence': seq,
                  'length': len(seq),
                  'patent_number': pat_match.group(1),
              })

  output = {
      'patent_number': args.patent_number or '',
      'keywords': args.keywords or [],
      'organism': args.organism or '',
      'total_results': len(results),
      'results': results,
  }
  _write_output(output, args.output)


def cmd_organism_length(args: argparse.Namespace) -> None:
  """Searches for protein by organism and exact length.

  Uses NCBI SLEN (Sequence Length) filter to find
  proteins of exact length from a given organism.
  Useful as a last-resort lookup when only organism
  and expected AA length are known.

  Args:
    args: Parsed CLI args with organism, length, retmax, and output attributes.
  """
  organism = args.organism
  length = args.length
  results = []

  query = f'"{organism}"[Organism] AND {length}[SLEN]'
  ids, count = esearch('protein', query, retmax=args.retmax)
  if ids:
    batch_size = 20
    for i in range(0, len(ids), batch_size):
      batch = ids[i : i + batch_size]
      fasta = efetch('protein', ','.join(batch))
      if fasta:
        entries = parse_fasta(fasta)
        for header, seq in entries:
          results.append({
              'header': header,
              'sequence': seq,
              'length': len(seq),
          })

  output = {
      'organism': organism,
      'target_length': length,
      'query': query,
      'total_count': count,
      'results': results,
  }
  _write_output(output, args.output)


def _pick_best_by_length(
    entries: list[tuple[str, str]],
    target_len: int,
) -> tuple[str, str] | None:
  """Picks the FASTA entry closest to target length, or longest if no target."""
  if not entries:
    return None
  valid = [(h, s) for h, s in entries if len(s) >= 50]
  if not valid:
    return None
  if target_len and target_len > 0:
    return min(valid, key=lambda x: abs(len(x[1]) - target_len))
  return max(valid, key=lambda x: len(x[1]))


def _get_dominant_cds(accession: str) -> str | None:
  """Returns the 'dominant isoform' CDS translation if annotated in GenBank."""
  xml_data = efetch('nuccore', accession, retmode='xml', rettype='gb')
  if not xml_data:
    return None
  cds_list = extract_cds_translations(xml_data)
  for cds in cds_list:
    if 'dominant' in cds.get('note', '').lower():
      return cds['translation']
  return None


def _write_output(
    data: dict[str, Any] | list[Any],
    output_path: str | None,
) -> None:
  """Writes results to JSON file and prints summary to stdout."""
  if output_path:
    with open(output_path, 'w') as f:
      json.dump(data, f, indent=2)
    print(f'Results written to: {output_path}')
  else:
    print(json.dumps(data, indent=2))

  # Print summary to stdout
  if isinstance(data, list):
    for item in data:
      if 'error' in item:
        print(f"  {item.get('accession', '?')}: ERROR — {item['error']}")
      elif 'sequence' in item:
        print(f"  {item.get('accession', '?')}: {item['length']} AA")
  elif isinstance(data, dict):
    if 'results' in data:
      n = len(data['results'])
      print(f'  Found {n} result(s)')
      for r in data['results'][:5]:
        if 'sequence' in r:
          desc = r.get('header', r.get('definition', ''))[:80]
          print(f"    {r['length']} AA — {desc}")
    elif 'sequences' in data:
      n = len(data['sequences'])
      print(f'  Found {n} sequence(s)')
    elif 'ids' in data:
      print(
          f'  {data.get("total_count", 0)} total,'
          f' {data.get("returned_count", 0)} returned'
      )


def main() -> None:
  dotenv.load_dotenv(os.path.expanduser('~/.env'))
  parser = argparse.ArgumentParser(
      description='NCBI E-utilities wrapper for sequence retrieval.',
  )
  sub = parser.add_subparsers(dest='command', help='Subcommand')

  # fetch-protein
  p = sub.add_parser('fetch-protein', help='Fetch protein by accession')
  p.add_argument('accession', nargs='+', help='Protein accession(s)')
  p.add_argument('--output', '-o', help='Output JSON file')

  # fetch-nucleotide
  p = sub.add_parser('fetch-nucleotide', help='Fetch nucleotide by accession')
  p.add_argument('accession', nargs='+', help='Nucleotide accession(s)')
  p.add_argument('--output', '-o', help='Output JSON file')

  # cds-translate
  p = sub.add_parser(
      'cds-translate',
      help='Fetch CDS accession and translate to protein',
  )
  p.add_argument('accession', nargs='+', help='CDS/nucleotide accession(s)')
  p.add_argument(
      '--target-length',
      type=int,
      default=0,
      help='Expected protein length (AA) for best-match selection',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # search
  p = sub.add_parser('search', help='Search an NCBI database')
  p.add_argument('query', help='Entrez search query')
  p.add_argument(
      '--database',
      '-d',
      default='protein',
      help='Database (protein, nuccore, gene, pubmed, etc.)',
  )
  p.add_argument(
      '--retmax', type=int, default=20, help='Maximum results to return'
  )
  p.add_argument(
      '--fetch-sequences',
      action='store_true',
      help='Also fetch FASTA sequences for returned IDs',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # elink
  p = sub.add_parser('elink', help='Follow cross-database links')
  p.add_argument('id', help='Source ID')
  p.add_argument(
      '--dbfrom', required=True, help='Source database (e.g., pubmed)'
  )
  p.add_argument('--db', required=True, help='Target database (e.g., protein)')
  p.add_argument(
      '--fetch-sequences',
      action='store_true',
      help='Also fetch FASTA sequences for linked IDs',
  )
  p.add_argument('--output', '-o', help='Output JSON file')

  # gene-protein
  p = sub.add_parser(
      'gene-protein',
      help='Search protein by gene name + organism',
  )
  p.add_argument('gene', help='Gene name')
  p.add_argument('--organism', required=True, help='Organism name')
  p.add_argument(
      '--target-length',
      type=int,
      default=0,
      help='Expected protein length (AA)',
  )
  p.add_argument('--retmax', type=int, default=10, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # locus-protein
  p = sub.add_parser(
      'locus-protein',
      help='Search protein by locus tag',
  )
  p.add_argument('locus', help='Locus tag or accession')
  p.add_argument('--organism', help='Organism name')
  p.add_argument('--retmax', type=int, default=10, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # pubmed-proteins
  p = sub.add_parser(
      'pubmed-proteins',
      help='Find protein sequences linked to a PubMed article',
  )
  p.add_argument('pmid', help='PubMed ID')
  p.add_argument('--identifier', help='Gene/protein name to match')
  p.add_argument('--output', '-o', help='Output JSON file')

  # patent-search
  p = sub.add_parser(
      'patent-search',
      help='Search for protein sequences from patents',
  )
  p.add_argument('--patent-number', help='Patent number (e.g., US12345678)')
  p.add_argument(
      '--keywords', nargs='+', help='Search keywords for patent protein search'
  )
  p.add_argument('--organism', help='Organism filter')
  p.add_argument('--retmax', type=int, default=50, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  # organism-length
  p = sub.add_parser(
      'organism-length',
      help='Search protein by organism and exact AA length',
  )
  p.add_argument('--organism', required=True, help='Organism name')
  p.add_argument(
      '--length', type=int, required=True, help='Exact protein length in AA'
  )
  p.add_argument('--retmax', type=int, default=50, help='Max results')
  p.add_argument('--output', '-o', help='Output JSON file')

  args = parser.parse_args()
  if not args.command:
    parser.print_help()
    sys.exit(1)

  cmd_map = {
      'fetch-protein': cmd_fetch_protein,
      'fetch-nucleotide': cmd_fetch_nucleotide,
      'cds-translate': cmd_cds_translate,
      'search': cmd_search,
      'elink': cmd_elink,
      'gene-protein': cmd_gene_protein,
      'locus-protein': cmd_locus_protein,
      'pubmed-proteins': cmd_pubmed_proteins,
      'patent-search': cmd_patent_search,
      'organism-length': cmd_organism_length,
  }
  cmd_map[args.command](args)


if __name__ == '__main__':
  main()