Skip to main content

argenus/
flanking_db.rs

1//! Flanking Database Builder Module
2//!
3//! Downloads prokaryotic genomes from NCBI and PLSDB, aligns ARG sequences,
4//! and extracts flanking sequences for genus classification.
5//!
6//! # Data Sources
7//! - **NCBI**: GenBank (bacteria + archaea), Complete Genome + Chromosome level
8//! - **PLSDB**: Standalone circular complete plasmids (ASSEMBLY_UID == -1)
9//!
10//! # Pipeline Overview
11//! 1. Download NCBI taxonomy database (taxdump)
12//! 2. Fetch prokaryotic genome assemblies
13//! 3. Download PLSDB plasmid sequences
14//! 4. Align ARG database to genomes (minimap2)
15//! 5. Extract flanking sequences with genus annotation
16//! 6. Build compressed FDB index
17//!
18//! # Output Files
19//! - `genome_catalogue.tsv`: Accession → genus mapping
20//! - `flanking_db.tsv`: Raw flanking sequences
21//! - `flanking.fdb`: Compressed indexed database
22//!
23//! # External Dependencies
24//! - minimap2 (alignment)
25
26use anyhow::{Context, Result};
27use rustc_hash::{FxHashMap, FxHashSet};
28use std::fs::File;
29use std::io::{BufRead, BufReader, BufWriter, Read, Write};
30use std::path::{Path, PathBuf};
31use std::process::Command;
32use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};
33use std::sync::{Arc, Mutex, mpsc};
34use std::time::Duration;
35use std::thread;
36
37const NCBI_FTP_BASE: &str = "https://ftp.ncbi.nlm.nih.gov/genomes";
38const NCBI_TAXDUMP_URL: &str = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz";
39const NCBI_DATASETS_API: &str = "https://api.ncbi.nlm.nih.gov/datasets/v2";
40const PLSDB_META_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_meta.tar.gz";
41const PLSDB_FASTA_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_fasta";
42// Default flanking length: 1000bp (configurable via -n flag)
43const API_BATCH_SIZE: usize = 1000; // NCBI recommends <1000 per request
44
45/// PLSDB options for flanking database build
46#[derive(Clone, Debug, Default)]
47pub struct PlsdbOptions {
48    /// Pre-downloaded PLSDB directory (contains meta.tar.gz and sequences.fasta)
49    pub dir: Option<PathBuf>,
50    /// Skip PLSDB processing entirely
51    pub skip: bool,
52}
53
54/// Configuration for flanking database build
55#[derive(Clone, Debug)]
56pub struct FlankBuildConfig {
57    /// Flanking region length in bp (upstream + downstream)
58    pub flanking_length: usize,
59    /// Queue buffer size in GB for backpressure control
60    pub queue_buffer_gb: u32,
61    /// PLSDB options
62    pub plsdb: PlsdbOptions,
63}
64
65impl Default for FlankBuildConfig {
66    fn default() -> Self {
67        Self {
68            flanking_length: 1000,
69            queue_buffer_gb: 30,
70            plsdb: PlsdbOptions::default(),
71        }
72    }
73}
74
75/// Assembly metadata from NCBI
76#[derive(Clone, Debug)]
77pub struct AssemblyInfo {
78    pub accession: String,
79    pub taxid: String,
80    pub species_taxid: String,
81    pub organism_name: String,
82}
83
84/// Plasmid metadata from PLSDB
85#[derive(Clone, Debug)]
86pub struct PlasmidInfo {
87    pub accession: String,
88    pub taxonomy_uid: String,
89    pub genus: String,
90    pub species: String,
91}
92
93/// Parsed PAF alignment hit for tie-breaking
94#[derive(Clone, Debug)]
95struct PafHit {
96    query_name: String,
97    query_start: usize,
98    query_end: usize,
99    gene_name: String,
100    gene_length: usize,
101    score: i64,
102    mapq: u8,
103    divergence: f32,
104    gap_count: usize,
105    raw_line: String,
106}
107
108impl PafHit {
109    fn from_paf_line(line: &str) -> Option<Self> {
110        let fields: Vec<&str> = line.split('\t').collect();
111        if fields.len() < 12 {
112            return None;
113        }
114
115        let query_name = fields[0].to_string();
116        let query_start: usize = fields[2].parse().ok()?;
117        let query_end: usize = fields[3].parse().ok()?;
118        let gene_name = fields[5].to_string();
119        let gene_length: usize = fields[6].parse().ok()?;
120        let mapq: u8 = fields[11].parse().unwrap_or(0);
121
122        let matches: usize = fields[9].parse().unwrap_or(0);
123        let block_len: usize = fields[10].parse().unwrap_or(0);
124
125        let mut score: i64 = 0;
126        let mut divergence: f32 = 1.0;
127
128        for field in &fields[12..] {
129            if let Some(val) = field.strip_prefix("AS:i:") {
130                score = val.parse().unwrap_or(0);
131            } else if let Some(val) = field.strip_prefix("de:f:") {
132                divergence = val.parse().unwrap_or(1.0);
133            }
134        }
135
136        let gap_count = block_len.saturating_sub(matches);
137
138        Some(PafHit {
139            query_name,
140            query_start,
141            query_end,
142            gene_name,
143            gene_length,
144            score,
145            mapq,
146            divergence,
147            gap_count,
148            raw_line: line.to_string(),
149        })
150    }
151
152    fn overlaps(&self, other: &PafHit) -> bool {
153        if self.query_name != other.query_name {
154            return false;
155        }
156        let start = self.query_start.max(other.query_start);
157        let end = self.query_end.min(other.query_end);
158        if start >= end {
159            return false;
160        }
161        let overlap = end - start;
162        let self_len = self.query_end - self.query_start;
163        let other_len = other.query_end - other.query_start;
164        overlap * 2 > self_len || overlap * 2 > other_len
165    }
166}
167
168/// Compare hits: Score↓ → Gene length↓ → MapQ↓ → Divergence↑ → Gap↑ → Name
169fn compare_paf_hits(a: &PafHit, b: &PafHit) -> std::cmp::Ordering {
170    use std::cmp::Ordering;
171    b.score.cmp(&a.score)
172        .then_with(|| b.gene_length.cmp(&a.gene_length))
173        .then_with(|| b.mapq.cmp(&a.mapq))
174        .then_with(|| a.divergence.partial_cmp(&b.divergence).unwrap_or(Ordering::Equal))
175        .then_with(|| a.gap_count.cmp(&b.gap_count))
176        .then_with(|| a.gene_name.cmp(&b.gene_name))
177}
178
179/// Deduplicate overlapping hits using tie-breaking rules
180fn deduplicate_paf_hits(hits: Vec<PafHit>) -> Vec<PafHit> {
181    if hits.is_empty() {
182        return hits;
183    }
184
185    let mut groups: Vec<Vec<PafHit>> = Vec::new();
186    for hit in hits {
187        let mut found = false;
188        for group in &mut groups {
189            if group.iter().any(|h| h.overlaps(&hit)) {
190                group.push(hit.clone());
191                found = true;
192                break;
193            }
194        }
195        if !found {
196            groups.push(vec![hit]);
197        }
198    }
199
200    groups
201        .into_iter()
202        .map(|mut g| {
203            g.sort_by(compare_paf_hits);
204            g.into_iter().next().unwrap()
205        })
206        .collect()
207}
208
209// ============================================================================
210// Adaptive Thread Bucketing
211// ============================================================================
212
213/// Genome size category for adaptive thread allocation
214/// Larger genomes benefit from more minimap2 threads
215#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
216pub enum GenomeSizeCategory {
217    /// Small genomes (<5MB): 4 threads for minimap2
218    Small,
219    /// Medium genomes (5-20MB): 6 threads for minimap2
220    Medium,
221    /// Large genomes (>20MB): 8 threads for minimap2
222    Large,
223}
224
225/// Size thresholds in bytes
226const SMALL_GENOME_THRESHOLD: u64 = 5 * 1024 * 1024;   // 5MB
227const LARGE_GENOME_THRESHOLD: u64 = 20 * 1024 * 1024;  // 20MB
228
229impl GenomeSizeCategory {
230    /// Determine category from file size in bytes
231    pub fn from_size(size_bytes: u64) -> Self {
232        if size_bytes < SMALL_GENOME_THRESHOLD {
233            GenomeSizeCategory::Small
234        } else if size_bytes < LARGE_GENOME_THRESHOLD {
235            GenomeSizeCategory::Medium
236        } else {
237            GenomeSizeCategory::Large
238        }
239    }
240
241    /// Get recommended minimap2 thread count for this category
242    pub fn thread_count(&self, max_threads: usize) -> usize {
243        let recommended = match self {
244            GenomeSizeCategory::Small => 4,
245            GenomeSizeCategory::Medium => 6,
246            GenomeSizeCategory::Large => 8,
247        };
248        // Don't exceed available threads
249        recommended.min(max_threads)
250    }
251
252    /// Get category name for logging
253    pub fn name(&self) -> &'static str {
254        match self {
255            GenomeSizeCategory::Small => "small (<5MB)",
256            GenomeSizeCategory::Medium => "medium (5-20MB)",
257            GenomeSizeCategory::Large => "large (>20MB)",
258        }
259    }
260}
261
262/// Bucket genome files by size category
263/// Returns a map of category -> list of genome paths
264fn bucket_genomes_by_size(genome_files: &[PathBuf]) -> FxHashMap<GenomeSizeCategory, Vec<PathBuf>> {
265    let mut buckets: FxHashMap<GenomeSizeCategory, Vec<PathBuf>> = FxHashMap::default();
266
267    for path in genome_files {
268        let size = std::fs::metadata(path)
269            .map(|m| m.len())
270            .unwrap_or(0);
271        let category = GenomeSizeCategory::from_size(size);
272        buckets.entry(category).or_default().push(path.clone());
273    }
274
275    buckets
276}
277
278// ============================================================================
279// Genome Catalogue
280// ============================================================================
281
282/// Unified genome catalogue entry
283/// Maps accession to taxonomy information
284#[derive(Clone, Debug)]
285pub struct CatalogEntry {
286    pub accession: String,      // GCF_/GCA_ or NZ_/CP accession
287    pub taxid: String,          // NCBI taxonomy ID
288    pub species_taxid: String,  // Species-level taxonomy ID
289    pub genus: String,          // Genus name (parsed from organism name)
290    pub species: String,        // Species name
291    pub organism_name: String,  // Full organism name
292    pub source: String,         // "refseq", "genbank", or "plsdb"
293}
294
295/// Genome catalog for accession-to-genus lookup
296pub struct GenomeCatalog {
297    entries: FxHashMap<String, CatalogEntry>,
298}
299
300impl Default for GenomeCatalog {
301    fn default() -> Self {
302        Self::new()
303    }
304}
305
306impl GenomeCatalog {
307    pub fn new() -> Self {
308        Self {
309            entries: FxHashMap::default(),
310        }
311    }
312
313    /// Add an assembly to the catalog with taxonomy-based genus lookup
314    pub fn add_assembly_with_taxonomy(
315        &mut self,
316        asm: &AssemblyInfo,
317        source: &str,
318        taxonomy: Option<&TaxonomyDB>,
319    ) {
320        let (parsed_genus, species) = parse_organism_name(&asm.organism_name);
321
322        // Resolve genus with smart fallback (uncultured, taxonomy, higher ranks)
323        let genus = resolve_genus(&asm.organism_name, &asm.taxid, &parsed_genus, taxonomy);
324
325        let entry = CatalogEntry {
326            accession: asm.accession.clone(),
327            taxid: asm.taxid.clone(),
328            species_taxid: asm.species_taxid.clone(),
329            genus,
330            species,
331            organism_name: asm.organism_name.clone(),
332            source: source.to_string(),
333        };
334
335        // Index by full accession and base accession (without version)
336        self.entries.insert(asm.accession.clone(), entry.clone());
337        if let Some(base) = asm.accession.split('.').next() {
338            self.entries.insert(base.to_string(), entry);
339        }
340    }
341
342    /// Add a PLSDB plasmid to the catalog with taxonomy-based genus lookup
343    pub fn add_plasmid_with_taxonomy(
344        &mut self,
345        plasmid: &PlasmidInfo,
346        taxonomy: Option<&TaxonomyDB>,
347    ) {
348        // Construct organism name from genus + species
349        let organism_name = format!("{} {}", plasmid.genus, plasmid.species);
350
351        // Resolve genus with smart fallback (uncultured, taxonomy, higher ranks)
352        let genus = resolve_genus(&organism_name, &plasmid.taxonomy_uid, &plasmid.genus, taxonomy);
353
354        let entry = CatalogEntry {
355            accession: plasmid.accession.clone(),
356            taxid: plasmid.taxonomy_uid.clone(),
357            species_taxid: String::new(),
358            genus,
359            species: plasmid.species.clone(),
360            organism_name,
361            source: "plsdb".to_string(),
362        };
363
364        // Index by full accession and variants
365        self.entries.insert(plasmid.accession.clone(), entry.clone());
366
367        // Also index by base accession (without NZ_ prefix and version)
368        let base = plasmid.accession
369            .strip_prefix("NZ_")
370            .unwrap_or(&plasmid.accession);
371        let no_ver = base.split('.').next().unwrap_or(base);
372        self.entries.insert(no_ver.to_string(), entry.clone());
373
374        // Also index by underscore-replaced version (for filename-derived lookups)
375        // e.g., NZ_CP157204.1 -> NZ_CP157204_1
376        let underscore_key = plasmid.accession.replace('.', "_");
377        self.entries.insert(underscore_key, entry);
378    }
379
380    /// Look up genus by accession (tries multiple variants)
381    pub fn get_genus(&self, accession: &str) -> Option<&str> {
382        // Try exact match first
383        if let Some(entry) = self.entries.get(accession) {
384            return Some(&entry.genus);
385        }
386
387        // Try without version
388        if let Some(base) = accession.split('.').next() {
389            if let Some(entry) = self.entries.get(base) {
390                return Some(&entry.genus);
391            }
392        }
393
394        // Try without NZ_ prefix
395        let stripped = accession.strip_prefix("NZ_").unwrap_or(accession);
396        if let Some(entry) = self.entries.get(stripped) {
397            return Some(&entry.genus);
398        }
399        let stripped_base = stripped.split('.').next().unwrap_or(stripped);
400        if let Some(entry) = self.entries.get(stripped_base) {
401            return Some(&entry.genus);
402        }
403
404        // Try with underscore-to-dot conversion (for PLSDB filenames like NZ_CP157204_1)
405        // Pattern: NZ_ABC123_1 -> NZ_ABC123.1
406        if let Some(last_underscore) = accession.rfind('_') {
407            let suffix = &accession[last_underscore + 1..];
408            // Check if suffix looks like a version number (digits only)
409            if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
410                let with_dot = format!("{}.{}", &accession[..last_underscore], suffix);
411                if let Some(entry) = self.entries.get(&with_dot) {
412                    return Some(&entry.genus);
413                }
414                // Also try without NZ_ prefix
415                let stripped_with_dot = with_dot.strip_prefix("NZ_").unwrap_or(&with_dot);
416                if let Some(entry) = self.entries.get(stripped_with_dot) {
417                    return Some(&entry.genus);
418                }
419            }
420        }
421
422        None
423    }
424
425    /// Save catalog to TSV file
426    pub fn save(&self, path: &Path) -> Result<()> {
427        let file = File::create(path)?;
428        let mut writer = BufWriter::new(file);
429
430        writeln!(writer, "accession\ttaxid\tspecies_taxid\tgenus\tspecies\torganism_name\tsource")?;
431
432        // Deduplicate entries (same accession may be indexed multiple ways)
433        let mut seen: FxHashSet<&str> = FxHashSet::default();
434        for entry in self.entries.values() {
435            if seen.contains(entry.accession.as_str()) {
436                continue;
437            }
438            seen.insert(&entry.accession);
439
440            writeln!(
441                writer,
442                "{}\t{}\t{}\t{}\t{}\t{}\t{}",
443                entry.accession,
444                entry.taxid,
445                entry.species_taxid,
446                entry.genus,
447                entry.species,
448                entry.organism_name.replace('\t', " "),
449                entry.source
450            )?;
451        }
452
453        Ok(())
454    }
455
456    /// Get entry count (unique accessions)
457    pub fn len(&self) -> usize {
458        let mut seen: FxHashSet<&str> = FxHashSet::default();
459        for entry in self.entries.values() {
460            seen.insert(&entry.accession);
461        }
462        seen.len()
463    }
464
465    /// Check if catalog is empty (required by clippy for len())
466    #[allow(dead_code)]
467    pub fn is_empty(&self) -> bool {
468        self.entries.is_empty()
469    }
470}
471
472/// Parse organism name to extract genus and species
473fn parse_organism_name(organism: &str) -> (String, String) {
474    let parts: Vec<&str> = organism.split_whitespace().collect();
475
476    let genus = parts.first()
477        .map(|s| s.to_string())
478        .unwrap_or_else(|| "Unknown".to_string());
479
480    let species = if parts.len() >= 2 {
481        // Join remaining parts as species (handles "sp." and strain info)
482        parts[1..].join(" ")
483    } else {
484        "Unknown".to_string()
485    };
486
487    (clean_genus(&genus), species)
488}
489
490/// Clean genus name: remove parentheses and citations
491fn clean_genus(genus: &str) -> String {
492    // Remove everything after '(' including the parenthesis
493    // e.g., "Methylacidiphilum (ex Ratnadevi et al. 2023)" -> "Methylacidiphilum"
494    let cleaned = if let Some(idx) = genus.find('(') {
495        genus[..idx].trim()
496    } else {
497        genus.trim()
498    };
499
500    cleaned.to_string()
501}
502
503/// Determine genus with smart fallback logic
504/// Priority:
505/// 1. If organism contains "uncultured" -> "uncultured"
506/// 2. Use taxonomy database to find genus or higher rank
507/// 3. Fall back to parsed genus from organism name
508fn resolve_genus(
509    organism_name: &str,
510    taxid_str: &str,
511    _parsed_genus: &str,
512    taxonomy: Option<&TaxonomyDB>,
513) -> String {
514    // Check for uncultured
515    let organism_lower = organism_name.to_lowercase();
516    if organism_lower.contains("uncultured") {
517        return "uncultured".to_string();
518    }
519
520    // Try taxonomy database
521    if let Some(tax_db) = taxonomy {
522        if let Ok(taxid) = taxid_str.parse::<u32>() {
523            // First try to get genus directly
524            if let Some(genus) = tax_db.get_genus(taxid) {
525                return clean_genus(&genus);
526            }
527
528            // If no genus, try to get higher rank (family, order, class)
529            if let Some((name, _rank)) = tax_db.get_genus_or_higher(taxid) {
530                return clean_genus(&name);
531            }
532        }
533    }
534
535    // Fall back: use full organism name (cleaned) when taxonomy lookup fails
536    // This preserves more information (e.g., "Dehalococcoidia bacterium" instead of just "Dehalococcoidia")
537    let cleaned = clean_genus(organism_name);
538    if cleaned.is_empty() {
539        "unknown".to_string()
540    } else {
541        cleaned
542    }
543}
544
545/// NCBI Taxonomy database for taxid → genus lookup
546/// Parses names.dmp and nodes.dmp from taxdump
547pub struct TaxonomyDB {
548    /// taxid → scientific name
549    names: FxHashMap<u32, String>,
550    /// taxid → (parent_taxid, rank)
551    nodes: FxHashMap<u32, (u32, String)>,
552}
553
554impl Default for TaxonomyDB {
555    fn default() -> Self {
556        Self::new()
557    }
558}
559
560impl TaxonomyDB {
561    pub fn new() -> Self {
562        Self {
563            names: FxHashMap::default(),
564            nodes: FxHashMap::default(),
565        }
566    }
567
568    /// Load taxonomy from taxdump directory
569    pub fn load(taxdump_dir: &Path) -> Result<Self> {
570        let mut db = Self::new();
571
572        let names_path = taxdump_dir.join("names.dmp");
573        let nodes_path = taxdump_dir.join("nodes.dmp");
574
575        if !names_path.exists() || !nodes_path.exists() {
576            anyhow::bail!("Taxonomy files not found in {}", taxdump_dir.display());
577        }
578
579        // Parse names.dmp
580        // Format: tax_id | name_txt | unique name | name class |
581        eprintln!("    Loading names.dmp...");
582        let file = File::open(&names_path)?;
583        let reader = BufReader::new(file);
584
585        for line in reader.lines() {
586            let line = line?;
587            let fields: Vec<&str> = line.split("\t|\t").collect();
588            if fields.len() < 4 {
589                continue;
590            }
591
592            // Only keep scientific names
593            let name_class = fields[3].trim_end_matches("\t|");
594            if name_class != "scientific name" {
595                continue;
596            }
597
598            let taxid: u32 = match fields[0].parse() {
599                Ok(id) => id,
600                Err(_) => continue,
601            };
602            let name = fields[1].to_string();
603
604            db.names.insert(taxid, name);
605        }
606
607        // Parse nodes.dmp
608        // Format: tax_id | parent_tax_id | rank | ...
609        eprintln!("    Loading nodes.dmp...");
610        let file = File::open(&nodes_path)?;
611        let reader = BufReader::new(file);
612
613        for line in reader.lines() {
614            let line = line?;
615            let fields: Vec<&str> = line.split("\t|\t").collect();
616            if fields.len() < 3 {
617                continue;
618            }
619
620            let taxid: u32 = match fields[0].parse() {
621                Ok(id) => id,
622                Err(_) => continue,
623            };
624            let parent_taxid: u32 = match fields[1].parse() {
625                Ok(id) => id,
626                Err(_) => continue,
627            };
628            let rank = fields[2].to_string();
629
630            db.nodes.insert(taxid, (parent_taxid, rank));
631        }
632
633        eprintln!("    Loaded {} taxa, {} nodes", db.names.len(), db.nodes.len());
634        Ok(db)
635    }
636
637    /// Get genus name for a taxid by traversing up the tree
638    pub fn get_genus(&self, taxid: u32) -> Option<String> {
639        let mut current = taxid;
640        let mut visited = 0;
641
642        // Traverse up to 50 levels (safety limit)
643        while visited < 50 {
644            if let Some((parent, rank)) = self.nodes.get(&current) {
645                if rank == "genus" {
646                    return self.names.get(&current).cloned();
647                }
648
649                // Check if we're at the root (parent == self)
650                if *parent == current {
651                    return None;
652                }
653
654                current = *parent;
655                visited += 1;
656            } else {
657                return None;
658            }
659        }
660
661        None
662    }
663
664    /// Get genus or higher rank (family, order, class) if genus not found
665    /// Returns (name, rank) tuple
666    pub fn get_genus_or_higher(&self, taxid: u32) -> Option<(String, String)> {
667        let mut current = taxid;
668        let mut visited = 0;
669
670        // Priority ranks to return if genus not found
671        let target_ranks = ["genus", "family", "order", "class", "phylum"];
672        let mut best_match: Option<(String, String, usize)> = None; // (name, rank, priority)
673
674        // Traverse up to 50 levels (safety limit)
675        while visited < 50 {
676            if let Some((parent, rank)) = self.nodes.get(&current) {
677                // Check if this rank is one we want
678                if let Some(priority) = target_ranks.iter().position(|r| r == rank) {
679                    if let Some(name) = self.names.get(&current) {
680                        // Return immediately if genus found
681                        if rank == "genus" {
682                            return Some((name.clone(), rank.clone()));
683                        }
684                        // Otherwise, keep track of best (highest priority = lowest index)
685                        if best_match.is_none() || priority < best_match.as_ref().unwrap().2 {
686                            best_match = Some((name.clone(), rank.clone(), priority));
687                        }
688                    }
689                }
690
691                // Check if we're at the root (parent == self)
692                if *parent == current {
693                    break;
694                }
695
696                current = *parent;
697                visited += 1;
698            } else {
699                break;
700            }
701        }
702
703        best_match.map(|(name, rank, _)| (name, rank))
704    }
705}
706
707/// Build state for resume/tracking
708#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
709pub struct BuildState {
710    pub started_at: String,
711    pub last_updated: String,
712    pub current_step: String,
713    pub completed_steps: Vec<String>,
714    pub genomes_downloaded: usize,
715    pub plsdb_extracted: usize,
716    pub alignments_done: bool,
717    pub flanking_extracted: bool,
718}
719
720impl Default for BuildState {
721    fn default() -> Self {
722        Self::new()
723    }
724}
725
726impl BuildState {
727    pub fn new() -> Self {
728        let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
729        Self {
730            started_at: now.clone(),
731            last_updated: now,
732            current_step: String::new(),
733            completed_steps: Vec::new(),
734            genomes_downloaded: 0,
735            plsdb_extracted: 0,
736            alignments_done: false,
737            flanking_extracted: false,
738        }
739    }
740
741    pub fn load(path: &Path) -> Option<Self> {
742        std::fs::read_to_string(path)
743            .ok()
744            .and_then(|s| serde_json::from_str(&s).ok())
745    }
746
747    pub fn save(&self, path: &Path) -> Result<()> {
748        let json = serde_json::to_string_pretty(self)?;
749        std::fs::write(path, json)?;
750        Ok(())
751    }
752
753    pub fn update_step(&mut self, step: &str) {
754        self.current_step = step.to_string();
755        self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
756    }
757
758    pub fn complete_step(&mut self, step: &str) {
759        if !self.completed_steps.contains(&step.to_string()) {
760            self.completed_steps.push(step.to_string());
761        }
762        self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
763    }
764
765    pub fn is_completed(&self, step: &str) -> bool {
766        self.completed_steps.contains(&step.to_string())
767    }
768}
769
770/// Builder configuration
771pub struct FlankingDbBuilder {
772    output_dir: PathBuf,
773    threads: usize,
774    amr_db_path: PathBuf,
775    /// NCBI email for API access (required, enables higher rate limits)
776    email: String,
777    /// Flanking build configuration
778    config: FlankBuildConfig,
779}
780
781impl FlankingDbBuilder {
782    pub fn new(amr_db: &Path, output_dir: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Self {
783        Self {
784            output_dir: output_dir.to_path_buf(),
785            threads,
786            amr_db_path: amr_db.to_path_buf(),
787            email: email.to_string(),
788            config,
789        }
790    }
791
792    /// Run the complete pipeline with resume support
793    pub fn build(&self) -> Result<()> {
794        eprintln!("\n============================================================");
795        eprintln!(" ARGenus Flanking Database Builder (Rust)");
796        eprintln!("============================================================");
797        eprintln!("Output directory: {}", self.output_dir.display());
798        eprintln!("NCBI Email: {}", self.email);
799        eprintln!("Download method: NCBI Datasets API (batch)");
800        eprintln!("Threads: {}", self.threads);
801        eprintln!("Flanking length: {} bp", self.config.flanking_length);
802        eprintln!("Queue buffer: {} GB", self.config.queue_buffer_gb);
803        eprintln!("Alignment: minimap2 (AMR indexed)");
804        eprintln!();
805
806        std::fs::create_dir_all(&self.output_dir)?;
807        let genomes_dir = self.output_dir.join("genomes");
808        std::fs::create_dir_all(&genomes_dir)?;
809        // Use user-provided PLSDB directory or create default
810        let plsdb_dir = self.config.plsdb.dir.clone().unwrap_or_else(|| self.output_dir.join("plsdb"));
811        if !self.config.plsdb.skip {
812            std::fs::create_dir_all(&plsdb_dir)?;
813        }
814        let taxonomy_dir = self.output_dir.join("taxonomy");
815        let temp_dir = self.output_dir.join("temp");
816        std::fs::create_dir_all(&temp_dir)?;
817        let state_path = self.output_dir.join("build_state.json");
818
819        // Load or create build state
820        let mut state = BuildState::load(&state_path).unwrap_or_else(|| {
821            eprintln!("Starting fresh build...");
822            BuildState::new()
823        });
824
825        if !state.completed_steps.is_empty() {
826            eprintln!("Resuming from previous state:");
827            eprintln!("  Started: {}", state.started_at);
828            eprintln!("  Last updated: {}", state.last_updated);
829            eprintln!("  Completed steps: {}", state.completed_steps.join(", "));
830            eprintln!("  Genomes downloaded: {}", state.genomes_downloaded);
831            eprintln!("  PLSDB extracted: {}", state.plsdb_extracted);
832            eprintln!();
833        }
834
835        // Initialize unified catalog
836        let mut catalog = GenomeCatalog::new();
837        let catalog_path = self.output_dir.join("genome_catalog.tsv");
838
839        // Step 1: Download NCBI taxdump (for authoritative taxonomy)
840        if !state.is_completed("taxonomy_download") {
841            state.update_step("taxonomy_download");
842            state.save(&state_path)?;
843            eprintln!("[1/9] Downloading NCBI taxonomy database...");
844            self.download_taxdump(&taxonomy_dir)?;
845            state.complete_step("taxonomy_download");
846            state.save(&state_path)?;
847        } else {
848            eprintln!("[1/9] Taxonomy database already downloaded, skipping...");
849        }
850
851        // Step 2: Load taxonomy database
852        eprintln!("[2/9] Loading taxonomy database...");
853        let taxonomy = match TaxonomyDB::load(&taxonomy_dir) {
854            Ok(db) => {
855                eprintln!("    Taxonomy database loaded successfully");
856                Some(db)
857            }
858            Err(e) => {
859                eprintln!("    Warning: Failed to load taxonomy: {}", e);
860                eprintln!("    Genus will be parsed from organism names instead");
861                None
862            }
863        };
864
865        // Step 3: Download NCBI assembly summary (GenBank only)
866        eprintln!("[3/9] Downloading NCBI assembly summaries (GenBank)...");
867        let assemblies = self.download_assembly_summaries()?;
868        eprintln!("    GenBank assemblies: {} (Complete Genome + Chromosome)", assemblies.len());
869
870        // Add assemblies to catalog (with taxonomy-based genus if available)
871        for asm in &assemblies {
872            catalog.add_assembly_with_taxonomy(asm, "genbank", taxonomy.as_ref());
873        }
874
875        // Step 4: Download PLSDB (skip if --skip-plsdb or --plsdb-dir provided)
876        let standalone_plasmids = if self.config.plsdb.skip {
877            eprintln!("[4/9] PLSDB skipped (--skip-plsdb)");
878            eprintln!("[5/9] PLSDB metadata skipped");
879            Vec::new()
880        } else if self.config.plsdb.dir.is_some() {
881            // User provided pre-downloaded PLSDB directory
882            eprintln!("[4/9] Using pre-downloaded PLSDB: {}", plsdb_dir.display());
883            // Validate required files exist
884            let nuccore_csv = plsdb_dir.join("nuccore.csv");
885            let fasta_path = plsdb_dir.join("sequences.fasta");
886            if !nuccore_csv.exists() {
887                // Check if meta.tar.gz needs extraction
888                let meta_tar = plsdb_dir.join("meta.tar.gz");
889                if meta_tar.exists() {
890                    eprintln!("    Extracting meta.tar.gz...");
891                    std::process::Command::new("tar")
892                        .args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
893                        .status()
894                        .with_context(|| "Failed to extract PLSDB metadata")?;
895                } else {
896                    anyhow::bail!("PLSDB directory missing nuccore.csv and meta.tar.gz: {}", plsdb_dir.display());
897                }
898            }
899            if !fasta_path.exists() {
900                anyhow::bail!("PLSDB directory missing sequences.fasta: {}", plsdb_dir.display());
901            }
902            state.complete_step("plsdb_download");
903            state.save(&state_path)?;
904
905            // Step 5: Load PLSDB metadata
906            eprintln!("[5/9] Loading PLSDB metadata...");
907            let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
908            eprintln!("    Standalone circular+complete plasmids: {} (not in any assembly)",
909                     plasmids.len());
910            plasmids
911        } else {
912            // Download from server
913            if !state.is_completed("plsdb_download") {
914                state.update_step("plsdb_download");
915                state.save(&state_path)?;
916                eprintln!("[4/9] Downloading PLSDB database...");
917                self.download_plsdb(&plsdb_dir)?;
918                state.complete_step("plsdb_download");
919                state.save(&state_path)?;
920            } else {
921                eprintln!("[4/9] PLSDB database already downloaded, skipping...");
922            }
923
924            // Step 5: Load PLSDB metadata (standalone plasmids only)
925            eprintln!("[5/9] Loading PLSDB metadata...");
926            let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
927            eprintln!("    Standalone circular+complete plasmids: {} (not in any assembly)",
928                     plasmids.len());
929            plasmids
930        };
931
932        // Add PLSDB plasmids to catalog (with taxonomy-based genus)
933        for plasmid in &standalone_plasmids {
934            catalog.add_plasmid_with_taxonomy(plasmid, taxonomy.as_ref());
935        }
936
937        // Step 6: Save unified catalog
938        eprintln!("[6/9] Saving unified genome catalog...");
939        catalog.save(&catalog_path)?;
940        eprintln!("    Catalog entries: {}", catalog.len());
941        eprintln!("    Saved to: {}", catalog_path.display());
942
943        // Step 7-9: Batch pipeline (download+align per batch, no combined FASTA)
944        // Saves ~350GB disk by avoiding combined FASTA
945        let paf_output = self.output_dir.join("all_alignments.paf");
946        let merged_hits = self.output_dir.join("merged_alignment_hits.tsv");
947        let output_tsv = self.output_dir.join("all_flanking_sequences.tsv");
948
949        if !state.is_completed("alignment_flanking") {
950            // Step 7: Download NCBI genomes and align in batches
951            state.update_step("genome_download");
952            state.save(&state_path)?;
953
954            eprintln!("[7/9] Batch processing: Download + Align NCBI genomes...");
955            eprintln!("    Producer-consumer pattern with {} GB queue buffer", self.config.queue_buffer_gb);
956
957            let downloaded = self.download_and_align_batches(
958                &assemblies,
959                &genomes_dir,
960                &temp_dir,
961                &paf_output,
962                &catalog,
963            )?;
964            state.genomes_downloaded = downloaded;
965            state.complete_step("genome_download");
966            state.save(&state_path)?;
967            eprintln!("    Processed {} genomes", downloaded);
968
969            // Step 8: Process PLSDB sequences (batch align)
970            if self.config.plsdb.skip {
971                eprintln!("[8/9] PLSDB processing skipped (--skip-plsdb)");
972                state.plsdb_extracted = 0;
973            } else {
974                state.update_step("plsdb_extract");
975                state.save(&state_path)?;
976                eprintln!("[8/9] Processing PLSDB sequences...");
977                let plsdb_count = self.align_plsdb(
978                    &standalone_plasmids,
979                    &plsdb_dir,
980                    &temp_dir,
981                    &paf_output,
982                )?;
983                state.plsdb_extracted = plsdb_count;
984                state.complete_step("plsdb_extract");
985                state.save(&state_path)?;
986            }
987
988            // Step 9: Convert PAF to merged format and extract flanking
989            state.update_step("alignment_flanking");
990            state.save(&state_path)?;
991
992            // Step 9a: Convert PAF to merged hits format
993            eprintln!("[9/9] Processing alignment results...");
994            self.convert_paf_to_merged(&paf_output, &merged_hits)?;
995            state.alignments_done = true;
996            state.save(&state_path)?;
997
998            // Step 9b: Extract flanking sequences from individual genome files
999            eprintln!("    Extracting flanking sequences (flanking_length: {} bp)...",
1000                self.config.flanking_length);
1001            self.extract_flanking_sequences(
1002                &merged_hits,
1003                &genomes_dir,
1004                &output_tsv,
1005                &catalog,
1006            )?;
1007            // Also extract from PLSDB (if not skipped)
1008            if !self.config.plsdb.skip {
1009                self.extract_flanking_from_plsdb(
1010                    &merged_hits,
1011                    &plsdb_dir,
1012                    &output_tsv,
1013                    &catalog,
1014                )?;
1015            }
1016            state.flanking_extracted = true;
1017            state.complete_step("alignment_flanking");
1018            state.save(&state_path)?;
1019
1020            // Cleanup: only delete temp files on success (keep genome files for resume)
1021            eprintln!("    Cleaning up PAF file...");
1022            std::fs::remove_file(&paf_output).ok();
1023        } else {
1024            eprintln!("[7-9] Batch pipeline already completed, skipping...");
1025        }
1026
1027        // Cleanup temp directory
1028        if temp_dir.exists() {
1029            eprintln!("\nCleaning up temp files...");
1030            std::fs::remove_dir_all(&temp_dir).ok();
1031        }
1032
1033        // Step 10: Build FDB from TSV (external sort + zstd compression)
1034        let fdb_path = self.output_dir.join("flanking.fdb");
1035        if !state.is_completed("fdb_build") && output_tsv.exists() {
1036            state.update_step("fdb_build");
1037            state.save(&state_path)?;
1038            eprintln!("\n[10/10] Building FDB from TSV (external sort + zstd)...");
1039
1040            // Use 1GB buffer for external sort, all available threads
1041            let buffer_size_mb = 1024;
1042            crate::fdb::build(&output_tsv, &fdb_path, buffer_size_mb, self.threads)?;
1043
1044            state.complete_step("fdb_build");
1045            state.save(&state_path)?;
1046        } else if fdb_path.exists() {
1047            eprintln!("[10/10] FDB already built, skipping...");
1048        }
1049
1050        eprintln!("\n============================================================");
1051        eprintln!(" Build Complete");
1052        eprintln!("============================================================");
1053        eprintln!("Output files:");
1054        eprintln!("  - {}", output_tsv.display());
1055        if fdb_path.exists() {
1056            let fdb_size = std::fs::metadata(&fdb_path).map(|m| m.len()).unwrap_or(0);
1057            eprintln!("  - {} ({:.1} MB)", fdb_path.display(), fdb_size as f64 / 1024.0 / 1024.0);
1058        }
1059        eprintln!("  - {}", catalog_path.display());
1060
1061        Ok(())
1062    }
1063
1064    /// Download GenBank-only assembly summaries (bacteria + archaea)
1065    /// GenBank is a superset of RefSeq with 0 FTP NA entries
1066    fn download_assembly_summaries(&self) -> Result<Vec<AssemblyInfo>> {
1067        let mut assemblies = Vec::new();
1068
1069        // Download GenBank only (no RefSeq deduplication needed)
1070        for kingdom in ["bacteria", "archaea"] {
1071            let url = format!(
1072                "{}/genbank/{}/assembly_summary.txt",
1073                NCBI_FTP_BASE, kingdom
1074            );
1075            eprintln!("    Downloading genbank/{}...", kingdom);
1076
1077            match self.download_and_parse_assembly_summary(&url, None) {
1078                Ok((mut asm, _)) => {
1079                    eprintln!("      Found {} assemblies (Complete Genome + Chromosome)", asm.len());
1080                    assemblies.append(&mut asm);
1081                }
1082                Err(e) => {
1083                    eprintln!("      Warning: Failed to download {}: {}", url, e);
1084                }
1085            }
1086        }
1087
1088        eprintln!("    Total assemblies: {}", assemblies.len());
1089        Ok(assemblies)
1090    }
1091
1092    /// Download PLSDB meta and fasta files
1093    fn download_plsdb(&self, plsdb_dir: &Path) -> Result<()> {
1094        let meta_tar = plsdb_dir.join("meta.tar.gz");
1095        let fasta_path = plsdb_dir.join("sequences.fasta");
1096
1097        // Download meta archive if needed
1098        if !plsdb_dir.join("nuccore.csv").exists() {
1099            eprintln!("    Downloading PLSDB metadata...");
1100            self.download_file(PLSDB_META_URL, &meta_tar)?;
1101
1102            // Extract tar.gz
1103            eprintln!("    Extracting metadata...");
1104            let status = Command::new("tar")
1105                .args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
1106                .status()
1107                .with_context(|| "Failed to extract PLSDB metadata")?;
1108
1109            if !status.success() {
1110                anyhow::bail!("tar extraction failed");
1111            }
1112
1113            // Clean up tar file
1114            std::fs::remove_file(&meta_tar).ok();
1115        } else {
1116            eprintln!("    PLSDB metadata already exists, skipping download...");
1117        }
1118
1119        // Download FASTA if needed
1120        if !fasta_path.exists() {
1121            eprintln!("    Downloading PLSDB sequences (~7GB)...");
1122            self.download_file(PLSDB_FASTA_URL, &fasta_path)?;
1123        } else {
1124            eprintln!("    PLSDB sequences already exist, skipping download...");
1125        }
1126
1127        Ok(())
1128    }
1129
1130    /// Download and extract NCBI taxdump
1131    fn download_taxdump(&self, taxonomy_dir: &Path) -> Result<()> {
1132        let tar_path = taxonomy_dir.join("taxdump.tar.gz");
1133        let names_path = taxonomy_dir.join("names.dmp");
1134
1135        // Skip if already extracted
1136        if names_path.exists() {
1137            eprintln!("    Taxdump already exists, skipping download...");
1138            return Ok(());
1139        }
1140
1141        std::fs::create_dir_all(taxonomy_dir)?;
1142
1143        // Download taxdump
1144        eprintln!("    Downloading NCBI taxdump (~60MB)...");
1145        self.download_file(NCBI_TAXDUMP_URL, &tar_path)?;
1146
1147        // Extract
1148        eprintln!("    Extracting taxdump...");
1149        let status = Command::new("tar")
1150            .args(["-xzf", tar_path.to_str().unwrap(), "-C", taxonomy_dir.to_str().unwrap()])
1151            .status()
1152            .with_context(|| "Failed to extract taxdump")?;
1153
1154        if !status.success() {
1155            anyhow::bail!("tar extraction failed");
1156        }
1157
1158        // Clean up tar file
1159        std::fs::remove_file(&tar_path).ok();
1160
1161        Ok(())
1162    }
1163
1164    /// Download a file from URL with retry
1165    fn download_file(&self, url: &str, output_path: &Path) -> Result<()> {
1166        for attempt in 0..3 {
1167            match self.download_file_once(url, output_path) {
1168                Ok(_) => return Ok(()),
1169                Err(e) if attempt < 2 => {
1170                    eprintln!("      Download failed (attempt {}): {}", attempt + 1, e);
1171                    eprintln!("      Retrying in 5 seconds...");
1172                    std::thread::sleep(Duration::from_secs(5));
1173                    continue;
1174                }
1175                Err(e) => return Err(e),
1176            }
1177        }
1178        Ok(())
1179    }
1180
1181    /// Download a file from URL (single attempt)
1182    fn download_file_once(&self, url: &str, output_path: &Path) -> Result<()> {
1183        let response = ureq::get(url)
1184            .set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
1185            .timeout(Duration::from_secs(7200)) // 2 hour timeout for dynamically generated files
1186            .call()
1187            .with_context(|| format!("Failed to download {}", url))?;
1188
1189        let mut file = File::create(output_path)?;
1190        let mut reader = response.into_reader();
1191
1192        // Stream download with progress
1193        let mut buffer = [0u8; 65536];
1194        let mut total = 0usize;
1195        loop {
1196            match reader.read(&mut buffer) {
1197                Ok(0) => break,
1198                Ok(n) => {
1199                    file.write_all(&buffer[..n])?;
1200                    total += n;
1201                    if total.is_multiple_of(100 * 1024 * 1024) {
1202                        eprintln!("      Downloaded {} MB...", total / (1024 * 1024));
1203                    }
1204                }
1205                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
1206                Err(e) => return Err(e.into()),
1207            }
1208        }
1209        eprintln!("      Total: {} MB", total / (1024 * 1024));
1210
1211        Ok(())
1212    }
1213
1214    /// Download and parse assembly_summary.txt (streaming for large files)
1215    /// Returns (assemblies, paired_gca_accessions)
1216    /// If exclude_set is provided, skip assemblies in that set (for GenBank deduplication)
1217    fn download_and_parse_assembly_summary(
1218        &self,
1219        url: &str,
1220        exclude_set: Option<&FxHashSet<String>>,
1221    ) -> Result<(Vec<AssemblyInfo>, Vec<String>)> {
1222        let response = ureq::get(url)
1223            .timeout(Duration::from_secs(600))
1224            .call()
1225            .with_context(|| format!("Failed to download {}", url))?;
1226
1227        let reader = BufReader::new(response.into_reader());
1228        let mut assemblies = Vec::new();
1229        let mut paired_gca = Vec::new();
1230        let mut skipped_duplicates = 0usize;
1231
1232        for line in reader.lines() {
1233            let line = line?;
1234            if line.starts_with('#') {
1235                continue;
1236            }
1237
1238            let fields: Vec<&str> = line.split('\t').collect();
1239            if fields.len() < 20 {
1240                continue;
1241            }
1242
1243            let accession = fields[0];
1244            let assembly_level = fields[11];
1245
1246            // Only include Complete Genome and Chromosome level
1247            if assembly_level != "Complete Genome" && assembly_level != "Chromosome" {
1248                continue;
1249            }
1250
1251            let ftp_path = fields[19];
1252            if ftp_path == "na" || ftp_path.is_empty() {
1253                continue;
1254            }
1255
1256            // Check for deduplication (GenBank vs RefSeq)
1257            if let Some(exclude) = exclude_set {
1258                // Strip version for matching (GCA_000005845.2 -> GCA_000005845)
1259                let acc_base = accession.split('.').next().unwrap_or(accession);
1260                if exclude.contains(accession) || exclude.contains(acc_base) {
1261                    skipped_duplicates += 1;
1262                    continue;
1263                }
1264            }
1265
1266            // For RefSeq, collect paired GCA accessions with identical status
1267            // Field 17: gbrs_paired_asm (e.g., GCA_000005845.2)
1268            // Field 18: paired_asm_comp (identical, different, na)
1269            if fields.len() > 18 && accession.starts_with("GCF_") {
1270                let paired_asm = fields[17];
1271                let paired_comp = fields[18];
1272                if paired_asm != "na" && !paired_asm.is_empty() && paired_comp == "identical" {
1273                    paired_gca.push(paired_asm.to_string());
1274                    // Also add without version
1275                    if let Some(base) = paired_asm.split('.').next() {
1276                        paired_gca.push(base.to_string());
1277                    }
1278                }
1279            }
1280
1281            assemblies.push(AssemblyInfo {
1282                accession: accession.to_string(),
1283                taxid: fields[5].to_string(),
1284                species_taxid: fields[6].to_string(),
1285                organism_name: fields[7].to_string(),
1286            });
1287        }
1288
1289        if skipped_duplicates > 0 {
1290            eprintln!("      Skipped {} duplicates (have identical RefSeq)", skipped_duplicates);
1291        }
1292
1293        Ok((assemblies, paired_gca))
1294    }
1295
1296    /// Load PLSDB plasmids (circular + complete only)
1297    fn load_plsdb_plasmids(&self, plsdb_dir: &Path) -> Result<Vec<PlasmidInfo>> {
1298        let nuccore_path = plsdb_dir.join("nuccore.csv");
1299        let taxonomy_path = plsdb_dir.join("taxonomy.csv");
1300
1301        // Load taxonomy mapping
1302        let mut taxonomy: FxHashMap<String, (String, String)> = FxHashMap::default();
1303        if taxonomy_path.exists() {
1304            let tax_file = File::open(&taxonomy_path)?;
1305            let tax_reader = BufReader::new(tax_file);
1306
1307            for (idx, line) in tax_reader.lines().enumerate() {
1308                let line = line?;
1309                if idx == 0 {
1310                    continue; // Skip header
1311                }
1312
1313                let fields: Vec<&str> = line.split(',').collect();
1314                if fields.len() >= 10 {
1315                    let uid = fields[0].to_string();
1316                    let genus = fields[8].to_string();
1317                    let species = fields[9].to_string();
1318                    taxonomy.insert(uid, (genus, species));
1319                }
1320            }
1321        }
1322
1323        // Load nuccore and filter
1324        let mut plasmids = Vec::new();
1325        let nuc_file = File::open(&nuccore_path)?;
1326        let nuc_reader = BufReader::new(nuc_file);
1327
1328        for (idx, line) in nuc_reader.lines().enumerate() {
1329            let line = line?;
1330            if idx == 0 {
1331                continue; // Skip header
1332            }
1333
1334            // Parse CSV with quotes handling
1335            let fields = parse_csv_line(&line);
1336            if fields.len() < 15 {
1337                continue;
1338            }
1339
1340            let completeness = &fields[4];
1341            let topology = &fields[14];
1342            let assembly_uid = &fields[9];  // ASSEMBLY_UID
1343
1344            // Filter: circular + complete only
1345            if completeness != "complete" {
1346                continue;
1347            }
1348            if topology != "circular" {
1349                continue;
1350            }
1351
1352            // Filter: standalone plasmids only (not linked to any assembly)
1353            // Plasmids with ASSEMBLY_UID != "-1" are already included in GenBank assemblies
1354            if assembly_uid != "-1" {
1355                continue;
1356            }
1357
1358            let taxonomy_uid = &fields[12];
1359            let (genus, species) = taxonomy.get(taxonomy_uid)
1360                .cloned()
1361                .unwrap_or_else(|| ("Unknown".to_string(), "Unknown".to_string()));
1362
1363            plasmids.push(PlasmidInfo {
1364                accession: fields[1].clone(),
1365                taxonomy_uid: taxonomy_uid.clone(),
1366                genus,
1367                species,
1368            });
1369        }
1370
1371        Ok(plasmids)
1372    }
1373
1374    /// Batch processing: Download and align in batches without combined FASTA
1375    /// Producer-consumer pattern with backpressure based on queue_buffer_gb
1376    /// Saves ~350GB disk by avoiding combined FASTA
1377    fn download_and_align_batches(
1378        &self,
1379        assemblies: &[AssemblyInfo],
1380        genomes_dir: &Path,
1381        temp_dir: &Path,
1382        paf_output: &Path,
1383        _catalog: &GenomeCatalog,
1384    ) -> Result<usize> {
1385        let all_accessions_path = temp_dir.join("accessions_all.txt");
1386        let done_accessions_path = temp_dir.join("accessions_done.txt");
1387
1388        // Step 1: Save all accessions if not exists
1389        let accessions: Vec<String> = assemblies.iter()
1390            .map(|a| a.accession.clone())
1391            .collect();
1392
1393        if !all_accessions_path.exists() {
1394            eprintln!("    Saving {} accessions to temp file...", accessions.len());
1395            let mut file = File::create(&all_accessions_path)?;
1396            for acc in &accessions {
1397                writeln!(file, "{}", acc)?;
1398            }
1399        }
1400
1401        // Step 2: Load already downloaded accessions and verify PAF integrity
1402        let mut done_set: FxHashSet<String> = FxHashSet::default();
1403        if done_accessions_path.exists() {
1404            let file = File::open(&done_accessions_path)?;
1405            let reader = BufReader::new(file);
1406            for acc in reader.lines().map_while(Result::ok) {
1407                done_set.insert(acc.trim().to_string());
1408            }
1409
1410            // Verify PAF file integrity on resume
1411            if !done_set.is_empty() {
1412                if !paf_output.exists() {
1413                    eprintln!("    WARNING: PAF file missing but {} genomes marked as done", done_set.len());
1414                    eprintln!("    Clearing done list and reprocessing all genomes...");
1415                    done_set.clear();
1416                    std::fs::remove_file(&done_accessions_path).ok();
1417                } else {
1418                    // Verify PAF has content (basic sanity check)
1419                    let paf_size = std::fs::metadata(paf_output)?.len();
1420                    if paf_size == 0 {
1421                        eprintln!("    WARNING: PAF file is empty but {} genomes marked as done", done_set.len());
1422                        eprintln!("    Clearing done list and reprocessing all genomes...");
1423                        done_set.clear();
1424                        std::fs::remove_file(&done_accessions_path).ok();
1425                    } else {
1426                        // Verify genome files exist for done accessions
1427                        let mut missing_genomes = Vec::new();
1428                        for acc in &done_set {
1429                            let genome_path = genomes_dir.join(format!("{}.fna", acc));
1430                            if !genome_path.exists() {
1431                                missing_genomes.push(acc.clone());
1432                            }
1433                        }
1434                        if !missing_genomes.is_empty() {
1435                            eprintln!("    WARNING: {} genome files missing for done accessions", missing_genomes.len());
1436                            for acc in &missing_genomes {
1437                                done_set.remove(acc);
1438                            }
1439                            eprintln!("    Removed missing genomes from done list, will re-download");
1440                        }
1441
1442                        if !done_set.is_empty() {
1443                            eprintln!("    Resume: {} already processed (PAF: {} MB)",
1444                                     done_set.len(), paf_size / 1024 / 1024);
1445                        }
1446                    }
1447                }
1448            }
1449        }
1450
1451        // Step 3: Filter remaining accessions
1452        let remaining: Vec<&String> = accessions.iter()
1453            .filter(|a| !done_set.contains(*a))
1454            .collect();
1455
1456        if remaining.is_empty() {
1457            eprintln!("    All {} genomes already processed", accessions.len());
1458            return Ok(accessions.len());
1459        }
1460
1461        eprintln!("    Remaining: {}/{} genomes to process", remaining.len(), accessions.len());
1462
1463        // Step 4: Calculate channel capacity based on queue_buffer_gb
1464        // Each batch ~3GB (1000 genomes × 3MB avg), queue_buffer_gb=30 → ~10 batches
1465        let bytes_per_batch = API_BATCH_SIZE * 3 * 1024 * 1024; // ~3GB per batch
1466        let queue_capacity = std::cmp::max(
1467            1,
1468            (self.config.queue_buffer_gb as usize * 1024 * 1024 * 1024) / bytes_per_batch
1469        );
1470        eprintln!("    Queue capacity: {} batches (based on {} GB buffer)", queue_capacity, self.config.queue_buffer_gb);
1471
1472        // Step 5: Setup producer-consumer with bounded channel
1473        let batches: Vec<Vec<String>> = remaining.chunks(API_BATCH_SIZE)
1474            .map(|chunk| chunk.iter().map(|s| (*s).clone()).collect())
1475            .collect();
1476        let total_batches = batches.len();
1477
1478        // Batch info: (batch_idx, genome_files)
1479        let (tx, rx) = mpsc::sync_channel::<(usize, Vec<PathBuf>)>(queue_capacity);
1480
1481        // Shared state for progress tracking
1482        let downloaded = Arc::new(AtomicUsize::new(done_set.len()));
1483        let aligned = Arc::new(AtomicUsize::new(done_set.len()));
1484        let download_error = Arc::new(AtomicBool::new(false));
1485
1486        // Clone necessary data for producer thread
1487        let genomes_dir = genomes_dir.to_path_buf();
1488        let temp_dir_producer = temp_dir.to_path_buf();
1489        let temp_dir_consumer = temp_dir.to_path_buf();
1490        let email = self.email.clone();
1491        let downloaded_clone = Arc::clone(&downloaded);
1492        let download_error_clone = Arc::clone(&download_error);
1493        let total_accessions = accessions.len();
1494
1495        // Producer thread: downloads batches and sends to channel
1496        let producer = thread::spawn(move || -> Result<()> {
1497            for (batch_idx, batch) in batches.into_iter().enumerate() {
1498                if download_error_clone.load(Ordering::Relaxed) {
1499                    break;
1500                }
1501
1502                let zip_path = temp_dir_producer.join(format!("batch_{:04}.zip", batch_idx));
1503
1504                // Download batch via NCBI Datasets API
1505                let url = format!("{}/genome/download", NCBI_DATASETS_API);
1506                let acc_list: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
1507                let request_body = serde_json::json!({
1508                    "accessions": acc_list,
1509                    "include_annotation_type": ["GENOME_FASTA"]
1510                });
1511
1512                let response = ureq::post(&url)
1513                    .set("Content-Type", "application/json")
1514                    .set("ncbi-client-id", &email)
1515                    .timeout(Duration::from_secs(3600))
1516                    .send_json(&request_body);
1517
1518                match response {
1519                    Ok(resp) => {
1520                        // Save ZIP file
1521                        let mut zip_file = File::create(&zip_path)?;
1522                        let mut reader = resp.into_reader();
1523                        std::io::copy(&mut reader, &mut zip_file)?;
1524                        drop(zip_file);
1525
1526                        // Extract genomes to individual files
1527                        let genome_files = Self::extract_batch_to_files_static(&zip_path, &genomes_dir)?;
1528
1529                        // Clean up ZIP
1530                        std::fs::remove_file(&zip_path).ok();
1531
1532                        downloaded_clone.fetch_add(batch.len(), Ordering::Relaxed);
1533                        eprintln!("    [Download] Batch {}/{}: {} genomes ({}/{})",
1534                                 batch_idx + 1, total_batches, batch.len(),
1535                                 downloaded_clone.load(Ordering::Relaxed), total_accessions);
1536
1537                        // Send to consumer (blocks if channel is full - backpressure)
1538                        if tx.send((batch_idx, genome_files)).is_err() {
1539                            break; // Consumer dropped
1540                        }
1541                    }
1542                    Err(e) => {
1543                        eprintln!("    [Download] Batch {} failed: {}", batch_idx + 1, e);
1544                        download_error_clone.store(true, Ordering::Relaxed);
1545                        break;
1546                    }
1547                }
1548
1549                // Rate limit between batches
1550                std::thread::sleep(Duration::from_millis(200));
1551            }
1552            Ok(())
1553        });
1554
1555        // Consumer: receives batches, runs minimap2, appends to PAF
1556        let paf_file = Arc::new(Mutex::new(
1557            std::fs::OpenOptions::new()
1558                .create(true)
1559                .append(true)
1560                .open(paf_output)?
1561        ));
1562        let done_file = Arc::new(Mutex::new(
1563            std::fs::OpenOptions::new()
1564                .create(true)
1565                .append(true)
1566                .open(&done_accessions_path)?
1567        ));
1568
1569        let amr_db = self.amr_db_path.clone();
1570        let max_threads = self.threads;
1571
1572        while let Ok((batch_idx, genome_files)) = rx.recv() {
1573            if genome_files.is_empty() {
1574                continue;
1575            }
1576
1577            // Adaptive thread bucketing: group genomes by size category
1578            let buckets = bucket_genomes_by_size(&genome_files);
1579            let mut batch_hits: Vec<PafHit> = Vec::new();
1580
1581            // Process each size category with appropriate thread count
1582            for (category, bucket_files) in &buckets {
1583                if bucket_files.is_empty() {
1584                    continue;
1585                }
1586
1587                let bucket_threads = category.thread_count(max_threads);
1588                let bucket_suffix = format!("batch_{:04}_{:?}", batch_idx, category);
1589                let bucket_fasta = temp_dir_consumer.join(format!("{}.fas", bucket_suffix));
1590
1591                // Create bucket FASTA with contig|filename headers
1592                {
1593                    let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
1594                    for genome_path in bucket_files {
1595                        let filename = genome_path.file_name()
1596                            .and_then(|n| n.to_str())
1597                            .unwrap_or("unknown.fna");
1598
1599                        let file = File::open(genome_path)?;
1600                        let reader = BufReader::new(file);
1601                        for line in reader.lines() {
1602                            let line = line?;
1603                            if let Some(stripped) = line.strip_prefix('>') {
1604                                let contig_id = stripped.split_whitespace().next().unwrap_or(stripped);
1605                                writeln!(fasta_writer, ">{}|{}", contig_id, filename)?;
1606                            } else {
1607                                writeln!(fasta_writer, "{}", line)?;
1608                            }
1609                        }
1610                    }
1611                    fasta_writer.flush()?;
1612                }
1613
1614                // Run minimap2 with adaptive thread count
1615                let bucket_paf = temp_dir_consumer.join(format!("{}.paf", bucket_suffix));
1616                let status = Command::new("minimap2")
1617                    .args([
1618                        "-cx", "asm20",
1619                        "-t", &bucket_threads.to_string(),
1620                        amr_db.to_str().unwrap(),
1621                        bucket_fasta.to_str().unwrap(),
1622                        "-o", bucket_paf.to_str().unwrap(),
1623                    ])
1624                    .stdout(std::process::Stdio::null())
1625                    .stderr(std::process::Stdio::null())
1626                    .status();
1627
1628                if let Ok(s) = status {
1629                    if s.success() && bucket_paf.exists() {
1630                        // Parse PAF hits from bucket
1631                        let paf_content = std::fs::read_to_string(&bucket_paf)?;
1632                        let hits: Vec<PafHit> = paf_content
1633                            .lines()
1634                            .filter_map(PafHit::from_paf_line)
1635                            .collect();
1636                        batch_hits.extend(hits);
1637                    }
1638                }
1639
1640                // Clean up bucket temp files
1641                std::fs::remove_file(&bucket_fasta).ok();
1642                std::fs::remove_file(&bucket_paf).ok();
1643            }
1644
1645            // Deduplicate combined hits from all buckets
1646            let dedup_hits = deduplicate_paf_hits(batch_hits);
1647
1648            // Append deduplicated hits to main PAF and flush immediately
1649            {
1650                let mut paf = paf_file.lock().unwrap();
1651                for hit in &dedup_hits {
1652                    writeln!(paf, "{}", hit.raw_line)?;
1653                }
1654                paf.flush()?;  // Ensure PAF is written before marking as done
1655            }
1656
1657            // Record completed accessions (only after PAF is flushed)
1658            {
1659                let mut done = done_file.lock().unwrap();
1660                for genome_path in &genome_files {
1661                    if let Some(stem) = genome_path.file_stem().and_then(|s| s.to_str()) {
1662                        writeln!(done, "{}", stem)?;
1663                    }
1664                }
1665                done.flush()?;  // Ensure done list is persisted
1666            }
1667
1668            let count = genome_files.len();
1669            aligned.fetch_add(count, Ordering::Relaxed);
1670
1671            // Log with bucket distribution
1672            let bucket_info: Vec<String> = buckets.iter()
1673                .map(|(cat, files)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), files.len()))
1674                .collect();
1675            eprintln!("    [Align] Batch {}/{}: {} genomes [{}] ({}/{})",
1676                     batch_idx + 1, total_batches, count,
1677                     bucket_info.join(", "),
1678                     aligned.load(Ordering::Relaxed), total_accessions);
1679        }
1680
1681        // Wait for producer
1682        let _ = producer.join();
1683
1684        Ok(aligned.load(Ordering::Relaxed))
1685    }
1686
1687    /// Extract genomes from ZIP to individual files in genomes_dir
1688    fn extract_batch_to_files_static(zip_path: &Path, genomes_dir: &Path) -> Result<Vec<PathBuf>> {
1689        let zip_file = File::open(zip_path)?;
1690        let mut archive = zip::ZipArchive::new(zip_file)?;
1691        let mut genome_files = Vec::new();
1692
1693        for i in 0..archive.len() {
1694            let mut file = archive.by_index(i)?;
1695            let name = file.name().to_string();
1696
1697            if name.ends_with("_genomic.fna") || name.ends_with("_genomic.fasta") {
1698                let parts: Vec<&str> = name.split('/').collect();
1699                if let Some(acc_dir) = parts.iter().find(|p| p.starts_with("GCA_") || p.starts_with("GCF_")) {
1700                    let filename = format!("{}.fna", acc_dir);
1701                    let output_path = genomes_dir.join(&filename);
1702
1703                    // Extract to file
1704                    let mut content = Vec::new();
1705                    file.read_to_end(&mut content)?;
1706                    std::fs::write(&output_path, &content)?;
1707
1708                    genome_files.push(output_path);
1709                }
1710            }
1711        }
1712
1713        Ok(genome_files)
1714    }
1715
1716    /// PLSDB processing: batch align PLSDB sequences
1717    fn align_plsdb(
1718        &self,
1719        plasmids: &[PlasmidInfo],
1720        plsdb_dir: &Path,
1721        temp_dir: &Path,
1722        paf_output: &Path,
1723    ) -> Result<usize> {
1724        let fasta_path = plsdb_dir.join("sequences.fasta");
1725
1726        if !fasta_path.exists() {
1727            eprintln!("    Warning: sequences.fasta not found in PLSDB directory");
1728            return Ok(0);
1729        }
1730
1731        // Build set of accessions to extract
1732        let target_accs: FxHashSet<_> = plasmids.iter()
1733            .map(|p| p.accession.clone())
1734            .collect();
1735
1736        if target_accs.is_empty() {
1737            eprintln!("    No PLSDB plasmids to process");
1738            return Ok(0);
1739        }
1740
1741        // Process in batches
1742        let acc_list: Vec<_> = target_accs.iter().collect();
1743        let batches: Vec<_> = acc_list.chunks(API_BATCH_SIZE).collect();
1744        let total_batches = batches.len();
1745        let mut total_processed = 0usize;
1746
1747        eprintln!("    Processing {} PLSDB plasmids in {} batches...",
1748                 target_accs.len(), total_batches);
1749
1750        // Open source FASTA and build index for random access
1751        let file = File::open(&fasta_path)?;
1752        let reader = BufReader::new(file);
1753
1754        // Build sequence map (memory intensive but simpler)
1755        let mut seq_map: FxHashMap<String, String> = FxHashMap::default();
1756        let mut current_acc: Option<String> = None;
1757        let mut current_seq = String::new();
1758
1759        for line in reader.lines() {
1760            let line = line?;
1761            if line.starts_with('>') {
1762                if let Some(acc) = current_acc.take() {
1763                    if target_accs.contains(&acc) {
1764                        seq_map.insert(acc, std::mem::take(&mut current_seq));
1765                    }
1766                }
1767                let header = line.trim_start_matches('>');
1768                let acc = header.split_whitespace().next().unwrap_or("").to_string();
1769                current_acc = Some(acc);
1770                current_seq.clear();
1771            } else {
1772                current_seq.push_str(line.trim());
1773            }
1774        }
1775        if let Some(acc) = current_acc {
1776            if target_accs.contains(&acc) {
1777                seq_map.insert(acc, current_seq);
1778            }
1779        }
1780
1781        eprintln!("    Loaded {} target sequences from PLSDB", seq_map.len());
1782
1783        // Process batches with adaptive thread bucketing
1784        for (batch_idx, batch) in batches.iter().enumerate() {
1785            // Bucket sequences by size for adaptive thread allocation
1786            let mut size_buckets: FxHashMap<GenomeSizeCategory, Vec<(&str, &str)>> = FxHashMap::default();
1787
1788            for acc in *batch {
1789                if let Some(seq) = seq_map.get(*acc) {
1790                    // Use sequence length as size estimate (1 byte per base)
1791                    let category = GenomeSizeCategory::from_size(seq.len() as u64);
1792                    size_buckets.entry(category)
1793                        .or_default()
1794                        .push((acc.as_str(), seq.as_str()));
1795                }
1796            }
1797
1798            let mut batch_hits: Vec<PafHit> = Vec::new();
1799
1800            // Process each size category with appropriate thread count
1801            for (category, bucket_seqs) in &size_buckets {
1802                if bucket_seqs.is_empty() {
1803                    continue;
1804                }
1805
1806                let bucket_threads = category.thread_count(self.threads);
1807                let bucket_suffix = format!("plsdb_batch_{:04}_{:?}", batch_idx, category);
1808                let bucket_fasta = temp_dir.join(format!("{}.fas", bucket_suffix));
1809
1810                // Create bucket FASTA with contig|filename headers
1811                {
1812                    let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
1813                    for (acc, seq) in bucket_seqs {
1814                        let filename = format!("{}.fna", acc.replace('.', "_"));
1815                        writeln!(fasta_writer, ">{}|{}", acc, filename)?;
1816                        writeln!(fasta_writer, "{}", seq)?;
1817                    }
1818                    fasta_writer.flush()?;
1819                }
1820
1821                // Run minimap2 with adaptive thread count
1822                let bucket_paf = temp_dir.join(format!("{}.paf", bucket_suffix));
1823                let status = Command::new("minimap2")
1824                    .args([
1825                        "-cx", "asm20",
1826                        "-t", &bucket_threads.to_string(),
1827                        self.amr_db_path.to_str().unwrap(),
1828                        bucket_fasta.to_str().unwrap(),
1829                        "-o", bucket_paf.to_str().unwrap(),
1830                    ])
1831                    .stdout(std::process::Stdio::null())
1832                    .stderr(std::process::Stdio::null())
1833                    .status();
1834
1835                if let Ok(s) = status {
1836                    if s.success() && bucket_paf.exists() {
1837                        // Parse PAF hits from bucket
1838                        let paf_content = std::fs::read_to_string(&bucket_paf)?;
1839                        let hits: Vec<PafHit> = paf_content
1840                            .lines()
1841                            .filter_map(PafHit::from_paf_line)
1842                            .collect();
1843                        batch_hits.extend(hits);
1844                    }
1845                }
1846
1847                // Clean up bucket temp files
1848                std::fs::remove_file(&bucket_fasta).ok();
1849                std::fs::remove_file(&bucket_paf).ok();
1850            }
1851
1852            // Deduplicate combined hits from all buckets
1853            let dedup_hits = deduplicate_paf_hits(batch_hits);
1854
1855            // Append deduplicated hits to main PAF
1856            {
1857                let mut paf_file = std::fs::OpenOptions::new()
1858                    .create(true)
1859                    .append(true)
1860                    .open(paf_output)?;
1861                for hit in &dedup_hits {
1862                    writeln!(paf_file, "{}", hit.raw_line)?;
1863                }
1864            }
1865
1866            total_processed += batch.len();
1867
1868            // Log with bucket distribution
1869            let bucket_info: Vec<String> = size_buckets.iter()
1870                .map(|(cat, seqs)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), seqs.len()))
1871                .collect();
1872            eprintln!("    [PLSDB] Batch {}/{}: {} plasmids [{}] ({}/{})",
1873                     batch_idx + 1, total_batches, batch.len(),
1874                     bucket_info.join(", "),
1875                     total_processed, target_accs.len());
1876        }
1877
1878        // Also write PLSDB sequences to plsdb_dir as individual files for flanking extraction
1879        for (acc, seq) in &seq_map {
1880            let out_path = plsdb_dir.join(format!("{}.fna", acc.replace('.', "_")));
1881            if !out_path.exists() {
1882                let mut out_file = File::create(&out_path)?;
1883                writeln!(out_file, ">{}", acc)?;
1884                writeln!(out_file, "{}", seq)?;
1885            }
1886        }
1887
1888        eprintln!("    Aligned {} PLSDB sequences", total_processed);
1889        Ok(total_processed)
1890    }
1891
1892    /// Extract flanking sequences from PLSDB files
1893    fn extract_flanking_from_plsdb(
1894        &self,
1895        hits_path: &Path,
1896        plsdb_dir: &Path,
1897        output_path: &Path,
1898        catalog: &GenomeCatalog,
1899    ) -> Result<()> {
1900        let hits_file = File::open(hits_path)?;
1901        let reader = BufReader::new(hits_file);
1902
1903        // Group hits by genome file (only PLSDB files)
1904        let mut plsdb_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
1905
1906        for (i, line) in reader.lines().enumerate() {
1907            let line = line?;
1908            // Skip header
1909            if i == 0 && line.starts_with("gene") {
1910                continue;
1911            }
1912
1913            let fields: Vec<&str> = line.split('\t').collect();
1914            if fields.len() < 4 {
1915                continue;
1916            }
1917
1918            let gene = fields[0];
1919            let contig_file = fields[1];
1920            let start: usize = fields[2].parse().unwrap_or(0);
1921            let end: usize = fields[3].parse().unwrap_or(0);
1922
1923            // Parse contig_id|filename format
1924            let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
1925                (contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
1926            } else {
1927                continue;
1928            };
1929
1930            // Only process PLSDB files (NZ_*, CP*, etc.)
1931            if genome_file.starts_with("NZ_") || genome_file.starts_with("CP") ||
1932               genome_file.starts_with("AP") || genome_file.starts_with("NC_") {
1933                plsdb_hits.entry(genome_file)
1934                    .or_default()
1935                    .push((gene.to_string(), contig_id, start, end));
1936            }
1937        }
1938
1939        if plsdb_hits.is_empty() {
1940            return Ok(());
1941        }
1942
1943        // Append to output file
1944        let mut writer = std::fs::OpenOptions::new()
1945            .create(true)
1946            .append(true)
1947            .open(output_path)?;
1948
1949        let mut extracted = 0usize;
1950
1951        for (genome_file, hits) in &plsdb_hits {
1952            let genome_path = plsdb_dir.join(genome_file);
1953            if !genome_path.exists() {
1954                continue;
1955            }
1956
1957            // Load genome sequences
1958            let sequences = match self.load_genome_sequences(&genome_path) {
1959                Ok(seqs) => seqs,
1960                Err(_) => continue,
1961            };
1962
1963            let seq_map: FxHashMap<&str, &str> = sequences.iter()
1964                .map(|(header, seq)| {
1965                    let key = header.split_whitespace().next().unwrap_or(header.as_str());
1966                    (key, seq.as_str())
1967                })
1968                .collect();
1969
1970            let genome_acc = genome_file.trim_end_matches(".fna");
1971            let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
1972
1973            for (gene, contig_id, start, end) in hits {
1974                let contig_seq = match seq_map.get(contig_id.as_str()) {
1975                    Some(seq) => *seq,
1976                    None => continue,
1977                };
1978
1979                let contig_len = contig_seq.len();
1980                if *start >= contig_len || *end > contig_len || start >= end {
1981                    continue;
1982                }
1983
1984                // Get genus
1985                let genus = base_genus.clone().unwrap_or_else(|| "Unknown".to_string());
1986
1987                // Extract flanking sequences
1988                let upstream_start = start.saturating_sub(self.config.flanking_length);
1989                let upstream = &contig_seq[upstream_start..*start];
1990
1991                let downstream_end = std::cmp::min(*end + self.config.flanking_length, contig_len);
1992                let downstream = &contig_seq[*end..downstream_end];
1993
1994                writeln!(writer, "{}\t{}\t{}\t{}\t{}\t{}\t{}",
1995                        gene, contig_id, genus, start, end, upstream, downstream)?;
1996                extracted += 1;
1997            }
1998        }
1999
2000        if extracted > 0 {
2001            eprintln!("    Extracted {} PLSDB flanking sequences", extracted);
2002        }
2003        Ok(())
2004    }
2005
2006    /// Convert PAF to merged hits format (gene, contig|filename, start, end)
2007    fn convert_paf_to_merged(
2008        &self,
2009        paf_path: &Path,
2010        output_path: &Path,
2011    ) -> Result<()> {
2012        let mut hits: FxHashSet<(String, String, usize, usize)> = FxHashSet::default();
2013
2014        // Process PAF: query=contig|file, target=AMR
2015        // Columns: 0=query, 2=qstart, 3=qend, 5=target
2016        if paf_path.exists() {
2017            let file = File::open(paf_path)?;
2018            let reader = BufReader::new(file);
2019            for line in reader.lines() {
2020                let line = line?;
2021                let fields: Vec<&str> = line.split('\t').collect();
2022                if fields.len() < 12 {
2023                    continue;
2024                }
2025                let contig_file = fields[0].to_string(); // contig|filename
2026                let gene = fields[5].to_string();        // AMR gene
2027                let start: usize = fields[2].parse().unwrap_or(0);
2028                let end: usize = fields[3].parse().unwrap_or(0);
2029                hits.insert((gene, contig_file, start, end));
2030            }
2031        }
2032
2033        // Write output
2034        let mut writer = BufWriter::new(File::create(output_path)?);
2035        for (gene, contig_file, start, end) in &hits {
2036            writeln!(writer, "{}\t{}\t{}\t{}", gene, contig_file, start, end)?;
2037        }
2038
2039        eprintln!("    Found {} unique hits", hits.len());
2040        Ok(())
2041    }
2042
2043    /// Extract flanking sequences from merged hits TSV
2044    /// Input format: gene\tcontig_file\tstart\tend (with header)
2045    /// contig_file format: contig_id|filename.fna
2046    fn extract_flanking_sequences(
2047        &self,
2048        hits_path: &Path,
2049        genomes_dir: &Path,
2050        output_path: &Path,
2051        catalog: &GenomeCatalog,
2052    ) -> Result<()> {
2053        // Skip if output already exists and is non-empty
2054        if output_path.exists() {
2055            let metadata = std::fs::metadata(output_path)?;
2056            if metadata.len() > 0 {
2057                eprintln!("    Flanking sequences file already exists ({} MB), skipping extraction...",
2058                         metadata.len() / (1024 * 1024));
2059                return Ok(());
2060            }
2061        }
2062
2063        let hits_file = File::open(hits_path)?;
2064        let reader = BufReader::new(hits_file);
2065
2066        // Group hits by genome file
2067        // Format: genome_file -> Vec<(gene, contig_id, start, end)>
2068        let mut genome_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
2069
2070        for (i, line) in reader.lines().enumerate() {
2071            let line = line?;
2072            // Skip header
2073            if i == 0 && line.starts_with("gene") {
2074                continue;
2075            }
2076
2077            let fields: Vec<&str> = line.split('\t').collect();
2078            if fields.len() < 4 {
2079                continue;
2080            }
2081
2082            let gene = fields[0];
2083            let contig_file = fields[1]; // format: contig_id|filename.fna
2084            let start: usize = fields[2].parse().unwrap_or(0);
2085            let end: usize = fields[3].parse().unwrap_or(0);
2086
2087            // Parse contig_id|filename format
2088            let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
2089                (contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
2090            } else {
2091                // Fallback: try to extract genome file from contig name
2092                (contig_file.to_string(), extract_genome_file(contig_file))
2093            };
2094
2095            genome_hits.entry(genome_file)
2096                .or_default()
2097                .push((gene.to_string(), contig_id, start, end));
2098        }
2099
2100        // Process each genome and extract flanking sequences
2101        let output_file = File::create(output_path)?;
2102        let mut writer = BufWriter::new(output_file);
2103
2104        // Write header
2105        writeln!(writer, "Gene\tContig\tGenus\tStart\tEnd\tUpstream\tDownstream")?;
2106
2107        let total_genomes = genome_hits.len();
2108        let mut processed = 0;
2109        let mut genus_found = 0usize;
2110        let mut genus_unknown = 0usize;
2111        let mut sequences_extracted = 0usize;
2112
2113        for (genome_file, hits) in &genome_hits {
2114            let genome_path = genomes_dir.join(genome_file);
2115            if !genome_path.exists() {
2116                continue;
2117            }
2118
2119            // Load genome sequences
2120            let sequences = match self.load_genome_sequences(&genome_path) {
2121                Ok(seqs) => seqs,
2122                Err(_) => continue,
2123            };
2124
2125            // Build a map for quick contig lookup
2126            let seq_map: FxHashMap<&str, &str> = sequences.iter()
2127                .map(|(header, seq)| {
2128                    let key = header.split_whitespace().next().unwrap_or(header.as_str());
2129                    (key, seq.as_str())
2130                })
2131                .collect();
2132
2133            // Try to get genus from genome filename (accession)
2134            let genome_acc = genome_file.trim_end_matches(".fna");
2135            let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
2136
2137            // Extract flanking for each hit
2138            for (gene, contig_id, start, end) in hits {
2139                // Look up contig sequence
2140                let contig_seq = match seq_map.get(contig_id.as_str()) {
2141                    Some(seq) => *seq,
2142                    None => continue,
2143                };
2144
2145                let contig_len = contig_seq.len();
2146
2147                // Bounds check
2148                if *start >= contig_len || *end > contig_len || start >= end {
2149                    continue;
2150                }
2151
2152                let upstream_start = start.saturating_sub(self.config.flanking_length);
2153                let downstream_end = std::cmp::min(end + self.config.flanking_length, contig_len);
2154
2155                let upstream = if *start > 0 && upstream_start < *start {
2156                    &contig_seq[upstream_start..*start]
2157                } else {
2158                    ""
2159                };
2160
2161                let downstream = if *end < contig_len && *end < downstream_end {
2162                    &contig_seq[*end..downstream_end]
2163                } else {
2164                    ""
2165                };
2166
2167                // Try multiple methods to get genus:
2168                // 1. From genome accession (base_genus from NCBI catalog)
2169                // 2. From contig accession (for PLSDB plasmids)
2170                let genus = if let Some(ref g) = base_genus {
2171                    genus_found += 1;
2172                    g.clone()
2173                } else if let Some(g) = catalog.get_genus(contig_id) {
2174                    genus_found += 1;
2175                    g.to_string()
2176                } else {
2177                    genus_unknown += 1;
2178                    "Unknown".to_string()
2179                };
2180
2181                writeln!(
2182                    writer,
2183                    "{}\t{}\t{}\t{}\t{}\t{}\t{}",
2184                    gene, contig_id, genus, start, end, upstream, downstream
2185                )?;
2186                sequences_extracted += 1;
2187            }
2188
2189            processed += 1;
2190            if processed % 10000 == 0 || processed == total_genomes {
2191                eprintln!("    Processed {}/{} genomes ({} sequences)",
2192                    processed, total_genomes, sequences_extracted);
2193            }
2194        }
2195
2196        eprintln!("    Extracted {} flanking sequences from {} genomes",
2197            sequences_extracted, processed);
2198        eprintln!("    Genus resolved: {}, Unknown: {}", genus_found, genus_unknown);
2199        Ok(())
2200    }
2201
2202    /// Load genome sequences from FASTA
2203    fn load_genome_sequences(&self, genome_path: &Path) -> Result<Vec<(String, String)>> {
2204        let file = File::open(genome_path)?;
2205        let reader = BufReader::new(file);
2206
2207        let mut sequences = Vec::new();
2208        let mut current_name: Option<String> = None;
2209        let mut current_seq = String::new();
2210
2211        for line in reader.lines() {
2212            let line = line?;
2213
2214            if let Some(stripped) = line.strip_prefix('>') {
2215                if let Some(name) = current_name.take() {
2216                    sequences.push((name, std::mem::take(&mut current_seq)));
2217                }
2218                current_name = Some(stripped.to_string());
2219                current_seq.clear();
2220            } else {
2221                current_seq.push_str(line.trim());
2222            }
2223        }
2224
2225        if let Some(name) = current_name {
2226            sequences.push((name, current_seq));
2227        }
2228
2229        Ok(sequences)
2230    }
2231}
2232
2233/// Parse CSV line handling quoted fields
2234fn parse_csv_line(line: &str) -> Vec<String> {
2235    let mut fields = Vec::new();
2236    let mut current = String::new();
2237    let mut in_quotes = false;
2238
2239    for ch in line.chars() {
2240        match ch {
2241            '"' => in_quotes = !in_quotes,
2242            ',' if !in_quotes => {
2243                fields.push(std::mem::take(&mut current));
2244            }
2245            _ => current.push(ch),
2246        }
2247    }
2248    fields.push(current);
2249    fields
2250}
2251
2252/// Extract genome filename from contig name
2253fn extract_genome_file(contig: &str) -> String {
2254    // Try to extract GCA/GCF accession from contig name
2255    if let Some(start) = contig.find("GC") {
2256        let remainder = &contig[start..];
2257        if let Some(end) = remainder.find(|c: char| !c.is_alphanumeric() && c != '_' && c != '.') {
2258            return format!("{}.fna", &remainder[..end]);
2259        }
2260        return format!("{}.fna", remainder);
2261    }
2262
2263    // For PLSDB accessions (NZ_*, CP*, etc.)
2264    let acc = contig.split_whitespace().next().unwrap_or(contig);
2265    format!("{}.fna", acc.replace('.', "_"))
2266}
2267
2268/// Builds the flanking sequence database.
2269///
2270/// This is the main entry point for flanking database construction.
2271/// Downloads genomes, aligns ARGs, and extracts flanking sequences.
2272///
2273/// # Arguments
2274/// * `output_dir` - Directory for output files
2275/// * `arg_db` - Path to ARG reference database (AMR_NCBI.fas or .mmi)
2276/// * `threads` - Number of threads for parallel processing
2277/// * `email` - NCBI API email for higher rate limits
2278/// * `config` - Flanking build configuration (flanking_length, queue_buffer_gb, plsdb)
2279///
2280/// # Output Files
2281/// - `genome_catalogue.tsv`: Accession → genus mapping
2282/// - `flanking_db.tsv`: Raw flanking data
2283/// - `flanking.fdb`: Compressed database
2284pub fn build(output_dir: &Path, arg_db: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Result<()> {
2285    let builder = FlankingDbBuilder::new(arg_db, output_dir, threads, email, config);
2286    builder.build()
2287}