use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::time::Instant;
use tracing::debug;

use obisys::progress_bar;

use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::ser::Serialize as EpSerialize;
use memmap2::Mmap;
use obicompactvec::PersistentCompactIntVecBuilder;
use obikseq::RoutableSuperKmer;
use obikseq::Sequence;
use obikseq::superkmer::SuperKmer;
use obiskio::{SKFileMeta, SKFileReader, SKFileWriter, SKResult};
use ptr_hash::{PtrHash, PtrHashParams, bucket_fn::CubicEps, hash::Xx64};
use rayon::prelude::*;
use remove_dir_all::remove_dir_all;
use sysinfo::System;

use niffler::Level;
use niffler::send::compression::Format;

use crate::kmer_sort::{chunk_size_from_ram, sort_unique_kmers};

type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;

pub struct KmerSpectrum {
    pub f0:     u64,
    pub f1:     u64,
    pub counts: BTreeMap<u32, u64>,
}

const SK_EXT: &str = "skmer.zst";
pub const PARTITIONS_SUBDIR: &str = "partitions";

pub struct KmerPartition {
    root_path: PathBuf,
    n_partitions: usize,
    partitions_mask: u64,
    kmer_size: usize,
    minimizer_size: usize,
    writers: Vec<Option<SKFileWriter>>,
    level: Level,
    closed: bool,
}

impl KmerPartition {
    pub fn create<P: AsRef<Path>>(
        path: P,
        n_bits: usize,
        kmer_size: usize,
        minimizer_size: usize,
        force: bool,
    ) -> SKResult<Self> {
        Self::create_with(path, n_bits, kmer_size, minimizer_size, Level::One, force)
    }

