use std::fs;
use std::io;

use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::{PersistentCompactIntMatrix, PersistentCompactIntVec};
use obidebruinj::GraphDeBruijn;
use obilayeredmap::meta::PartitionMeta;
use obilayeredmap::{IndexMode, layer::Layer};
use obiskio::{SKError, SKFileMeta, SKFileReader};
use ptr_hash::{PtrHash, bucket_fn::CubicEps, hash::Xx64};

use crate::common::olm_to_sk;
use crate::graph_pipeline::{materialize_layer, write_graph_as_unitigs};
use crate::partition::KmerPartition;

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

fn remove_if_exists(path: &std::path::Path) {
    if let Err(e) = fs::remove_file(path) {
        if e.kind() != io::ErrorKind::NotFound {
            eprintln!("warning: could not remove {}: {e}", path.display());
        }
    }
}

impl KmerPartition {
    /// Build the layered MPHF index for partition `i`.
    ///
    /// Returns the number of canonical k-mers indexed, or 0 if the partition
    /// has no data or its layer was already built (resume-safe).
    ///
    /// Abundance filtering is applied when `min_ab > 1` or `max_ab.is_some()`,
    /// using `mphf1.bin` + `counts1.bin` if they exist.
    /// Count payload is stored iff `with_counts` is true.
    pub fn build_index_layer(
        &self,
        i: usize,
        min_ab: u32,
        max_ab: Option<u32>,
        with_counts: bool,
        mode: &IndexMode,
        block_bits: u8,
    ) -> Result<usize, SKError> {
        let part_dir = self.part_dir(i);
        let dedup_path = part_dir.join("dereplicated.skmer.zst");
        if !dedup_path.exists() {
            return Ok(0);
        }

        let layer_dir = part_dir.join("index").join("layer_0");
        if layer_dir.join("mphf.bin").exists() {
            return Ok(0);
        }

        let filter_active = min_ab > 1 || max_ab.is_some();
        let need_counts = filter_active || with_counts;

        let mphf1_opt: Option<Mphf> = if need_counts {
            let p = part_dir.join("mphf1.bin");
            p.exists().then(|| Mphf::load_full(&p).ok()).flatten()
        } else {
            None
        };

        let counts1_opt: Option<PersistentCompactIntVec> = if need_counts {
            let p = part_dir.join("counts1.bin");
            p.exists()
                .then(|| PersistentCompactIntVec::open(&p).ok())
                .flatten()
        } else {
            None
        };

        let mut g = GraphDeBruijn::new();
        let mut reader = SKFileReader::open(&dedup_path)?;
        for sk in reader.iter() {
            for kmer in sk.iter_canonical_kmers() {
                let accept = if filter_active {
                    match (&mphf1_opt, &counts1_opt) {
                        (Some(mphf), Some(counts)) => {
                            let ab = counts.get(mphf.index(&kmer.raw()));
                            ab >= min_ab && max_ab.map_or(true, |max| ab <= max)
                        }
                        _ => true,
                    }
                } else {
                    true
                };
                if accept {
                    g.push(kmer);
                }
            }
        }

        let n_kmers = if with_counts {
            let n = write_graph_as_unitigs(g, &layer_dir)?;
            Layer::<PersistentCompactIntMatrix>::build(
                &layer_dir,
                block_bits,
                mode,
                |kmer| match (&mphf1_opt, &counts1_opt) {
                    (Some(mphf), Some(counts)) => counts.get(mphf.index(&kmer.raw())),
                    _ => 1,
                },
            )
            .map_err(|e| olm_to_sk(e, "layer build"))?;
            n
        } else {
            materialize_layer(g, &layer_dir, block_bits, mode)?
        };

        let index_dir = layer_dir.parent().expect("layer_dir has a parent");
        PartitionMeta {
            n_layers: 1,
            mode: mode.clone(),
        }
        .save(index_dir)
        .map_err(|e| olm_to_sk(e, "layer build"))?;

        Ok(n_kmers)
    }

    /// Remove intermediate build artifacts for partition `i`.
    ///
    /// Deletes `dereplicated.skmer.zst` (+ sidecar), `mphf1.bin`, `counts1.bin`.
    pub fn remove_build_artifacts(&self, i: usize) {
        let part_dir = self.part_dir(i);
        let dedup = part_dir.join("dereplicated.skmer.zst");
        remove_if_exists(&SKFileMeta::sidecar_path(&dedup));
        remove_if_exists(&dedup);
        remove_if_exists(&part_dir.join("mphf1.bin"));
        remove_if_exists(&part_dir.join("counts1.bin"));
    }
}
