use std::collections::HashSet;
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;

use clap::Args;
use obikindex::KmerIndex;
use tracing::info;

#[derive(Args)]
pub struct AnnotateArgs {
    /// Index directory to annotate (modified in-place)
    pub index: PathBuf,

    /// CSV file with genome metadata (must contain an id column)
    #[arg(long)]
    pub csv: Option<PathBuf>,

    /// CSV field separator
    #[arg(long, default_value = ",")]
    pub sep: char,

    /// Name of the column that contains genome labels
    #[arg(long, default_value = "id")]
    pub id_col: String,

    /// Value that means "delete / absent" (removes existing key if present)
    #[arg(long, default_value = "NA")]
    pub na_value: String,

    /// Do not overwrite existing metadata keys
    #[arg(long)]
    pub no_overwrite: bool,

    /// Dump all genome metadata as CSV (stdout) instead of reading a CSV
    #[arg(long)]
    pub dump: bool,
}

pub fn run(args: AnnotateArgs) {
    if args.dump {
        run_dump(&args);
    } else {
        run_annotate(&args);
    }
}

fn run_dump(args: &AnnotateArgs) {
    let idx = open_index(&args.index);

    let genomes = &idx.meta().genomes;

    // Collect all keys in stable order (sorted for determinism)
    let mut key_set: HashSet<String> = HashSet::new();
    for g in genomes {
        for k in g.meta.keys() {
            key_set.insert(k.clone());
        }
    }
    let mut keys: Vec<String> = key_set.into_iter().collect();
    keys.sort();

    let stdout = io::stdout();
    let mut out = BufWriter::new(stdout.lock());

    // Header
    write!(out, "id").unwrap();
    for k in &keys {
        write!(out, "{}{k}", args.sep).unwrap();
    }
    writeln!(out).unwrap();

    // Rows
    for g in genomes {
        write!(out, "{}", g.label).unwrap();
        for k in &keys {
            let v = g.meta.get(k).map(|s| s.as_str()).unwrap_or("NA");
            write!(out, "{}{v}", args.sep).unwrap();
        }
        writeln!(out).unwrap();
    }
}

fn run_annotate(args: &AnnotateArgs) {
    let csv_path = match &args.csv {
        Some(p) => p.clone(),
        None => {
            eprintln!("error: --csv is required unless --dump is used");
            std::process::exit(1);
        }
    };

    let mut idx = open_index(&args.index);

    // Build a label → genome index position map
    let label_to_pos: std::collections::HashMap<String, usize> = idx
        .meta()
        .genomes
        .iter()
        .enumerate()
        .map(|(i, g)| (g.label.clone(), i))
        .collect();

    let sep = args.sep as u8;
    let mut rdr = csv::ReaderBuilder::new()
        .delimiter(sep)
        .from_path(&csv_path)
        .unwrap_or_else(|e| {
            eprintln!("error opening {}: {e}", csv_path.display());
            std::process::exit(1);
        });

    let headers = rdr.headers().unwrap_or_else(|e| {
        eprintln!("error reading CSV headers: {e}");
        std::process::exit(1);
    }).clone();

    let id_col_idx = headers.iter().position(|h| h == args.id_col).unwrap_or_else(|| {
        eprintln!("error: id column '{}' not found in CSV", args.id_col);
        std::process::exit(1);
    });

    let meta_cols: Vec<(usize, String)> = headers
        .iter()
        .enumerate()
        .filter(|(i, _)| *i != id_col_idx)
        .map(|(i, h)| (i, h.to_string()))
        .collect();

    let mut updated = 0usize;
    let mut skipped = 0usize;

    for result in rdr.records() {
        let record = result.unwrap_or_else(|e| {
            eprintln!("error reading CSV record: {e}");
            std::process::exit(1);
        });

        let label = record.get(id_col_idx).unwrap_or("").to_string();
        let pos = match label_to_pos.get(&label) {
            Some(&p) => p,
            None => {
                skipped += 1;
                continue;
            }
        };

        let genome = &mut idx.meta_mut().genomes[pos];
        for (col_idx, key) in &meta_cols {
            let val = record.get(*col_idx).unwrap_or("");
            if val == args.na_value {
                genome.meta.remove(key);
            } else if args.no_overwrite && genome.meta.contains_key(key) {
                // skip
            } else {
                genome.meta.insert(key.clone(), val.to_string());
            }
        }
        updated += 1;
    }

    idx.meta_mut().write(&args.index).unwrap_or_else(|e| {
        eprintln!("error writing index metadata: {e}");
        std::process::exit(1);
    });

    info!("annotated {updated} genome(s), skipped {skipped} CSV row(s) with unknown label");
}

fn open_index(path: &PathBuf) -> KmerIndex {
    KmerIndex::open(path).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    })
}