    pub fn create_with<P: AsRef<Path>>(
        path: P,
        n_bits: usize,
        kmer_size: usize,
        minimizer_size: usize,
        level: Level,
        force: bool,
    ) -> SKResult<Self> {
        let root_path = path.as_ref().to_owned();
        if root_path.exists() {
            if force {
                remove_dir_all(&root_path)?;
            } else {
                return Err(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    format!(
                        "{}: partition directory already exists",
                        root_path.display()
                    ),
                )
                .into());
            }
        }
        fs::create_dir_all(root_path.join(PARTITIONS_SUBDIR))?;
        let n_partitions = 1usize << n_bits;
        let writers = (0..n_partitions).map(|_| None).collect();
        let partition = Self {
            root_path,
            n_partitions,
            partitions_mask: (1u64 << n_bits) - 1,
            kmer_size,
            minimizer_size,
            writers,
            level,
            closed: false,
        };
        Ok(partition)
    }

    pub fn open_with_config<P: AsRef<Path>>(
        path: P,
        kmer_size: usize,
        minimizer_size: usize,
        n_bits: usize,
    ) -> SKResult<Self> {
        let root_path = path.as_ref().to_owned();
        if !root_path.exists() {
            return Err(io::Error::new(
                io::ErrorKind::NotFound,
                format!("{}: partition directory not found", root_path.display()),
            )
            .into());
        }
        let n_partitions = 1usize << n_bits;
        let writers = (0..n_partitions).map(|_| None).collect();
        Ok(Self {
            root_path,
            n_partitions,
            partitions_mask: (1u64 << n_bits) - 1,
            kmer_size,
            minimizer_size,
            writers,
            level: Level::One,
            closed: true,
        })
    }

    /// Route and write one super-kmer to its partition file.
    pub fn write(&mut self, rsk: RoutableSuperKmer) -> SKResult<()> {
        self.check_not_closed()?;
        let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
        let sk = rsk.into_superkmer();
        self.ensure_writer(partition)?.write(&sk)
    }

    /// Route and write a batch of super-kmers.
    pub fn write_batch(&mut self, rsks: Vec<RoutableSuperKmer>) -> SKResult<()> {
        self.check_not_closed()?;
        for rsk in rsks {
            let partition = (rsk.minimizer().seq_hash() & self.partitions_mask) as usize;
            let sk = rsk.into_superkmer();
            self.ensure_writer(partition)?.write(&sk)?;
        }
        Ok(())
    }

    pub fn flush(&mut self) -> SKResult<()> {
        self.check_not_closed()?;
        for writer in self.writers.iter_mut().flatten() {
            writer.flush()?;
        }
        Ok(())
    }

    pub fn close(&mut self) -> SKResult<()> {
        if self.closed {
            return Ok(());
        }
        self.closed = true;
        for writer in self.writers.iter_mut().flatten() {
            writer.close()?;
        }
        Ok(())
    }

    pub fn is_open(&self) -> bool {
        !self.closed
    }

    pub fn path(&self) -> &Path {
        &self.root_path
    }

    /// Path of partition `i` directory.
    pub fn part_dir(&self, i: usize) -> PathBuf {
        self.root_path.join(PARTITIONS_SUBDIR).join(format!("part_{i:05}"))
    }

    pub fn kmer_size(&self) -> usize {
        self.kmer_size
    }

    pub fn minimizer_size(&self) -> usize {
        self.minimizer_size
    }

    pub fn n_partitions(&self) -> usize {
        self.n_partitions
    }

    /// Deduplicate all `raw.{ext}` files in parallel, replacing each with a
    /// `dereplicated.{ext}` file where identical canonical sequences are merged
    /// and their counts summed.
    ///
    /// Each partition file is processed in two phases to bound memory use:
    ///
    /// 1. **Split** — the raw file is scattered into `2^temp_bits` temporary
    ///    files routed by `hash(canonical_seq) & temp_mask`.  Because duplicates
    ///    always share the same hash, they always land in the same temp file.
    /// 2. **Merge** — each temp file is loaded fully into a `HashMap`, counts
    ///    are accumulated in `u64` (no 24-bit overflow risk), and the result is
    ///    appended to `dereplicated.{ext}`.
    ///
    /// If a merged count exceeds the 24-bit header limit, the sequence is
    /// emitted as multiple records whose counts sum to the true total.
    ///
    /// `temp_bits` controls the split fan-out (`2^temp_bits` temp files per
    /// partition).  Higher values reduce per-temp-file memory at the cost of
    /// more temporary file descriptors — all managed by the global fd pool.
    pub fn dereplicate(&self) -> SKResult<()> {
        let level = self.level;
        let sys = System::new_all();
        // available_memory() can return 0 on macOS when the compressor page count exceeds
        // free+inactive+purgeable pages (sysinfo saturating_sub). Fall back to half of total.
        let available = match sys.available_memory() {
            0 => sys.total_memory() / 2,
            n => n,
        };
        let n_threads = rayon::current_num_threads().max(1) as u64;
        let available_per_thread = available / n_threads;

        let pb = progress_bar("dereplication", self.n_partitions as u64, "partitions");

        let results: Vec<SKResult<()>> = (0..self.n_partitions)
            .into_par_iter()
            .map(|i| {
                let dir = self.part_dir(i);
                if !dir.exists() {
                    pb.inc(1);
                    return Ok(());
                }
                let raw_path = dir.join(format!("raw.{SK_EXT}"));
                let t = Instant::now();
                let n_buckets = optimal_buckets(&raw_path, available_per_thread);
                let result = dereplicate_partition(&dir, level, n_buckets);
                pb.set_message(format!("last {:.0}ms", t.elapsed().as_millis()));
                pb.inc(1);
                result
            })
            .collect();

        pb.finish_and_clear();
        for r in results {
            r?;
        }
        Ok(())
    }

    /// For each partition that has a `dereplicated.{ext}` file:
    /// 1. Enumerates all unique canonical kmers (two passes over the file).
    /// 2. Builds a provisional MPHF (FMPHGO) over those kmers.
    /// 3. Writes a flat binary count file (`counts1.bin`, one `u32` per slot,
    ///    memory-mapped) accumulating kmer abundances from the superkmer counts.
    /// 4. Persists the MPHF to `mphf1.bin` for downstream use.
    ///
    /// Returns the aggregated `KmerSpectrum`. Per-partition spectrum files are
    /// deleted after aggregation unless `keep_partial` is true.
    ///
    /// Partitions are processed in parallel via Rayon (one task per thread).
    /// Peak memory per partition is ~80 MB, so n_threads partitions run simultaneously.
    pub fn count_kmer(&self, keep_partial: bool) -> SKResult<KmerSpectrum> {
        let sys = System::new_all();
        let available = match sys.available_memory() {
            0 => sys.total_memory() / 2,
            n => n,
        };
        let n_threads = rayon::current_num_threads().max(1) as u64;
        let chunk_kmers = chunk_size_from_ram(available / n_threads);

        let pb = progress_bar("counting", self.n_partitions as u64, "partitions");

        let results: Vec<SKResult<()>> = (0..self.n_partitions)
            .into_par_iter()
            .map(|i| {
                let dir = self.part_dir(i);
                let dedup_path = dir.join(format!("dereplicated.{SK_EXT}"));
                if !dedup_path.exists() {
                    pb.inc(1);
                    return Ok(());
                }
                let t = Instant::now();
                let result = count_partition(&dir, &dedup_path, chunk_kmers);
                pb.set_message(format!("last {:.0}ms", t.elapsed().as_millis()));
                pb.inc(1);
                result
            })
            .collect();

        pb.finish_and_clear();
        for r in results {
            r?;
        }

        // Aggregate per-partition spectra.
        let mut counts: BTreeMap<u32, u64> = BTreeMap::new();
        let mut f0: u64 = 0;
        let mut f1: u64 = 0;

        for i in 0..self.n_partitions {
            let path = self.part_dir(i).join("kmer_spectrum_raw.json");
            if !path.exists() {
                continue;
            }
            let v: serde_json::Value =
                serde_json::from_str(&fs::read_to_string(&path)?).map_err(io::Error::other)?;
            f0 += v["f0"].as_u64().unwrap_or(0);
            f1 += v["f1"].as_u64().unwrap_or(0);
            if let Some(obj) = v["spectrum"].as_object() {
                for (c_str, freq) in obj {
                    if let (Ok(c), Some(f)) = (c_str.parse::<u32>(), freq.as_u64()) {
                        *counts.entry(c).or_insert(0) += f;
                    }
                }
            }
            if !keep_partial {
                let _ = fs::remove_file(&path);
            }
        }

        Ok(KmerSpectrum { f0, f1, counts })
    }

    // ── private ───────────────────────────────────────────────────────────────

    fn check_not_closed(&self) -> SKResult<()> {
        if self.closed {
            Err(io::Error::new(io::ErrorKind::BrokenPipe, "write to closed KmerPartition").into())
        } else {
            Ok(())
        }
    }

    fn ensure_writer(&mut self, partition: usize) -> SKResult<&mut SKFileWriter> {
        if self.writers[partition].is_none() {
            let dir = self.root_path.join(PARTITIONS_SUBDIR).join(format!("part_{:05}", partition));
            fs::create_dir_all(&dir)?;
            let file_path = dir.join(format!("raw.{SK_EXT}"));
            let writer = SKFileWriter::create_with(file_path, Format::Zstd, self.level)?;
            self.writers[partition] = Some(writer);
        }
        Ok(self.writers[partition].as_mut().unwrap())
    }
}

