use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::Path;
use std::time::Duration;

use indicatif::{ProgressBar, ProgressStyle};
use obisys::{Reporter, Stage};
use rayon::prelude::*;
use tracing::info;

use obilayeredmap::IndexMode;

use crate::error::{OKIError, OKIResult};
use crate::index::KmerIndex;
use crate::meta::{GenomeInfo, IndexMeta};
use crate::state::IndexState;

pub use obikpartitionner::MergeMode;

impl KmerIndex {
    /// Merge `sources` into a new index at `output`.
    ///
    /// All sources must be in `Indexed` state and share the same `kmer_size`,
    /// `minimizer_size`, and `n_partitions`. Count mode additionally requires
    /// every source to have `with_counts = true`.
    ///
    /// Genome labels must be unique across all sources. If `rename_duplicates`
    /// is true, repeated labels are disambiguated by appending `.1`, `.2`, …
    /// to the second and subsequent occurrences. Otherwise a
    /// `DuplicateGenomeLabel` error is returned on the first conflict.
    pub fn merge<P: AsRef<Path>>(
        output: P,
        sources: &[&KmerIndex],
        mode: MergeMode,
        force: bool,
        rename_duplicates: bool,
        rep: &mut Reporter,
    ) -> OKIResult<Self> {
        let output = output.as_ref();

        if sources.is_empty() {
            return Err(OKIError::Io(io::Error::new(
                io::ErrorKind::InvalidInput,
                "merge requires at least one source index",
            )));
        }

        // ── Validate config compatibility ─────────────────────────────────────
        let ref0 = sources[0];
        for src in sources {
            if src.state() != IndexState::Indexed {
                return Err(OKIError::NotIndexed(src.root_path.clone()));
            }
            if src.kmer_size()      != ref0.kmer_size()
            || src.minimizer_size() != ref0.minimizer_size()
            || src.n_partitions()   != ref0.n_partitions()
            {
                return Err(OKIError::IncompatibleConfig);
            }
            if mode == MergeMode::Count && !src.meta.config.with_counts {
                return Err(OKIError::MismatchedMode);
            }
        }

        // ── Validate evidence compatibility ───────────────────────────────────
        let evidence = validate_evidence_compat(sources)?;

        // ── Compute final genome labels (rename duplicates if requested) ───────
        let (source_labels, all_genomes) = compute_labels(sources, rename_duplicates)?;

        // ── Prepare output directory ──────────────────────────────────────────
        if output.exists() {
            if force {
                fs::remove_dir_all(output)?;
            } else {
                return Err(OKIError::Io(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    format!("{}: output directory already exists", output.display()),
                )));
            }
        }

        // ── Bootstrap: copy first source to output ────────────────────────────
        info!(
            "bootstrap: copying {} → {} ({} genome(s))",
            sources[0].root_path.display(),
            output.display(),
            sources[0].meta.genomes.len(),
        );
        let t = Stage::start("bootstrap");
        let pb = spinner("bootstrap — copying index …");
        copy_dir_all(&sources[0].root_path, output)?;

        // Rewrite index.meta with final genome labels and the effective mode.
        let mut meta = IndexMeta::read(output).map_err(OKIError::Io)?;
        meta.genomes = all_genomes;
        meta.config.with_counts = mode == MergeMode::Count;
        meta.config.evidence = evidence.clone();
        meta.write(output)?;

        // In presence/absence mode, purge counts/ directories inherited from
        // source_0 — they are stale data from the source's count index.
        if mode == MergeMode::Presence {
            remove_dirs_named(output, "counts")?;
        }
        pb.finish_and_clear();
        rep.push(t.stop());

        // Rebuild spectrums/ from all sources using the (possibly renamed) labels.
        // Drop the spectrums/ that were copied from source_0 and rebuild from scratch.
        info!("rebuilding spectrums for {} source(s)", sources.len());
        let t = Stage::start("spectrums");
        let pb = spinner("spectrums — copying …");
        let spectrums_dir = output.join("spectrums");
        if spectrums_dir.exists() {
            fs::remove_dir_all(&spectrums_dir)?;
        }
        for (src, new_labels) in sources.iter().zip(&source_labels) {
            let old_labels: Vec<String> = src.meta.genomes.iter().map(|g| g.label.clone()).collect();
            copy_spectrums(&src.root_path, output, &old_labels, new_labels)?;
        }
        pb.finish_and_clear();
        rep.push(t.stop());

        // Open the destination index.
        let dst = KmerIndex::open(output)?;
        let n_partitions = dst.n_partitions();
        let n_dst_genomes = sources[0].meta.genomes.len();

        // ── Merge each subsequent source partition-by-partition ───────────────
        let remaining_sources: Vec<&KmerIndex> = sources[1..].to_vec();
        if !remaining_sources.is_empty() {
            let n_src_genomes: usize = remaining_sources.iter().map(|s| s.meta.genomes.len()).sum();
            info!(
                "merging {} partition(s) × {} additional source genome(s) into {} destination genome(s)",
                n_partitions, n_src_genomes, n_dst_genomes,
            );
            let t = Stage::start("merge_partitions");
            let pb = partition_bar(n_partitions as u64);

            let dst_partition = &dst.partition;
            let block_bits = dst.meta.config.block_bits;

            let errors: Vec<obiskio::SKError> = (0..n_partitions)
                .into_par_iter()
                .filter_map(|i| {
                    let srcs: Vec<(&obikpartitionner::KmerPartition, usize)> =
                        remaining_sources.iter().map(|s| (&s.partition, s.meta.genomes.len())).collect();
                    let result = dst_partition.merge_partition(i, &srcs, mode, n_dst_genomes, block_bits, &evidence).err();
                    pb.inc(1);
                    result
                })
                .collect();

            pb.finish_and_clear();
            if let Some(e) = errors.into_iter().next() {
                return Err(OKIError::Partition(e));
            }

            rep.push(t.stop());
        }

