use std::fs;
use std::path::{Path, PathBuf};

use cacheline_ef::{CachelineEf, CachelineEfVec};
use epserde::prelude::*;
use obicompactvec::LayerMeta;
use obikseq::CanonicalKmer;
use obiskio::{CanonicalKmerIter, UnitigFileReader, UnitigFileWriter, build_unitig_idx};
use ptr_hash::{PtrHash, PtrHashParams, bucket_fn::CubicEps, hash::Xx64};

use crate::error::{OLMError, OLMResult};
use crate::evidence::{Evidence, EvidenceWriter};
use crate::fingerprint::{FingerprintVec, FingerprintVecWriter};
use crate::meta::IndexMode;

pub(crate) const MPHF_FILE:        &str = "mphf.bin";
pub(crate) const UNITIGS_FILE:     &str = "unitigs.bin";
pub(crate) const EVIDENCE_FILE:    &str = "evidence.bin";
pub(crate) const FINGERPRINT_FILE: &str = "fingerprint.bin";

/// Owned MPHF — used only at build time (construction + store).
pub(crate) type Mphf = PtrHash<u64, CubicEps, CachelineEfVec<Vec<CachelineEf>>, Xx64, Vec<u8>>;

/// Zero-copy MPHF for querying — ε-deserialized view into a memory-mapped file.
/// `MemCase` owns the mmap backing; `'static` is sound because MemCase pins the memory.
type MphfEps = PtrHash<u64, CubicEps, CachelineEfVec<&'static [CachelineEf]>, Xx64, &'static [u8]>;

// ── LayerEvidence ─────────────────────────────────────────────────────────────

enum LayerEvidence {
    Exact  { evidence: Evidence, unitigs: UnitigFileReader },
    Approx { fingerprint: FingerprintVec, unitigs_path: PathBuf },
    Hybrid { evidence: Evidence, unitigs: UnitigFileReader, fingerprint: FingerprintVec },
}

// ── MphfLayer ─────────────────────────────────────────────────────────────────

/// Autonomous kmer → slot mapping for one layer.
///
/// Two query methods:
/// - [`find`](Self::find) — O(1), uses fingerprint (Approx/Hybrid) or exact evidence (Exact).
/// - [`find_strict`](Self::find_strict) — always exact; O(1) on Exact/Hybrid layers,
///   O(n) sequential scan on Approx layers.
pub struct MphfLayer {
    mphf: MemCase<MphfEps>,
    ev:   LayerEvidence,
    n:    usize,
}

impl MphfLayer {
    /// Open a layer using the index-level `mode` determined at `LayeredMap` open time.
    /// No per-layer metadata file is read.
    pub fn open(dir: &Path, mode: &IndexMode) -> OLMResult<Self> {
        let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
            .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
        let (ev, n) = match mode {
            IndexMode::Exact => {
                let evidence = Evidence::open(&dir.join(EVIDENCE_FILE))?;
                let n = evidence.len();
                // open() auto-detects: uses direct access since exact layers always have .idx
                let unitigs = UnitigFileReader::open(&dir.join(UNITIGS_FILE))?;
                (LayerEvidence::Exact { evidence, unitigs }, n)
            }
            IndexMode::Approx { .. } => {
                let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
                let n = fingerprint.n();
                let unitigs_path = dir.join(UNITIGS_FILE);
                (LayerEvidence::Approx { fingerprint, unitigs_path }, n)
            }
            IndexMode::Hybrid { .. } => {
                let evidence    = Evidence::open(&dir.join(EVIDENCE_FILE))?;
                let fingerprint = FingerprintVec::open(&dir.join(FINGERPRINT_FILE))?;
                let n = evidence.len();
                let unitigs = UnitigFileReader::open(&dir.join(UNITIGS_FILE))?;
                (LayerEvidence::Hybrid { evidence, unitigs, fingerprint }, n)
            }
        };
        Ok(Self { mphf, ev, n })
    }

    // ── Query API ─────────────────────────────────────────────────────────────