// ── free helpers ─────────────────────────────────────────────────────────────

/// Estimate the number of in-memory buckets needed to deduplicate the partition
/// file at `raw_path` given `available_bytes` of free RAM.
///
/// Memory per HashMap entry:
///   key Box (1 + avg_seq_bytes) + SuperKmer header (4 B) + avg seq bytes + u64 count (8 B),
///   multiplied by 1.5 for hashbrown load-factor overhead.
///
/// Returns 1 if the partition fits comfortably in memory (no split needed).
/// Always returns a power of two.
/// Remove a SuperKmer file and its sidecar (if present).
fn remove_skmer_file(path: &Path) -> SKResult<()> {
    fs::remove_file(path)?;
    let sidecar = SKFileMeta::sidecar_path(path);
    match fs::remove_file(&sidecar) {
        Ok(()) => {}
        Err(e) if e.kind() == io::ErrorKind::NotFound => {}
        Err(e) => return Err(e.into()),
    }
    Ok(())
}

fn optimal_buckets(raw_path: &Path, available_bytes: u64) -> usize {
    // Use 60 % of available RAM to leave headroom for the rest of the process.
    let budget = (available_bytes as f64 * 0.60) as u64;

    let meta = match SKFileMeta::read(raw_path) {
        Ok(Some(m)) if m.instances > 0 => m,
        _ => return 1,
    };

    let avg_seq_bytes = ((meta.length_sum + meta.instances - 1) / meta.instances + 3) / 4;
    // SuperKmer: header (4 B) + Box<[u8]> ptr+len (16 B) + heap seq bytes; value: u64 (8 B); ×1.5 for hashbrown overhead.
    let bytes_per_entry = ((4 + 16 + avg_seq_bytes + 8) as f64 * 1.5) as u64;
    let estimated = meta.instances * bytes_per_entry;

    if estimated <= budget {
        debug!("Dereplication: estimated={estimated} budget={budget} n_temp=1");
        return 1;
    }

    // Round up to the next power of two.
    let n = (estimated + budget - 1) / budget;
    debug!("Dereplication: estimated={estimated} budget={budget} n_temp={n}");
    n.next_power_of_two() as usize
}


