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

use clap::Args;
use obikindex::{validate_label, IndexBitsPerKmer, KmerIndex};
use tracing::info;

#[derive(Args)]
pub struct UtilsArgs {
    /// Index directories to operate on (one or more)
    #[arg(required = true, num_args = 1..)]
    pub indexes: Vec<PathBuf>,

    /// Set a new genome label: NEW_LABEL=OLD_LABEL  (single-index only)
    #[arg(long, value_name = "NEW=OLD")]
    pub new_label: Option<String>,

    /// Add missing layer_meta.json files to each layer  (single-index only)
    #[arg(long)]
    pub upgrade_index: bool,

    /// Print bits-per-kmer statistics  (single-index only)
    #[arg(long)]
    pub bits_per_kmer: bool,

    /// Print per-genome k-mer counts as CSV  (single-index only)
    #[arg(long)]
    pub stats: bool,

    /// Print partition size distribution report (accepts multiple indexes)
    #[arg(long)]
    pub partition_stats: bool,

    /// Write per-(partition, source) raw data as CSV to FILE  (used with --partition-stats)
    #[arg(long, value_name = "FILE")]
    pub csv: Option<PathBuf>,
}

pub fn run(args: UtilsArgs) {
    let mut any = false;

    if let Some(spec) = &args.new_label {
        any = true;
        run_rename(single_index(&args), spec);
    }

    if args.upgrade_index {
        any = true;
        run_upgrade_index(single_index(&args));
    }

    if args.bits_per_kmer {
        any = true;
        run_bits_per_kmer(single_index(&args));
    }

    if args.stats {
        any = true;
        run_stats(single_index(&args));
    }

    if args.partition_stats {
        any = true;
        run_partition_stats(&args.indexes, args.csv.as_deref());
    }

    if !any {
        eprintln!(
            "utils: no operation specified. \
             Available: --new-label, --upgrade-index, --bits-per-kmer, --stats, --partition-stats"
        );
        std::process::exit(1);
    }
}

// ── helpers ───────────────────────────────────────────────────────────────────

fn single_index(args: &UtilsArgs) -> &PathBuf {
    if args.indexes.len() > 1 {
        eprintln!("utils: this option requires exactly one index (got {})", args.indexes.len());
        std::process::exit(1);
    }
    &args.indexes[0]
}

// ── --partition-stats ─────────────────────────────────────────────────────────

/// Per-partition, per-source byte count of all unitigs.bin files summed across layers.
struct PartRow {
    partition: usize,
    source: String,
    bytes: u64,
}

fn collect_rows(indexes: &[PathBuf]) -> Vec<PartRow> {
    let mut rows = Vec::new();
    for path in indexes {
        let idx = KmerIndex::open(path).unwrap_or_else(|e| {
            eprintln!("error opening index {}: {e}", path.display());
            std::process::exit(1);
        });
        let name = path
            .file_name()
            .map(|n| n.to_string_lossy().into_owned())
            .unwrap_or_else(|| path.display().to_string());
        let n_parts = idx.n_partitions();
        for i in 0..n_parts {
            let mut bytes = 0u64;
            for l in 0.. {
                let p = idx.layer_unitigs_path(i, l);
                if !p.exists() {
                    break;
                }
                if let Ok(m) = std::fs::metadata(&p) {
                    bytes += m.len();
                }
            }
            rows.push(PartRow { partition: i, source: name.clone(), bytes });
        }
    }
    rows
}

/// Sum bytes per partition across all sources.
fn partition_totals(rows: &[PartRow], n_parts: usize) -> Vec<u64> {
    let mut totals = vec![0u64; n_parts];
    for r in rows {
        totals[r.partition] += r.bytes;
    }
    totals
}

fn stats_summary(totals: &[u64]) -> (u64, u64, f64, f64, u64, u64, u64) {
    let mut sorted = totals.to_vec();
    sorted.sort_unstable();
    let n = sorted.len();
    let min = sorted[0];
    let max = sorted[n - 1];
    let mean = sorted.iter().sum::<u64>() as f64 / n as f64;
    let median = if n % 2 == 0 {
        (sorted[n / 2 - 1] + sorted[n / 2]) as f64 / 2.0
    } else {
        sorted[n / 2] as f64
    };
    let p95 = sorted[(n as f64 * 0.95) as usize];
    let p99 = sorted[(n as f64 * 0.99) as usize];
    let variance = sorted
        .iter()
        .map(|&v| (v as f64 - mean).powi(2))
        .sum::<f64>()
        / n as f64;
    let std_dev = variance.sqrt();
    (min, max, mean, median, p95, p99, std_dev as u64)
}

