use std::collections::{BTreeMap, HashMap};
use std::path::PathBuf;

use clap::{Args, ValueEnum};
use obikindex::{GenomeInfo, KmerIndex};
use obikpartitionner::{AggOp, OutputCol};
use obisys::Reporter;
use tracing::info;

use super::predicate::matching_genome_indices;

// ── CLI types ─────────────────────────────────────────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum AggOpArg {
    Any,
    All,
    None,
    Sum,
    Min,
    Max,
}

impl From<AggOpArg> for AggOp {
    fn from(a: AggOpArg) -> Self {
        match a {
            AggOpArg::Any  => AggOp::Any,
            AggOpArg::All  => AggOp::All,
            AggOpArg::None => AggOp::None,
            AggOpArg::Sum  => AggOp::Sum,
            AggOpArg::Min  => AggOp::Min,
            AggOpArg::Max  => AggOp::Max,
        }
    }
}

#[derive(Args)]
pub struct SelectArgs {
    /// Source index directory
    pub source: PathBuf,

    /// Output index directory (mutually exclusive with --in-place)
    #[arg(long, conflicts_with = "in_place")]
    pub output: Option<PathBuf>,

    /// Rewrite the source index in-place (mutually exclusive with --output)
    #[arg(long)]
    pub in_place: bool,

    /// Define a named group: `<name>:<pred>` (repeatable; mutually exclusive with --aggregate-by)
    #[arg(long, value_name = "NAME:PRED", conflicts_with = "aggregate_by")]
    pub group: Vec<String>,

    /// Per-group aggregation operator: `<name>:<op>` (repeatable)
    #[arg(long, value_name = "NAME:OP")]
    pub group_op: Vec<String>,

    /// Auto-create one group per unique value of metadata key <KEY>
    #[arg(long, value_name = "KEY", conflicts_with = "group")]
    pub aggregate_by: Option<String>,

    /// Aggregation operator for all auto-generated groups
    #[arg(long, value_name = "OP")]
    pub aggregate_op: Option<AggOpArg>,

    /// Output columns in order: group names or genome labels, comma-separated
    #[arg(long, value_name = "COL,...", value_delimiter = ',')]
    pub select: Option<Vec<String>>,

    /// Minimum count to consider a genome as "carrying" the k-mer (logical ops only)
    #[arg(long, default_value = "0")]
    pub presence_threshold: u32,

    /// Overwrite existing output directory
    #[arg(short, long)]
    pub force: bool,
}

// ── Helpers ───────────────────────────────────────────────────────────────────

fn parse_name_value(s: &str, flag: &str) -> (String, String) {
    match s.find(':') {
        Some(pos) => (s[..pos].trim().to_string(), s[pos + 1..].to_string()),
        std::option::Option::None => {
            eprintln!("error in {flag}: expected <name>:<value>, got: {s}");
            std::process::exit(1);
        }
    }
}

fn parse_agg_op(s: &str) -> AggOp {
    match s.to_lowercase().as_str() {
        "any"  => AggOp::Any,
        "all"  => AggOp::All,
        "none" => AggOp::None,
        "sum"  => AggOp::Sum,
        "min"  => AggOp::Min,
        "max"  => AggOp::Max,
        other  => {
            eprintln!("unknown aggregation operator: {other}; valid: any, all, none, sum, min, max");
            std::process::exit(1);
        }
    }
}

fn default_op(src_is_count: bool) -> AggOp {
    if src_is_count { AggOp::Sum } else { AggOp::Any }
}

// ── build_specs ───────────────────────────────────────────────────────────────