    /// O(1) lookup — dispatches automatically:
    /// - Exact:  evidence + `verify_canonical_kmer`, zero false positives.
    /// - Approx: fingerprint check, false-positive rate ≈ 1/2^b.
    /// - Hybrid: fingerprint check (fast path), zero false positives via `find_strict`.
    #[inline]
    pub fn find(&self, kmer: CanonicalKmer) -> Option<usize> {
        let slot = self.mphf.index(&kmer.raw());
        if slot >= self.n { return None; }
        match &self.ev {
            LayerEvidence::Exact { evidence, unitigs } => {
                let (chunk_id, rank) = evidence.decode(slot);
                if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
                    Some(slot)
                } else {
                    None
                }
            }
            LayerEvidence::Approx { fingerprint, .. } |
            LayerEvidence::Hybrid { fingerprint, .. } => {
                if fingerprint.matches(slot, kmer.seq_hash()) { Some(slot) } else { None }
            }
        }
    }

    /// Always-exact lookup — zero false positives regardless of mode.
    ///
    /// - Exact/Hybrid: O(1) via evidence + `verify_canonical_kmer`.
    /// - Approx: O(n) sequential scan of `unitigs.bin` to confirm the kmer
    ///   that owns the slot, then exact comparison.
    pub fn find_strict(&self, kmer: CanonicalKmer) -> Option<usize> {
        let slot = self.mphf.index(&kmer.raw());
        if slot >= self.n { return None; }
        match &self.ev {
            LayerEvidence::Exact  { evidence, unitigs } |
            LayerEvidence::Hybrid { evidence, unitigs, .. } => {
                let (chunk_id, rank) = evidence.decode(slot);
                if unitigs.verify_canonical_kmer(chunk_id as usize, rank as usize, kmer) {
                    Some(slot)
                } else {
                    None
                }
            }
            LayerEvidence::Approx { unitigs_path, .. } => {
                let reader = UnitigFileReader::open_sequential(unitigs_path).ok()?;
                for (stored, _, _) in reader.iter_indexed_canonical_kmers() {
                    if self.mphf.index(&stored.raw()) == slot {
                        return if stored == kmer { Some(slot) } else { None };
                    }
                }
                None
            }
        }
    }

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

// ── MphfOnly ──────────────────────────────────────────────────────────────────

/// Lightweight wrapper that loads only the MPHF file, without evidence or unitigs.
///
/// Use this when the caller guarantees that all queried kmers are in the MPHF
/// domain (e.g. when iterating the source's own unitigs during merge).
pub struct MphfOnly(MemCase<MphfEps>);

impl MphfOnly {
    pub fn open(dir: &Path) -> OLMResult<Self> {
        let mphf: MemCase<MphfEps> = Mphf::mmap(&dir.join(MPHF_FILE), Flags::empty())
            .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
        Ok(Self(mphf))
    }

    /// Return the slot for `kmer`. Only valid when `kmer` is in the MPHF domain.
    #[inline]
    pub fn index(&self, kmer: CanonicalKmer) -> usize {
        self.0.index(&kmer.raw())
    }
}

impl MphfLayer {
    // ── Build helpers ─────────────────────────────────────────────────────────

    pub fn unitig_writer(dir: &Path) -> OLMResult<UnitigFileWriter> {
        fs::create_dir_all(dir)?;
        Ok(UnitigFileWriter::create(&dir.join(UNITIGS_FILE))?)
    }

    /// Build `evidence.bin` + `unitigs.bin.idx` from `unitigs.bin` + `mphf.bin`.
    pub fn build_exact_evidence(dir: &Path, block_bits: u8) -> OLMResult<usize> {
        let unitig_path = dir.join(UNITIGS_FILE);
        let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
        let n = unitigs.n_kmers();

        if n == 0 {
            fs::File::create(dir.join(EVIDENCE_FILE))?;
            build_unitig_idx(&unitig_path, block_bits)?;
            return Ok(0);
        }

        let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
            .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;

        let mut ev   = EvidenceWriter::new(n);
        let mut seen = vec![0u8; (n + 7) / 8];

        for (kmer, chunk_id, rank) in unitigs.iter_indexed_canonical_kmers() {
            let slot = mphf.index(&kmer.raw());
            if slot >= n {
                return Err(OLMError::Mphf("slot out of bounds".into()));
            }
            let byte = slot / 8;
            let bit  = 1u8 << (slot % 8);
            if seen[byte] & bit != 0 {
                return Err(OLMError::Mphf("duplicate slot".into()));
            }
            seen[byte] |= bit;
            ev.set(slot, chunk_id as u32, rank as u8);
        }

        ev.write(&dir.join(EVIDENCE_FILE))?;
        build_unitig_idx(&unitig_path, block_bits)?;
        Ok(n)
    }

    /// Build `fingerprint.bin` from `unitigs.bin` + `mphf.bin`.
    pub fn build_approx_evidence(dir: &Path, b: u8, z: u8) -> OLMResult<usize> {
        if b == 0 || b > 64 {
            return Err(OLMError::InvalidLayer("fingerprint width must be 1..=64".into()));
        }
        if z == 0 {
            return Err(OLMError::InvalidLayer("z must be ≥ 1".into()));
        }
        let unitig_path = dir.join(UNITIGS_FILE);
        let unitigs = UnitigFileReader::open_sequential(&unitig_path)?;
        let n = unitigs.n_kmers();

        if n == 0 {
            FingerprintVecWriter::new(0, b).write(&dir.join(FINGERPRINT_FILE))?;
            return Ok(0);
        }

        let mphf: Mphf = Mphf::load_full(&dir.join(MPHF_FILE))
            .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;

        let mut fw = FingerprintVecWriter::new(n, b);

        for (kmer, _, _) in unitigs.iter_indexed_canonical_kmers() {
            let slot = mphf.index(&kmer.raw());
            if slot >= n {
                return Err(OLMError::Mphf("slot out of bounds".into()));
            }
            fw.set(slot, kmer.seq_hash());
        }

        fw.write(&dir.join(FINGERPRINT_FILE))?;
        Ok(n)
    }