fn human_bytes(b: u64) -> String {
    if b >= 1 << 30 {
        format!("{:.1} GB", b as f64 / (1u64 << 30) as f64)
    } else if b >= 1 << 20 {
        format!("{:.1} MB", b as f64 / (1u64 << 20) as f64)
    } else if b >= 1 << 10 {
        format!("{:.1} KB", b as f64 / (1u64 << 10) as f64)
    } else {
        format!("{b} B")
    }
}

fn ascii_histogram(totals: &[u64], n_buckets: usize, bar_width: usize) -> String {
    let min = *totals.iter().min().unwrap();
    let max = *totals.iter().max().unwrap();
    if min == max {
        return format!("  (all partitions identical: {})\n", human_bytes(min));
    }

    let bucket_size = (max - min).max(1) as f64 / n_buckets as f64;
    let mut counts = vec![0usize; n_buckets];
    for &v in totals {
        let b = (((v - min) as f64 / bucket_size) as usize).min(n_buckets - 1);
        counts[b] += 1;
    }

    let max_count = *counts.iter().max().unwrap();
    let mut out = String::new();
    for (i, &c) in counts.iter().enumerate() {
        let lo = min + (i as f64 * bucket_size) as u64;
        let hi = min + ((i + 1) as f64 * bucket_size) as u64;
        let bar_len = if max_count > 0 { c * bar_width / max_count } else { 0 };
        let bar = "█".repeat(bar_len);
        out.push_str(&format!(
            "  {:>8} – {:>8} │{:<width$} {}\n",
            human_bytes(lo),
            human_bytes(hi),
            bar,
            c,
            width = bar_width
        ));
    }
    out
}

fn run_partition_stats(indexes: &[PathBuf], csv_path: Option<&std::path::Path>) {
    let rows = collect_rows(indexes);
    if rows.is_empty() {
        eprintln!("partition-stats: no data found");
        std::process::exit(1);
    }

    let n_parts = rows.iter().map(|r| r.partition).max().unwrap() + 1;
    let totals = partition_totals(&rows, n_parts);
    let (min, max, mean, median, p95, p99, std_dev) = stats_summary(&totals);

    // outliers: > median + 1.5 × IQR (approximate via > 1.5 × median as fallback)
    let mut sorted_t = totals.clone();
    sorted_t.sort_unstable();
    let q1 = sorted_t[n_parts / 4] as f64;
    let q3 = sorted_t[3 * n_parts / 4] as f64;
    let iqr = q3 - q1;
    let outlier_threshold = q3 + 1.5 * iqr;

    let mut out = String::new();
    out.push_str("# Partition size report\n\n");
    out.push_str(&format!(
        "Sources: {}  \nPartitions: {}  \n\n",
        indexes.len(),
        n_parts
    ));

    out.push_str("## Summary statistics (total unitigs.bin bytes per partition, sum across sources)\n\n");
    out.push_str("| Stat | Value |\n|---|---|\n");
    out.push_str(&format!("| min    | {} |\n", human_bytes(min)));
    out.push_str(&format!("| max    | {} |\n", human_bytes(max)));
    out.push_str(&format!("| mean   | {} |\n", human_bytes(mean as u64)));
    out.push_str(&format!("| median | {} |\n", human_bytes(median as u64)));
    out.push_str(&format!("| p95    | {} |\n", human_bytes(p95)));
    out.push_str(&format!("| p99    | {} |\n", human_bytes(p99)));
    out.push_str(&format!("| std    | {} |\n", human_bytes(std_dev)));
    out.push_str(&format!("| max/median ratio | {:.2}× |\n\n", max as f64 / median));

    out.push_str("## Histogram\n\n```\n");
    out.push_str(&ascii_histogram(&totals, 30, 40));
    out.push_str("```\n\n");

    let outliers: Vec<(usize, u64)> = totals
        .iter()
        .enumerate()
        .filter(|(_, v)| **v as f64 > outlier_threshold)
        .map(|(i, v)| (i, *v))
        .collect();

    if outliers.is_empty() {
        out.push_str("## Outliers\n\nNone (threshold: Q3 + 1.5×IQR = ");
        out.push_str(&human_bytes(outlier_threshold as u64));
        out.push_str(").\n");
    } else {
        out.push_str(&format!(
            "## Outliers (> Q3 + 1.5×IQR = {})\n\n| Partition | Total size | Ratio to median |\n|---|---|---|\n",
            human_bytes(outlier_threshold as u64)
        ));
        for (i, v) in &outliers {
            out.push_str(&format!(
                "| {} | {} | {:.2}× |\n",
                i,
                human_bytes(*v),
                *v as f64 / median
            ));
        }
        out.push('\n');
    }

    print!("{out}");

    if let Some(csv_out) = csv_path {
        let file = std::fs::File::create(csv_out).unwrap_or_else(|e| {
            eprintln!("error creating CSV file {}: {e}", csv_out.display());
            std::process::exit(1);
        });
        let mut w = io::BufWriter::new(file);
        writeln!(w, "partition,source,bytes").unwrap();
        for r in &rows {
            writeln!(w, "{},{},{}", r.partition, r.source, r.bytes).unwrap();
        }
        eprintln!("CSV written to {}", csv_out.display());
    }
}

