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

use memmap2::Mmap;

use crate::format::{byte_count_nonzero, byte_sum, HEADER_SIZE, MAGIC, OVERFLOW_ENTRY_SIZE, parse_index_entry};
use crate::views::IntSliceView;

pub struct PersistentCompactIntVec {
    mmap: Mmap,
    n: usize,
    n_overflow: usize,
    pub step: usize,
    index: Vec<(usize, usize)>, // (slot, pos) — L1-resident sparse index
    primary_offset: usize,      // = HEADER_SIZE
    data_offset: usize,         // = HEADER_SIZE + n
    path: PathBuf,
}

impl PersistentCompactIntVec {
    pub fn open(path: &Path) -> io::Result<Self> {
        let mmap = unsafe { Mmap::map(&File::open(path)?)? };

        if mmap.len() < HEADER_SIZE {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "PCIV file too short"));
        }
        if &mmap[0..4] != &MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PCIV magic"));
        }

        let n          = u64::from_le_bytes(mmap[8..16].try_into().unwrap())  as usize;
        let n_overflow = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;
        let n_index    = u64::from_le_bytes(mmap[24..32].try_into().unwrap()) as usize;
        let step       = u64::from_le_bytes(mmap[32..40].try_into().unwrap()) as usize;

        let primary_offset = HEADER_SIZE;
        let data_offset    = primary_offset + n;
        let index_offset   = data_offset + n_overflow * OVERFLOW_ENTRY_SIZE;

        let mut index = Vec::with_capacity(n_index);
        for i in 0..n_index {
            index.push(parse_index_entry(&mmap, index_offset, i));
        }

        Ok(Self { mmap, n, n_overflow, step, index, primary_offset, data_offset, path: path.to_path_buf() })
    }

    pub fn path(&self) -> &Path { &self.path }
    pub fn len(&self)      -> usize { self.n }
    pub fn is_empty(&self) -> bool  { self.n == 0 }

    pub fn get(&self, slot: usize) -> u32 {
        match self.mmap[self.primary_offset + slot] {
            255 => self.overflow_get(slot),
            v   => v as u32,
        }
    }

    fn overflow_get(&self, slot: usize) -> u32 {
        let (pos_start, pos_end) = if self.step == 0 {
            (0, self.n_overflow)
        } else {
            let i = self.index.partition_point(|&(s, _)| s <= slot).saturating_sub(1);
            let start = self.index[i].1;
            let end = if i + 1 < self.index.len() { self.index[i + 1].1 } else { self.n_overflow };
            (start, end)
        };
        let mut lo = pos_start;
        let mut hi = pos_end;
        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            match self.data_slot(mid).cmp(&slot) {
                std::cmp::Ordering::Equal   => return self.data_value(mid),
                std::cmp::Ordering::Less    => lo = mid + 1,
                std::cmp::Ordering::Greater => hi = mid,
            }
        }
        panic!("slot {slot} marked as overflow but not found")
    }

    #[inline]
    fn data_slot(&self, i: usize) -> usize {
        let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE;
        u64::from_le_bytes(self.mmap[off..off + 8].try_into().unwrap()) as usize
    }

    #[inline]
    fn data_value(&self, i: usize) -> u32 {
        let off = self.data_offset + i * OVERFLOW_ENTRY_SIZE + 8;
        u32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap())
    }

    pub fn sum(&self) -> u64 {
        let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
        byte_sum(primary, (0..self.n_overflow).map(|i| self.data_value(i)))
    }

    pub fn count_nonzero(&self) -> u64 {
        let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
        byte_count_nonzero(primary)
    }

    /// Lightweight zero-copy view — primary and overflow point into the mmap.
    pub fn view(&self) -> IntSliceView<'_> {
        let primary = &self.mmap[self.primary_offset..self.primary_offset + self.n];
        let overflow_raw = &self.mmap[self.data_offset..self.data_offset + self.n_overflow * OVERFLOW_ENTRY_SIZE];
        IntSliceView::new(primary, overflow_raw, self.n_overflow, self.n)
    }

    pub fn iter(&self) -> Iter<'_> {
        Iter { pciv: self, slot: 0, overflow_pos: 0 }
    }

    // ── Distance methods ──────────────────────────────────────────────────────

    pub fn bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        let sum_min = self.partial_bray_dist(other);
        let denom = self.sum() + other.sum();
        if denom == 0 { 0.0 } else { 1.0 - 2.0 * sum_min as f64 / denom as f64 }
    }

    pub fn partial_bray_dist(&self, other: &PersistentCompactIntVec) -> u64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter()).map(|(a, b)| a.min(b) as u64).sum()
    }

    pub fn relfreq_bray_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        let sa = self.sum() as f64;
        let sb = other.sum() as f64;
        if sa == 0.0 && sb == 0.0 { return 0.0; }
        1.0 - self.partial_relfreq_bray_dist(other, sa, sb)
    }

    pub fn partial_relfreq_bray_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
                let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
                pa.min(pb)
            })
            .sum()
    }

    pub fn euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.partial_euclidean_dist(other).sqrt()
    }

    pub fn partial_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter())
            .map(|(a, b)| { let d = a as f64 - b as f64; d * d })
            .sum()
    }

    pub fn relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        let sa = self.sum() as f64;
        let sb = other.sum() as f64;
        if sa == 0.0 && sb == 0.0 { return 0.0; }
        self.partial_relfreq_euclidean_dist(other, sa, sb).sqrt()
    }

    pub fn partial_relfreq_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 { a as f64 / sum_a } else { 0.0 };
                let pb = if sum_b > 0.0 { b as f64 / sum_b } else { 0.0 };
                let d = pa - pb;
                d * d
            })
            .sum()
    }

    pub fn hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        let sa = self.sum() as f64;
        let sb = other.sum() as f64;
        if sa == 0.0 && sb == 0.0 { return 0.0; }
        self.partial_hellinger_euclidean_dist(other, sa, sb).sqrt()
    }

    pub fn partial_hellinger_euclidean_dist(&self, other: &PersistentCompactIntVec, sum_a: f64, sum_b: f64) -> f64 {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter())
            .map(|(a, b)| {
                let pa = if sum_a > 0.0 { (a as f64 / sum_a).sqrt() } else { 0.0 };
                let pb = if sum_b > 0.0 { (b as f64 / sum_b).sqrt() } else { 0.0 };
                let d = pa - pb;
                d * d
            })
            .sum()
    }

    pub fn hellinger_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.hellinger_euclidean_dist(other) / std::f64::consts::SQRT_2
    }

    pub fn threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> f64 {
        let (intersection, union) = self.partial_threshold_jaccard_dist(other, threshold);
        if union == 0 { 0.0 } else { 1.0 - intersection as f64 / union as f64 }
    }

    pub fn partial_threshold_jaccard_dist(&self, other: &PersistentCompactIntVec, threshold: u32) -> (u64, u64) {
        assert_eq!(self.n, other.len(), "length mismatch");
        self.iter().zip(other.iter())
            .fold((0u64, 0u64), |(inter, uni), (a, b)| {
                let ap = a >= threshold;
                let bp = b >= threshold;
                (inter + (ap & bp) as u64, uni + (ap | bp) as u64)
            })
    }

    pub fn jaccard_dist(&self, other: &PersistentCompactIntVec) -> f64 {
        self.threshold_jaccard_dist(other, 1)
    }
}

impl<'a> IntoIterator for &'a PersistentCompactIntVec {
    type Item = u32;
    type IntoIter = Iter<'a>;
    fn into_iter(self) -> Iter<'a> { self.iter() }
}

pub struct Iter<'a> {
    pciv: &'a PersistentCompactIntVec,
    slot: usize,
    overflow_pos: usize,
}

impl ExactSizeIterator for Iter<'_> {}

impl Iterator for Iter<'_> {
    type Item = u32;

    fn next(&mut self) -> Option<u32> {
        if self.slot >= self.pciv.n { return None; }
        let v = self.pciv.mmap[self.pciv.primary_offset + self.slot];
        self.slot += 1;
        if v < 255 {
            Some(v as u32)
        } else {
            let val = self.pciv.data_value(self.overflow_pos);
            self.overflow_pos += 1;
            Some(val)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.pciv.n - self.slot;
        (remaining, Some(remaining))
    }
}