/// Maximum value that fits in the 24-bit COUNT field of a SuperKmer header.
const MAX_SK_COUNT: u64 = (1 << 24) - 1;

/// Deduplicate one partition directory in place (two-phase split + merge).
fn dereplicate_partition(dir: &Path, level: Level, n_temp: usize) -> SKResult<()> {
    let raw_path = dir.join(format!("raw.{SK_EXT}"));
    if !raw_path.exists() {
        return Ok(());
    }

    let out_path = dir.join(format!("dereplicated.{SK_EXT}"));
    let mut writer = SKFileWriter::create_with(&out_path, Format::Zstd, level)?;

    if n_temp == 1 {
        // ── Direct path: partition fits in memory, no split needed ────────────
        let map = load_bucket(&raw_path)?;
        remove_skmer_file(&raw_path)?;
        flush_map(map, &mut writer)?;
    } else {
        // ── Phase 1: split raw file into temp buckets ─────────────────────────
        let temp_mask = (n_temp as u64) - 1;
        let temp_paths: Vec<PathBuf> = (0..n_temp)
            .map(|j| dir.join(format!("temp_{j:04}.{SK_EXT}")))
            .collect();

        {
            let mut writers: Vec<SKFileWriter> = temp_paths
                .iter()
                .map(|p| SKFileWriter::create_with(p, Format::Zstd, level))
                .collect::<SKResult<_>>()?;

            let mut reader = SKFileReader::open(&raw_path)?;
            while let Some(sk) = reader.read()? {
                let bucket = (sk.seq_hash() & temp_mask) as usize;
                writers[bucket].write(&sk)?;
            }
            for w in &mut writers {
                w.close()?;
            }
        }
        remove_skmer_file(&raw_path)?;

        // ── Phase 2: merge each temp bucket into the output ───────────────────
        for temp_path in &temp_paths {
            let map = load_bucket(temp_path)?;
            remove_skmer_file(temp_path)?;
            flush_map(map, &mut writer)?;
        }
    }

    writer.close()?;
    Ok(())
}

/// Read a SuperKmer file into a deduplication map (already canonical).
fn load_bucket(path: &Path) -> SKResult<HashMap<SuperKmer, u64>> {
    let capacity = SKFileMeta::read(path)
        .ok()
        .flatten()
        .map(|m| m.instances as usize)
        .unwrap_or(0);
    let mut map: HashMap<SuperKmer, u64> = HashMap::with_capacity(capacity);
    let mut reader = SKFileReader::open(path)?;
    while let Some(sk) = reader.read()? {
        let count = sk.count() as u64;
        *map.entry(sk).or_insert(0) += count;
    }
    Ok(map)
}