    /// Build MPHF + evidence from `unitigs.bin` already present in `dir`.
    ///
    /// `fill_slot(slot, kmer)` is called once per kmer in all modes.
    /// No `layer_meta.json` is written — the mode is an index-level property
    /// stored in `PartitionMeta`.
    pub(crate) fn build(
        dir: &Path,
        block_bits: u8,
        mode: &IndexMode,
        fill_slot: &mut impl FnMut(usize, CanonicalKmer) -> OLMResult<()>,
    ) -> OLMResult<usize> {
        use rayon::prelude::*;

        let unitig_path = dir.join(UNITIGS_FILE);
        let n = UnitigFileReader::open_sequential(&unitig_path)?.n_kmers();

        let sk_to_olm = |e: obiskio::SKError| match e {
            obiskio::SKError::Io(io) => OLMError::Io(io),
            e => OLMError::InvalidLayer(e.to_string()),
        };

        // ── Empty layer ───────────────────────────────────────────────────────
        if n == 0 {
            let mphf: Mphf =
                Mphf::try_new(&[] as &[u64], PtrHashParams::<CubicEps>::default())
                    .ok_or_else(|| OLMError::Mphf("construction failed".into()))?;
            mphf.store(&dir.join(MPHF_FILE))
                .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;
            match mode {
                IndexMode::Exact | IndexMode::Hybrid { .. } => {
                    fs::File::create(dir.join(EVIDENCE_FILE))?;
                    build_unitig_idx(&unitig_path, block_bits)?;
                }
                IndexMode::Approx { b, .. } => {
                    FingerprintVecWriter::new(0, *b).write(&dir.join(FINGERPRINT_FILE))?;
                }
            }
            if let IndexMode::Hybrid { b, .. } = mode {
                FingerprintVecWriter::new(0, *b).write(&dir.join(FINGERPRINT_FILE))?;
            }
            return Ok(0);
        }

        // ── Pass 1: MPHF via clonable mmap iterator ───────────────────────────
        let keys = CanonicalKmerIter::new(&unitig_path).map_err(sk_to_olm)?;
        let mphf: Mphf =
            Mphf::new_from_par_iter(n, keys.map(|k| k.raw()).par_bridge(),
                                    PtrHashParams::<CubicEps>::default());
        mphf.store(&dir.join(MPHF_FILE))
            .map_err(|e| OLMError::InvalidLayer(e.to_string()))?;

        // ── Pass 2: fill evidence files + callback ────────────────────────────
        let unitigs2 = UnitigFileReader::open_sequential(&unitig_path)?;
        let mut seen = vec![0u8; (n + 7) / 8];

        match mode {
            IndexMode::Exact => {
                let mut ev = EvidenceWriter::new(n);
                for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
                    let slot = mphf.index(&kmer.raw());
                    if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
                    let byte = slot / 8; let bit = 1u8 << (slot % 8);
                    if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
                    seen[byte] |= bit;
                    ev.set(slot, chunk_id as u32, rank as u8);
                    fill_slot(slot, kmer)?;
                }
                ev.write(&dir.join(EVIDENCE_FILE))?;
                build_unitig_idx(&unitig_path, block_bits)?;
            }

            IndexMode::Approx { b, .. } => {
                let mut fw = FingerprintVecWriter::new(n, *b);
                for (kmer, _, _) in unitigs2.iter_indexed_canonical_kmers() {
                    let slot = mphf.index(&kmer.raw());
                    if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
                    let byte = slot / 8; let bit = 1u8 << (slot % 8);
                    if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
                    seen[byte] |= bit;
                    fw.set(slot, kmer.seq_hash());
                    fill_slot(slot, kmer)?;
                }
                fw.write(&dir.join(FINGERPRINT_FILE))?;
            }

            IndexMode::Hybrid { b, .. } => {
                let mut ev = EvidenceWriter::new(n);
                let mut fw = FingerprintVecWriter::new(n, *b);
                for (kmer, chunk_id, rank) in unitigs2.iter_indexed_canonical_kmers() {
                    let slot = mphf.index(&kmer.raw());
                    if slot >= n { return Err(OLMError::Mphf("slot out of bounds".into())); }
                    let byte = slot / 8; let bit = 1u8 << (slot % 8);
                    if seen[byte] & bit != 0 { return Err(OLMError::Mphf("duplicate slot".into())); }
                    seen[byte] |= bit;
                    ev.set(slot, chunk_id as u32, rank as u8);
                    fw.set(slot, kmer.seq_hash());
                    fill_slot(slot, kmer)?;
                }
                ev.write(&dir.join(EVIDENCE_FILE))?;
                fw.write(&dir.join(FINGERPRINT_FILE))?;
                build_unitig_idx(&unitig_path, block_bits)?;
            }
        }

        LayerMeta::save(dir, n)?;
        Ok(n)
    }
}