// ── existing single-index operations ─────────────────────────────────────────

fn run_stats(index_path: &PathBuf) {
    let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    });
    let (total, per_genome) = idx.genome_kmer_counts().unwrap_or_else(|e| {
        eprintln!("error computing stats: {e}");
        std::process::exit(1);
    });
    println!("genome,n_kmers");
    for (g, &n) in idx.meta().genomes.iter().zip(per_genome.iter()) {
        println!("{},{}", g.label, n);
    }
    println!("total,{total}");
}

fn run_bits_per_kmer(index_path: &PathBuf) {
    let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    });
    let stats: IndexBitsPerKmer = idx.bits_per_kmer().unwrap_or_else(|e| {
        eprintln!("error computing bits/kmer: {e}");
        std::process::exit(1);
    });
    println!("k-mers    : {}", stats.n_kmers);
    println!("genomes   : {}", stats.n_genomes);
    println!("mphf      : {:6.2} bits/kmer", stats.mphf);
    println!("evidence  : {:6.2} bits/kmer", stats.evidence);
    println!(
        "matrix    : {:6.2} bits/kmer  ({:.2} bits/kmer/genome)",
        stats.matrix, stats.matrix_per_genome
    );
    println!("total     : {:6.2} bits/kmer", stats.total);
}

fn run_upgrade_index(index_path: &PathBuf) {
    let idx = KmerIndex::open(index_path).unwrap_or_else(|e| {
        eprintln!("error opening index: {e}");
        std::process::exit(1);
    });
    idx.upgrade_layer_meta().unwrap_or_else(|e| {
        eprintln!("upgrade error: {e}");
        std::process::exit(1);
    });
    info!("upgrade-index: layer_meta.json written to all layers that were missing it");
}

fn run_rename(index_path: &PathBuf, spec: &str) {
    let (old_label, new_label) = parse_rename_spec(spec);

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

    let pos = idx
        .meta()
        .genomes
        .iter()
        .position(|g| g.label == old_label)
        .unwrap_or_else(|| {
            eprintln!("error: genome '{old_label}' not found in index");
            std::process::exit(1);
        });

    validate_label(&new_label).unwrap_or_else(|e| {
        eprintln!("error: --new-label: {e}");
        std::process::exit(1);
    });

    if idx.meta().genomes.iter().any(|g| g.label == new_label) {
        eprintln!("error: label '{new_label}' already exists in index");
        std::process::exit(1);
    }

    idx.meta_mut().genomes[pos].label = new_label.clone();
    idx.meta_mut().write(index_path).unwrap_or_else(|e| {
        eprintln!("error writing index metadata: {e}");
        std::process::exit(1);
    });

    let spectrums_dir = index_path.join("spectrums");
    let old_spectrum = spectrums_dir.join(format!("{old_label}.json"));
    let new_spectrum = spectrums_dir.join(format!("{new_label}.json"));
    if old_spectrum.exists() {
        std::fs::rename(&old_spectrum, &new_spectrum).unwrap_or_else(|e| {
            eprintln!("warning: could not rename spectrum file: {e}");
        });
    }

    info!("renamed genome '{old_label}' → '{new_label}'");
}

fn parse_rename_spec(spec: &str) -> (String, String) {
    let eq = spec.find('=').unwrap_or_else(|| {
        eprintln!("error: --new-label expects NEW_LABEL=OLD_LABEL, got '{spec}'");
        std::process::exit(1);
    });
    let new = spec[..eq].trim().to_string();
    let old = spec[eq + 1..].trim().to_string();
    if old.is_empty() || new.is_empty() {
        eprintln!("error: --new-label: both old and new labels must be non-empty");
        std::process::exit(1);
    }
    (old, new)
}
