use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::fs;
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};

use memmap2::Mmap;
use obiskio::{SKFileReader, SKResult};

/// Extract all canonical kmers from a dereplicated superkmer file,
/// sort them with an external merge sort, deduplicate, and write
/// unique u64 values to `unique_path`.  Returns f0 (distinct kmers).
pub fn sort_unique_kmers(
    dedup_path: &Path,
    work_dir: &Path,
    unique_path: &Path,
    chunk_kmers: usize,
) -> SKResult<usize> {
    let chunk_paths = extract_and_sort_chunks(dedup_path, work_dir, chunk_kmers)?;

    if chunk_paths.is_empty() {
        fs::write(unique_path, [])?;
        return Ok(0);
    }

    let f0 = merge_and_dedup(&chunk_paths, unique_path)?;
    for p in &chunk_paths {
        fs::remove_file(p)?;
    }
    Ok(f0)
}

/// Number of kmers per sort chunk from available RAM (per thread).
pub fn chunk_size_from_ram(available_bytes: u64) -> usize {
    // 40% of available RAM; each kmer is 8 bytes (u64)
    let n = ((available_bytes as f64 * 0.40) / 8.0) as usize;
    n.max(1 << 20) // minimum 1M kmers (~8 MB)
}

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

fn extract_and_sort_chunks(
    dedup_path: &Path,
    work_dir: &Path,
    chunk_kmers: usize,
) -> SKResult<Vec<PathBuf>> {
    let mut reader = SKFileReader::open(dedup_path)?;
    let mut buf: Vec<u64> = Vec::with_capacity(chunk_kmers);
    let mut paths: Vec<PathBuf> = Vec::new();

    while let Some(sk) = reader.read()? {
        for kmer in sk.iter_canonical_kmers() {
            buf.push(kmer.raw());
            if buf.len() >= chunk_kmers {
                paths.push(flush_sorted_chunk(&mut buf, work_dir, paths.len())?);
            }
        }
    }
    if !buf.is_empty() {
        paths.push(flush_sorted_chunk(&mut buf, work_dir, paths.len())?);
    }
    Ok(paths)
}

fn flush_sorted_chunk(buf: &mut Vec<u64>, work_dir: &Path, idx: usize) -> io::Result<PathBuf> {
    buf.sort_unstable();
    let path = work_dir.join(format!("kmer_sort_{idx:05}.bin"));
    let mut w = BufWriter::new(fs::File::create(&path)?);
    for &v in buf.iter() {
        w.write_all(&v.to_le_bytes())?;
    }
    buf.clear();
    Ok(path)
}

/// K-way merge of sorted chunk files + in-line dedup → `unique_path`.
fn merge_and_dedup(chunk_paths: &[PathBuf], unique_path: &Path) -> io::Result<usize> {
    struct ChunkReader {
        mmap: Mmap,
        pos:  usize,
    }
    impl ChunkReader {
        fn open(path: &Path) -> io::Result<Self> {
            let f = fs::File::open(path)?;
            Ok(Self { mmap: unsafe { Mmap::map(&f)? }, pos: 0 })
        }
        fn peek(&self) -> Option<u64> {
            let off = self.pos * 8;
            if off + 8 <= self.mmap.len() {
                Some(u64::from_le_bytes(self.mmap[off..off + 8].try_into().unwrap()))
            } else {
                None
            }
        }
        fn advance(&mut self) { self.pos += 1; }
    }

    let mut readers: Vec<ChunkReader> = chunk_paths.iter()
        .map(|p| ChunkReader::open(p))
        .collect::<io::Result<_>>()?;

    let mut heap: BinaryHeap<Reverse<(u64, usize)>> = BinaryHeap::new();
    for (i, r) in readers.iter().enumerate() {
        if let Some(v) = r.peek() {
            heap.push(Reverse((v, i)));
        }
    }

    let mut w   = BufWriter::new(fs::File::create(unique_path)?);
    let mut f0  = 0usize;
    let mut prev: Option<u64> = None;

    while let Some(Reverse((val, idx))) = heap.pop() {
        if prev != Some(val) {
            w.write_all(&val.to_le_bytes())?;
            prev = Some(val);
            f0 += 1;
        }
        readers[idx].advance();
        if let Some(next_val) = readers[idx].peek() {
            heap.push(Reverse((next_val, idx)));
        }
    }

    Ok(f0)
}
