/// Pipeline: concurrent download + parse with graceful shutdown and resume.
///
/// Architecture:
///   Download threads (I/O-bound, share one ureq Agent)
///     --> parse_channel (bounded) <-- cached files also feed here
///       --> Parse workers (CPU-bound, one per core)
///         --> Parquet writers (zero contention, per-worker dirs)
///
///   Downloads and parsing overlap: workers never wait for all downloads
///   to finish. Cached files feed instantly, downloads feed as they arrive.
///   Bounded channel provides backpressure (slows downloads if parse lags).
use crate::fetch::{self, FtpFile};
use crate::parse;
use crate::writer::WorkerWriter;
use crossbeam_channel::bounded;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use std::collections::HashSet;
use std::io::{self, BufRead, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

/// Estimated space requirements (bytes).
const EST_CACHE_BYTES: u64 = 110 << 30;   // ~110 GB compressed XML
const EST_OUTPUT_BYTES: u64 = 45 << 30;   // ~45 GB parquet output
const MIN_RAM_BYTES: u64 = 2 << 30;       // 2 GB absolute minimum

pub struct Config {
    pub cache_dir: PathBuf,
    pub output_dir: PathBuf,
    pub workers: usize,
    pub dl_jobs: usize,
    pub baseline: bool,
    pub updates: bool,
}

impl Config {
    pub fn auto(output_dir: PathBuf, cache_dir: PathBuf) -> Self {
        let cpus = thread::available_parallelism().map(|n| n.get()).unwrap_or(4);
        let cap = coren::MachCap::read_opts(&output_dir, false, false);
        let ram_gb = cap.mem_total as f64 / (1u64 << 30) as f64;
        // ~200MB RSS per worker (quick-xml + arrow builders)
        let by_ram = (ram_gb / 0.2) as usize;
        let workers = cpus.min(by_ram).max(1);
        let dl_jobs = 8.min(cpus * 2);
        Config {
            cache_dir, output_dir, workers, dl_jobs,
            baseline: true, updates: true,
        }
    }
}

pub fn default_data_dir() -> PathBuf {
    dirs::data_dir().unwrap_or_else(|| ".".into()).join("pmq")
}

pub fn default_cache_dir() -> PathBuf {
    dirs::cache_dir().unwrap_or_else(|| ".".into()).join("pmq").join("downloads")
}

/// Glob pattern for querying a table across all workers.
pub fn table_glob(output_dir: &Path, table: &str) -> String {
    let p = output_dir.join("_workers").join("*").join(format!("{table}.parquet"));
    p.to_string_lossy().into_owned()
}

/// List available table names by scanning worker dirs.
pub fn list_tables(output_dir: &Path) -> Vec<String> {
    let mut tables = HashSet::new();
    let workers = output_dir.join("_workers");
    if let Ok(entries) = std::fs::read_dir(&workers) {
        for entry in entries.flatten() {
            if entry.path().is_dir() {
                if let Ok(files) = std::fs::read_dir(entry.path()) {
                    for f in files.flatten() {
                        if let Some(name) = f.path().file_stem().and_then(|s| s.to_str()) {
                            tables.insert(name.to_owned());
                        }
                    }
                }
            }
        }
    }
    let mut v: Vec<_> = tables.into_iter().collect();
    v.sort();
    v
}

// -- Pre-flight resource check ------------------------------------------------

fn disk_free_for(path: &Path) -> u64 {
    std::fs::create_dir_all(path).ok();
    coren::MachCap::read_opts(path, false, false).disk_free
}

fn fmt_gb(bytes: u64) -> String {
    format!("{:.1} GB", bytes as f64 / (1u64 << 30) as f64)
}

fn prompt_line(msg: &str) -> String {
    eprint!("{msg}");
    io::stderr().flush().ok();
    let mut line = String::new();
    io::stdin().lock().read_line(&mut line).ok();
    line.trim().to_owned()
}

pub fn preflight(output_dir: &Path, cache_dir: &Path) -> anyhow::Result<(PathBuf, PathBuf)> {
    let cap = coren::MachCap::read_opts(output_dir, false, false);
    let ram = cap.mem_available;

    eprintln!("[pmq] Pre-flight check...");
    eprintln!("[pmq]   RAM available: {}", fmt_gb(ram));

    if ram < MIN_RAM_BYTES {
        eprintln!("[pmq]   WARNING: very low RAM ({}).  May OOM with large files.", fmt_gb(ram));
    }

    let out = check_dir_space("output", output_dir, EST_OUTPUT_BYTES)?;
    let cache = check_dir_space("cache", cache_dir, EST_CACHE_BYTES)?;
    Ok((out, cache))
}

fn check_dir_space(label: &str, dir: &Path, need: u64) -> anyhow::Result<PathBuf> {
    let free = disk_free_for(dir);
    eprintln!("[pmq]   {} dir: {} (free: {}, need: ~{})",
        label, dir.display(), fmt_gb(free), fmt_gb(need));

    if free >= need {
        return Ok(dir.to_path_buf());
    }

    let short = need.saturating_sub(free);
    eprintln!("[pmq]   INSUFFICIENT: {} volume short by ~{}", label, fmt_gb(short));
    eprintln!("  Enter alternative {} path, or press Enter to override:", label);

    let answer = prompt_line(&format!("  {label} path> "));
    if answer.is_empty() {
        eprintln!("[pmq]   Continuing (user override)");
        return Ok(dir.to_path_buf());
    }

    let alt = PathBuf::from(&answer);
    std::fs::create_dir_all(&alt)?;
    check_dir_space(label, &alt, need)
}

// -- State file (resume support) ----------------------------------------------

fn state_path(output_dir: &Path) -> PathBuf { output_dir.join("_state") }

fn load_state(output_dir: &Path) -> HashSet<String> {
    let path = state_path(output_dir);
    let mut set = HashSet::new();
    if let Ok(file) = std::fs::File::open(&path) {
        for line in io::BufReader::new(file).lines().flatten() {
            let t = line.trim().to_owned();
            if !t.is_empty() { set.insert(t); }
        }
    }
    set
}

fn append_state(output_dir: &Path, filename: &str) {
    let path = state_path(output_dir);
    if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(path) {
        let _ = writeln!(f, "{filename}");
    }
}

// -- Stats --------------------------------------------------------------------

struct Stats {
    files_done: AtomicU64,
    records: AtomicU64,
    parse_errors: AtomicU64,
    dl_done: AtomicU64,
    dl_errors: AtomicU64,
}

impl Stats {
    fn new() -> Self {
        Stats {
            files_done: AtomicU64::new(0),
            records: AtomicU64::new(0),
            parse_errors: AtomicU64::new(0),
            dl_done: AtomicU64::new(0),
            dl_errors: AtomicU64::new(0),
        }
    }
}

// -- Worker state tracking (per-worker progress) -----------------------------

struct WorkerState {
    names: Vec<std::sync::Mutex<String>>,
}

impl WorkerState {
    fn new(n: usize) -> Self {
        WorkerState { names: (0..n).map(|_| std::sync::Mutex::new(String::new())).collect() }
    }

    fn set(&self, wid: usize, name: &str) {
        if let Some(m) = self.names.get(wid) {
            if let Ok(mut s) = m.lock() { s.clear(); s.push_str(name); }
        }
    }

    fn clear(&self, wid: usize) {
        if let Some(m) = self.names.get(wid) {
            if let Ok(mut s) = m.lock() { s.clear(); }
        }
    }

    fn active_names(&self) -> Vec<String> {
        self.names.iter().filter_map(|m| {
            let s = m.lock().ok()?;
            if s.is_empty() { None } else { Some(s.clone()) }
        }).collect()
    }

    fn format_line(&self) -> String {
        self.names.iter().enumerate().map(|(i, m)| {
            let name = m.lock().map(|s| s.clone()).unwrap_or_default();
            if name.is_empty() {
                format!("W{}:.", i)
            } else {
                let short = name.strip_prefix("pubmed26").unwrap_or(&name);
                let short = short.strip_suffix(".xml.gz").unwrap_or(short);
                format!("W{}:{}", i, short)
            }
        }).collect::<Vec<_>>().join(" ")
    }
}

const DL_RETRIES: usize = 3;

// -- fd limit -----------------------------------------------------------------

fn raise_fd_limit() {
    match rlimit::increase_nofile_limit(8192) {
        Ok(n) => { if n < 512 { eprintln!("[pmq] WARNING: fd limit only {n}, may hit errors with many workers"); } }
        Err(e) => { eprintln!("[pmq] WARNING: could not raise fd limit: {e}"); }
    }
}

// -- Entry points -------------------------------------------------------------

pub fn run(config: &Config) -> anyhow::Result<()> {
    let t0 = Instant::now();

    let cap = coren::MachCap::read_opts(&config.output_dir, false, false);
    eprintln!("[pmq] workers={}  dl={}  RAM {:.1}GB avail={:.1}GB",
        config.workers, config.dl_jobs,
        cap.mem_total as f64 / (1u64 << 30) as f64,
        cap.mem_available as f64 / (1u64 << 30) as f64,
    );

    eprintln!("[pmq] Listing files from NCBI...");
    let all_files = fetch::list_files(config.baseline, config.updates)?;
    let total = all_files.len();
    eprintln!("[pmq] Found {total} files on server");
    if total == 0 { return Ok(()); }

    std::fs::create_dir_all(&config.cache_dir)?;
    let done_set = load_state(&config.output_dir);

    let mut cached = Vec::new();
    let mut to_download = Vec::new();
    let mut already_done = 0usize;

    for f in &all_files {
        if done_set.contains(&f.name) {
            already_done += 1;
            continue;
        }
        let path = config.cache_dir.join(&f.name);
        if path.exists() && std::fs::metadata(&path).map(|m| m.len() > 0).unwrap_or(false) {
            cached.push(path);
        } else {
            to_download.push(f.clone());
        }
    }

    if cached.is_empty() && to_download.is_empty() {
        eprintln!("[pmq] All {total} files already processed");
        return Ok(());
    }

    eprintln!(
        "[pmq] {} files total, {} done, {} cached, {} to download",
        total, already_done, cached.len(), to_download.len()
    );
    run_core(config, cached, to_download, t0)
}

pub fn run_local(config: &Config, files: Vec<PathBuf>) -> anyhow::Result<()> {
    let t0 = Instant::now();
    let done_set = load_state(&config.output_dir);
    let files: Vec<_> = files.into_iter().filter(|p| {
        !done_set.contains(&p.file_name().unwrap_or_default().to_string_lossy().to_string())
    }).collect();

    if files.is_empty() {
        eprintln!("[pmq] All files already processed");
        return Ok(());
    }
    eprintln!("[pmq] {} local files to process", files.len());
    run_core(config, files, vec![], t0)
}

/// Download with retry. Checks shutdown between attempts.
fn download_with_retry(
    agent: &ureq::Agent,
    f: &FtpFile,
    cache_dir: &Path,
    max_retries: usize,
    shutdown: &AtomicBool,
) -> Result<PathBuf, String> {
    for attempt in 0..=max_retries {
        if shutdown.load(Ordering::Relaxed) {
            return Err(format!("{}: shutdown", f.name));
        }
        match fetch::download_file(agent, f, cache_dir) {
            Ok(path) => return Ok(path),
            Err(e) => {
                if attempt < max_retries && !shutdown.load(Ordering::Relaxed) {
                    let delay = Duration::from_secs(2u64.pow(attempt as u32).min(30));
                    thread::sleep(delay);
                } else {
                    return Err(format!("{}: {e}", f.name));
                }
            }
        }
    }
    unreachable!()
}

// -- Core pipeline ------------------------------------------------------------

fn run_core(
    config: &Config,
    cached: Vec<PathBuf>,
    to_download: Vec<FtpFile>,
    t0: Instant,
) -> anyhow::Result<()> {
    let n_cached = cached.len();
    let n_download = to_download.len();
    let total = n_cached + n_download;
    if total == 0 { return Ok(()); }

    raise_fd_limit();

    std::fs::create_dir_all(&config.output_dir)?;
    let workers_dir = config.output_dir.join("_workers");
    std::fs::create_dir_all(&workers_dir)?;

    let stats = Arc::new(Stats::new());
    let shutdown = Arc::new(AtomicBool::new(false));
    let wstate = Arc::new(WorkerState::new(config.workers));

    // Adaptive flush threshold: use 40% of available RAM spread across workers.
    // ~3KB per article across all 27 tables (empirical average).
    let cap = coren::MachCap::read_opts(&config.output_dir, false, false);
    let budget_per_worker = (cap.mem_available as f64 * 0.4) as usize / config.workers.max(1);
    let flush_threshold = (budget_per_worker / 3000).clamp(5_000, 500_000);
    eprintln!("[pmq] flush_threshold={} (budget {:.0}MB/worker)",
        flush_threshold, budget_per_worker as f64 / (1 << 20) as f64);

    // Ctrl-C handler
    {
        let s = Arc::clone(&shutdown);
        ctrlc::set_handler(move || {
            if s.load(Ordering::Relaxed) {
                eprintln!("\n[pmq] Force quit");
                std::process::exit(1);
            }
            s.store(true, Ordering::Relaxed);
        }).ok();
    }

    // -- Channels --
    let (parse_tx, parse_rx) = bounded::<PathBuf>(config.workers * 4);
    let (dl_tx, dl_rx) = bounded::<FtpFile>(config.dl_jobs * 2);
    let (done_tx, done_rx) = bounded::<String>(256);

    // -- Progress bars --
    let mp = MultiProgress::new();

    let pb_main = mp.add(ProgressBar::new(total as u64));
    pb_main.set_style(ProgressStyle::default_bar()
        .template("[{elapsed_precise}] [{bar:30}] {pos}/{len} | {msg}")
        .expect("tpl").progress_chars("=>-"));
    pb_main.enable_steady_tick(Duration::from_millis(250));

    let pb_status = mp.add(ProgressBar::new_spinner());
    pb_status.set_style(ProgressStyle::default_spinner()
        .template("  {msg}").expect("tpl"));
    pb_status.enable_steady_tick(Duration::from_millis(250));

    // -- State persistence thread --
    let out_c = config.output_dir.clone();
    let state_thread = thread::spawn(move || {
        while let Ok(fname) = done_rx.recv() { append_state(&out_c, &fname); }
    });

    // -- Download threads --
    let dl_handles: Vec<_> = if n_download > 0 {
        let agent = ureq::AgentBuilder::new()
            .timeout_connect(Duration::from_secs(10))
            .timeout_read(Duration::from_secs(30))
            .build();

        // Feed download jobs
        let dl_tx_feeder = dl_tx.clone();
        let shutdown_f = Arc::clone(&shutdown);
        thread::Builder::new().name("dl-feed".into()).spawn(move || {
            for f in to_download {
                if shutdown_f.load(Ordering::Relaxed) { break; }
                if dl_tx_feeder.send(f).is_err() { break; }
            }
        }).ok();
        drop(dl_tx);

        (0..config.dl_jobs).map(|i| {
            let dl_rx = dl_rx.clone();
            let parse_tx = parse_tx.clone();
            let cache_dir = config.cache_dir.clone();
            let shutdown = Arc::clone(&shutdown);
            let stats = Arc::clone(&stats);
            let pb_main = pb_main.clone();
            let agent = agent.clone();

            thread::Builder::new()
                .name(format!("dl-{i}"))
                .spawn(move || {
                    while let Ok(f) = dl_rx.recv() {
                        if shutdown.load(Ordering::Relaxed) { break; }
                        match download_with_retry(&agent, &f, &cache_dir, DL_RETRIES, &shutdown) {
                            Ok(path) => {
                                stats.dl_done.fetch_add(1, Ordering::Relaxed);
                                let _ = parse_tx.send(path);
                            }
                            Err(msg) => {
                                if !shutdown.load(Ordering::Relaxed) {
                                    stats.dl_errors.fetch_add(1, Ordering::Relaxed);
                                    pb_main.suspend(|| {
                                        eprintln!("  dl error: {msg}");
                                    });
                                }
                            }
                        }
                    }
                })
                .expect("spawn dl thread")
        }).collect()
    } else {
        drop(dl_tx);
        vec![]
    };
    drop(dl_rx);

    // -- Feed cached files directly into parse channel --
    let cached_tx = parse_tx.clone();
    let shutdown_c = Arc::clone(&shutdown);
    let cache_feeder = thread::Builder::new()
        .name("cache-feed".into())
        .spawn(move || {
            for path in cached {
                if shutdown_c.load(Ordering::Relaxed) { break; }
                if cached_tx.send(path).is_err() { break; }
            }
        })
        .expect("spawn cache feeder");
    drop(parse_tx);

    // -- Parse workers --
    let parse_t0 = Instant::now();
    let mut parse_handles = Vec::new();
    for wid in 0..config.workers {
        let rx = parse_rx.clone();
        let wdir = workers_dir.clone();
        let stats = Arc::clone(&stats);
        let shutdown = Arc::clone(&shutdown);
        let pb_main = pb_main.clone();
        let wstate = Arc::clone(&wstate);
        let done_tx = done_tx.clone();

        parse_handles.push(thread::Builder::new()
            .name(format!("parse-{wid}"))
            .spawn(move || {
                let mut writer = match WorkerWriter::new(&wdir, wid, flush_threshold) {
                    Ok(w) => w,
                    Err(e) => { eprintln!("worker {wid}: init: {e}"); return; }
                };

                while let Ok(path) = rx.recv() {
                    if shutdown.load(Ordering::Relaxed) { break; }

                    let fname = path.file_name().unwrap_or_default().to_string_lossy().to_string();
                    wstate.set(wid, &fname);

                    match parse::parse_file(&path, &mut writer) {
                        Ok(n) => {
                            stats.records.fetch_add(n, Ordering::Relaxed);
                            let _ = done_tx.send(fname);
                        }
                        Err(e) => {
                            eprintln!("  PARSE ERROR: {fname}: {e}");
                            eprintln!("  Delete the cached file and run again to retry.");
                            std::fs::remove_file(&path).ok();
                            stats.parse_errors.fetch_add(1, Ordering::Relaxed);
                        }
                    }

                    wstate.clear(wid);
                    stats.files_done.fetch_add(1, Ordering::Relaxed);
                    pb_main.inc(1);
                }

                if let Err(e) = writer.finish() {
                    eprintln!("worker {wid}: flush error: {e}");
                }
            })
            .expect("spawn parse worker"));
    }
    drop(parse_rx);
    drop(done_tx);

    // -- Status update loop (main thread) --
    let mut shutdown_printed = false;
    loop {
        thread::sleep(Duration::from_millis(250));

        if shutdown.load(Ordering::Relaxed) && !shutdown_printed {
            shutdown_printed = true;
            pb_status.finish_and_clear();
            pb_main.finish_and_clear();
            eprintln!();
            eprintln!("[pmq] Ctrl-C received. Shutting down...");
            eprintln!("[pmq] Workers finishing current file. Ctrl-C again to force quit.");
            eprintln!();
            break;
        }

        let elapsed = parse_t0.elapsed().as_secs_f64();
        let recs = stats.records.load(Ordering::Relaxed);
        let rate = if elapsed > 0.5 { recs as f64 / elapsed } else { 0.0 };

        pb_main.set_message(format!("{:.1}M recs | {:.0}/s", recs as f64 / 1e6, rate));
        pb_status.set_message(wstate.format_line());

        let all_dl_done = dl_handles.iter().all(|h| h.is_finished());
        let cache_done = cache_feeder.is_finished();
        let all_parse_done = parse_handles.iter().all(|h| h.is_finished());

        if all_dl_done && cache_done && all_parse_done {
            break;
        }
    }

    // -- Join / shutdown drain --
    if shutdown_printed {
        loop {
            let active = wstate.active_names();
            if active.is_empty() && dl_handles.iter().all(|h| h.is_finished()) { break; }
            eprintln!("[pmq]   finishing: {}", active.join(" "));
            thread::sleep(Duration::from_secs(1));
        }
        eprintln!("[pmq] All tasks stopped");

        cache_feeder.join().ok();
        for h in dl_handles { h.join().ok(); }
        for h in parse_handles { h.join().ok(); }
        state_thread.join().ok();
    } else {
        cache_feeder.join().ok();
        for h in dl_handles { h.join().ok(); }
        for h in parse_handles { h.join().ok(); }
        state_thread.join().ok();

        pb_status.finish_and_clear();
        pb_main.finish();
    }

    let elapsed = t0.elapsed();
    let records = stats.records.load(Ordering::Relaxed);
    let errors = stats.parse_errors.load(Ordering::Relaxed);
    let dl_errors = stats.dl_errors.load(Ordering::Relaxed);
    let rate = records as f64 / elapsed.as_secs_f64().max(0.01);

    eprintln!();
    if shutdown.load(Ordering::Relaxed) {
        let done = stats.files_done.load(Ordering::Relaxed);
        eprintln!("[pmq] Interrupted. {done} files ({records} records) saved. Run again to resume.");
    } else {
        eprintln!(
            "[pmq] Done: {} records in {:.0}s ({:.0} rec/s)",
            records, elapsed.as_secs_f64(), rate,
        );
    }
    if errors > 0 {
        eprintln!("[pmq] ERROR: {errors} files failed to parse. Delete corrupt cached files and run again.");
    }
    if dl_errors > 0 { eprintln!("[pmq] {dl_errors} download errors (run again to retry)"); }
    eprintln!("[pmq] Query with: read_parquet('{}')", table_glob(&config.output_dir, "<table>"));
    Ok(())
}