/// Write all entries of a deduplication map to `writer`, splitting oversized counts.
fn flush_map(map: HashMap<SuperKmer, u64>, writer: &mut SKFileWriter) -> SKResult<()> {
    for (mut sk, mut total) in map {
        while total > MAX_SK_COUNT {
            sk.set_count(MAX_SK_COUNT as u32);
            writer.write(&sk)?;
            total -= MAX_SK_COUNT;
        }
        sk.set_count(total as u32);
        writer.write(&sk)?;
    }
    Ok(())
}

fn build_mphf(unique_path: &Path, f0: usize) -> io::Result<Mphf> {
    let file = fs::File::open(unique_path)?;
    let mmap = unsafe { Mmap::map(&file)? };
    let kmers: &[u64] = unsafe {
        std::slice::from_raw_parts(mmap.as_ptr() as *const u64, f0)
    };
    // Sequential constructor: the outer par_iter over partitions already saturates
    // the Rayon pool. new_from_par_iter would get no additional threads and adds
    // coordination overhead. try_new accesses the same mmap'd pages at zero extra cost.
    Mphf::try_new(kmers, PtrHashParams::<CubicEps>::default())
        .ok_or_else(|| io::Error::other("ptr_hash construction failed"))
}

fn count_partition(dir: &Path, dedup_path: &Path, chunk_kmers: usize) -> SKResult<()> {
    let unique_path = dir.join("sorted_unique.bin");
    let f0 = sort_unique_kmers(dedup_path, dir, &unique_path, chunk_kmers)?;
    if f0 == 0 {
        return Ok(());
    }
    debug!("{}: f0={f0} unique kmers sorted", dir.display());

    let mphf = build_mphf(&unique_path, f0)?;
    fs::remove_file(&unique_path)?;

    let counts_path = dir.join("counts1.bin");
    let mut builder = PersistentCompactIntVecBuilder::new(f0, &counts_path)?;

    {
        let mut reader = SKFileReader::open(dedup_path)?;
        while let Some(sk) = reader.read()? {
            let sk_count = sk.count();
            for kmer in sk.iter_canonical_kmers() {
                let slot = mphf.index(&kmer.raw());
                builder.set(slot, builder.get(slot).saturating_add(sk_count));
            }
        }
    }

    let mut spectrum: BTreeMap<u32, u64> = BTreeMap::new();
    for slot in 0..f0 {
        let c = builder.get(slot);
        if c > 0 {
            *spectrum.entry(c).or_insert(0) += 1;
        }
    }
    let f1: u64 = spectrum.iter().map(|(&c, &f)| c as u64 * f).sum();
    builder.close()?;

    let spectrum_map: BTreeMap<String, u64> = spectrum
        .iter()
        .map(|(&c, &f)| (format!("{c:010}"), f))
        .collect();
    serde_json::to_writer_pretty(
        fs::File::create(dir.join("kmer_spectrum_raw.json"))?,
        &serde_json::json!({ "f0": f0 as u64, "f1": f1, "spectrum": &spectrum_map }),
    )
    .map_err(io::Error::other)?;

    EpSerialize::store(&mphf, &dir.join("mphf1.bin"))
        .map_err(|e| io::Error::other(e.to_string()))?;

    Ok(())
}

impl Drop for KmerPartition {
    fn drop(&mut self) {
        let _ = self.close();
    }
}