/// Resolve CLI arguments into an ordered list of `OutputCol`.
///
/// Returns `(specs, output_presence)`.
fn build_specs(
    args: &SelectArgs,
    genomes: &[GenomeInfo],
    src_is_count: bool,
) -> (Vec<OutputCol>, bool) {
    // ── 1. Build group_indices: name → Vec<usize> ────────────────────────────
    // Also keep insertion order for the default `--select *` case.
    let mut group_order: Vec<String> = Vec::new();
    let mut group_indices: HashMap<String, Vec<usize>> = HashMap::new();

    if let Some(ref key) = args.aggregate_by {
        // One group per unique value of `key`, in sorted order.
        let mut value_to_indices: BTreeMap<String, Vec<usize>> = BTreeMap::new();
        for (i, g) in genomes.iter().enumerate() {
            if let Some(v) = g.meta.get(key) {
                value_to_indices.entry(v.clone()).or_default().push(i);
            }
        }
        for (v, idxs) in value_to_indices {
            group_order.push(v.clone());
            group_indices.insert(v, idxs);
        }
    } else {
        for raw in &args.group {
            let (name, pred) = parse_name_value(raw, "--group");
            let idxs = matching_genome_indices(&pred, genomes).unwrap_or_else(|e| {
                eprintln!("error in --group {name}: {e}");
                std::process::exit(1);
            });
            if !group_indices.contains_key(&name) {
                group_order.push(name.clone());
            }
            group_indices.insert(name, idxs);
        }
    }

    // ── 2. Build per-group ops ────────────────────────────────────────────────
    let global_op = args.aggregate_op.map(AggOp::from);
    let mut group_op: HashMap<String, AggOp> = HashMap::new();
    for raw in &args.group_op {
        let (name, op_str) = parse_name_value(raw, "--group-op");
        if !group_indices.contains_key(&name) {
            eprintln!("--group-op references undefined group: {name}");
            std::process::exit(1);
        }
        group_op.insert(name, parse_agg_op(&op_str));
    }

    // ── 3. Genome label → index map for pass-through columns ─────────────────
    let label_to_idx: HashMap<&str, usize> = genomes.iter().enumerate()
        .map(|(i, g)| (g.label.as_str(), i))
        .collect();

    // ── 4. Determine output column names ─────────────────────────────────────
    let col_names: Vec<String> = if let Some(ref sel) = args.select {
        sel.clone()
    } else if !group_order.is_empty() {
        group_order.clone()
    } else {
        // Identity: all genomes in original order
        genomes.iter().map(|g| g.label.clone()).collect()
    };

    // ── 5. Build OutputCol list ───────────────────────────────────────────────
    let mut specs: Vec<OutputCol> = Vec::with_capacity(col_names.len());

    for name in &col_names {
        if let Some(idxs) = group_indices.get(name) {
            let op = group_op.get(name)
                .copied()
                .or(global_op)
                .unwrap_or_else(|| default_op(src_is_count));
            specs.push(OutputCol { label: name.clone(), indices: idxs.clone(), op });
        } else if let Some(&idx) = label_to_idx.get(name.as_str()) {
            // Pass-through: single-element group with default op.
            let op = default_op(src_is_count);
            specs.push(OutputCol { label: name.clone(), indices: vec![idx], op });
        } else {
            eprintln!("--select: unknown column '{name}' (not a group name or genome label)");
            std::process::exit(1);
        }
    }

    if specs.is_empty() {
        eprintln!("select: no output columns defined");
        std::process::exit(1);
    }

    // ── 6. Determine output type ──────────────────────────────────────────────
    let output_presence = !src_is_count
        || specs.iter().all(|s| s.op.is_logical());

    (specs, output_presence)
}

// ── run ───────────────────────────────────────────────────────────────────────

pub fn run(args: SelectArgs) {
    if !args.in_place && args.output.is_none() {
        eprintln!("error: one of --output or --in-place must be specified");
        std::process::exit(1);
    }

    // Lock whichever directory actually gets written: the source itself in
    // --in-place mode, otherwise the (distinct) --output directory. Acquired
    // before opening the source so a concurrent writer can't slip in between
    // the open and the write below.
    let lock_target = if args.in_place { &args.source } else { args.output.as_ref().unwrap() };
    let _lock = obisys::DirLock::acquire(lock_target).unwrap_or_else(|e| {
        eprintln!("error locking {}: {e}", lock_target.display());
        std::process::exit(1);
    });

    let mut src = KmerIndex::open(&args.source).unwrap_or_else(|e| {
        eprintln!("error opening source index: {e}");
        std::process::exit(1);
    });

    let src_is_count = src.meta().config.with_counts;
    let (specs, output_presence) = build_specs(&args, &src.meta().genomes.clone(), src_is_count);

    info!(
        "select: {} genome(s) → {} output column(s), output={}",
        src.meta().genomes.len(),
        specs.len(),
        if output_presence { "presence" } else { "count" },
    );

    let mut rep = Reporter::new();

    if args.in_place {
        src.select_in_place(&specs, args.presence_threshold, output_presence, &mut rep)
            .unwrap_or_else(|e| {
                eprintln!("select error: {e}");
                std::process::exit(1);
            });
        rep.print();
        info!("selected in-place → {}", args.source.display());
    } else {
        let output = args.output.unwrap();
        KmerIndex::select(&output, &src, &specs, args.presence_threshold, output_presence, args.force, &mut rep)
            .unwrap_or_else(|e| {
                eprintln!("select error: {e}");
                std::process::exit(1);
            });
        rep.print();
        info!("selected index → {}", output.display());
    }
}