        // Re-open to get the updated state.
        KmerIndex::open(output)
    }
}

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

/// Compute the final genome label lists for all sources.
///
/// Returns `(per_source_labels, all_genomes_flat)`.
/// The first occurrence of a label keeps the original name. Subsequent
/// occurrences receive `.1`, `.2`, … suffixes when `rename_duplicates` is true,
/// or trigger a `DuplicateGenomeLabel` error otherwise.
fn compute_labels(
    sources: &[&KmerIndex],
    rename_duplicates: bool,
) -> OKIResult<(Vec<Vec<String>>, Vec<GenomeInfo>)> {
    let mut seen: HashMap<String, usize> = HashMap::new();
    let mut source_labels: Vec<Vec<String>> = Vec::with_capacity(sources.len());
    let mut all_genomes: Vec<GenomeInfo> = Vec::new();

    for src in sources {
        let mut labels = Vec::with_capacity(src.meta.genomes.len());
        for genome in &src.meta.genomes {
            let label = &genome.label;
            let count = seen.entry(label.clone()).or_insert(0);
            let new_label = if *count == 0 {
                label.clone()
            } else if rename_duplicates {
                format!("{label}.{count}")
            } else {
                return Err(OKIError::DuplicateGenomeLabel(label.clone()));
            };
            *count += 1;
            labels.push(new_label.clone());
            all_genomes.push(GenomeInfo { label: new_label, meta: genome.meta.clone() });
        }
        source_labels.push(labels);
    }

    Ok((source_labels, all_genomes))
}

/// Copy spectrum JSON files from `src_root/spectrums/` to `dst_root/spectrums/`,
/// mapping each `old_labels[i]` filename to `new_labels[i]`.
fn copy_spectrums(
    src_root: &Path,
    dst_root: &Path,
    old_labels: &[String],
    new_labels: &[String],
) -> io::Result<()> {
    let src_dir = src_root.join("spectrums");
    let dst_dir = dst_root.join("spectrums");
    fs::create_dir_all(&dst_dir)?;
    for (old, new) in old_labels.iter().zip(new_labels.iter()) {
        let src_file = src_dir.join(format!("{old}.json"));
        if src_file.exists() {
            fs::copy(&src_file, dst_dir.join(format!("{new}.json")))?;
        }
    }
    Ok(())
}

/// Recursively remove every directory named `name` under `root`.
fn remove_dirs_named(root: &Path, name: &str) -> io::Result<()> {
    for entry in fs::read_dir(root)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            if path.file_name().and_then(|n| n.to_str()) == Some(name) {
                fs::remove_dir_all(&path)?;
            } else {
                remove_dirs_named(&path, name)?;
            }
        }
    }
    Ok(())
}

fn spinner(msg: &'static str) -> ProgressBar {
    let pb = ProgressBar::new_spinner();
    pb.set_style(
        ProgressStyle::with_template("{spinner} {msg}  {elapsed}")
            .unwrap()
            .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]),
    );
    pb.set_message(msg);
    pb.enable_steady_tick(Duration::from_millis(100));
    pb
}

fn partition_bar(n: u64) -> ProgressBar {
    let pb = ProgressBar::new(n);
    pb.set_style(
        ProgressStyle::with_template(
            "{spinner} merge — {bar:40.cyan/blue} {pos}/{len} partitions  {elapsed}",
        )
        .unwrap()
        .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]),
    );
    pb.enable_steady_tick(Duration::from_millis(100));
    pb
}

/// Check that all sources share the same evidence kind.
///
/// Rules:
/// - all `Exact`                        → OK, returns `Exact`
/// - all `Approx { b, z }` same params  → OK, returns `Approx { b, z }`
/// - mixed exact/approx or different approx params → `IncompatibleEvidence`
fn validate_evidence_compat(sources: &[&KmerIndex]) -> OKIResult<IndexMode> {
    let ref_ev = &sources[0].meta.config.evidence;
    for src in &sources[1..] {
        let ev = &src.meta.config.evidence;
        let compat = match (ref_ev, ev) {
            (IndexMode::Exact, IndexMode::Exact) => true,
            (IndexMode::Approx { b: b1, z: z1 },
             IndexMode::Approx { b: b2, z: z2 }) => b1 == b2 && z1 == z2,
            (IndexMode::Hybrid { b: b1, z: z1 },
             IndexMode::Hybrid { b: b2, z: z2 }) => b1 == b2 && z1 == z2,
            _ => false,
        };
        if !compat {
            return Err(OKIError::IncompatibleEvidence(format!(
                "source {:?} has evidence {:?}, expected {:?} — \
                 convert all sources to the same evidence kind first \
                 (use the `reindex` command)",
                src.root_path.display(), ev, ref_ev,
            )));
        }
    }
    Ok(ref_ev.clone())
}

fn copy_dir_all(src: &Path, dst: &Path) -> io::Result<()> {
    fs::create_dir_all(dst)?;
    for entry in fs::read_dir(src)? {
        let entry = entry?;
        let src_path = entry.path();
        let dst_path = dst.join(entry.file_name());
        if src_path.is_dir() {
            copy_dir_all(&src_path, &dst_path)?;
        } else {
            fs::copy(&src_path, &dst_path)?;
        }
    }
    Ok(())
}