// ── integration tests ─────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    use obikrope::Rope;
    use obikseq::SuperKmer;
    use obiskbuilder::build_superkmers;

    const K: usize = 11;
    const M: usize = 5;

    fn setup() {
        obikseq::params::set_k(K);
        obikseq::params::set_m(M);
    }

    /// Direct canonical k-mer counts from ASCII sequences — ground truth.
    fn direct_counts(seqs: &[&[u8]]) -> (u64, u64) {
        let mut counts: HashMap<Vec<u8>, u64> = HashMap::new();
        for seq in seqs {
            for i in 0..seq.len().saturating_sub(K - 1) {
                let km = SuperKmer::from_ascii(&seq[i..i + K]).to_ascii();
                *counts.entry(km).or_insert(0) += 1;
            }
        }
        let f0 = counts.len() as u64;
        let f1: u64 = counts.values().sum();
        (f0, f1)
    }

    /// Run the full pipeline on a list of sequences and return (f0, f1) from
    /// the `kmer_spectrum_raw.json` produced by `count_partition`.
    fn pipeline_counts(seqs: &[&[u8]]) -> (u64, u64) {
        setup();

        let mut rope_data: Vec<u8> = Vec::new();
        for seq in seqs {
            rope_data.extend_from_slice(seq);
            rope_data.push(0x00);
        }
        let mut rope = Rope::new(None);
        rope.push(rope_data);

        let superkmers: Vec<_> = build_superkmers(rope, K, 1, 0.0);

        let dir = tempfile::tempdir().unwrap();
        let mut kp = KmerPartition::create(dir.path(), 0, K, M, true).unwrap();
        kp.write_batch(superkmers).unwrap();
        kp.close().unwrap();
        kp.dereplicate().unwrap();

        let part_dir = dir.path().join(PARTITIONS_SUBDIR).join("part_00000");
        let dedup_path = part_dir.join("dereplicated.skmer.zst");
        if !dedup_path.exists() {
            return (0, 0);
        }
        count_partition(&part_dir, &dedup_path, 1 << 20).unwrap();

        let spec: serde_json::Value = serde_json::from_reader(
            fs::File::open(part_dir.join("kmer_spectrum_raw.json")).unwrap(),
        ).unwrap();
        let f0 = spec["f0"].as_u64().unwrap_or(0);
        let f1 = spec["f1"].as_u64().unwrap_or(0);
        (f0, f1)
    }

    #[test]
    fn single_sequence_f0_f1_match() {
        let seqs: &[&[u8]] = &[b"ACGTACGTACGTACGTACGT"];
        let (ef0, ef1) = direct_counts(seqs);
        let (gf0, gf1) = pipeline_counts(seqs);
        assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
        assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
    }

    #[test]
    fn two_sequences_f0_f1_match() {
        let seqs: &[&[u8]] = &[
            b"ACGTACGTACGTACGTACGT",
            b"TGCATGCATGCATGCATGCA",
        ];
        let (ef0, ef1) = direct_counts(seqs);
        let (gf0, gf1) = pipeline_counts(seqs);
        assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
        assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
    }

    #[test]
    fn repeated_sequence_f1_doubles() {
        let seq = b"ACGTACGTACGTACGTACGT";
        let seqs: &[&[u8]] = &[seq, seq];
        let (ef0, ef1) = direct_counts(seqs);
        let (gf0, gf1) = pipeline_counts(seqs);
        assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
        assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
    }

    #[test]
    fn many_sequences_f0_f1_match() {
        // 20 distinct sequences of length 40 — forces multiple super-kmers and
        // multiple minimizer boundaries per sequence.
        let bases = b"ACGT";
        let seqs: Vec<Vec<u8>> = (0..20u32)
            .map(|i| (0..40).map(|j| bases[((i * 7 + j * 3) % 4) as usize]).collect())
            .collect();
        let seq_refs: Vec<&[u8]> = seqs.iter().map(|v| v.as_slice()).collect();
        let (ef0, ef1) = direct_counts(&seq_refs);
        let (gf0, gf1) = pipeline_counts(&seq_refs);
        assert_eq!(gf0, ef0, "f0 wrong: expected {ef0}, got {gf0}");
        assert_eq!(gf1, ef1, "f1 wrong: expected {ef1}, got {gf1}");
    }
}
