use std::fs::{self, File, OpenOptions};
use std::io::{self, Seek, SeekFrom, Write as _};
use std::path::{Path, PathBuf};

use memmap2::{Mmap, MmapMut};

use crate::reader::PersistentCompactIntVec;
use crate::views::{BitSliceIter, BitSliceView, IntSliceView};

const MAGIC: [u8; 4] = *b"PBIV";

// Header: magic(4) + _pad(4) + n(8) = 16 bytes.
// Data starts at offset 16, u64-aligned (mmap base is page-aligned, 16 % 8 == 0).
const HEADER_SIZE: usize = 16;

#[inline]
pub(crate) fn n_words(n: usize) -> usize { n.div_ceil(64) }

#[inline]
fn n_bytes_for_words(n: usize) -> usize { n_words(n) * 8 }

// ── PersistentBitVec ──────────────────────────────────────────────────────────

pub struct PersistentBitVec {
    mmap: Mmap,
    n:    usize,
    path: PathBuf,
}

impl PersistentBitVec {
    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, "PBIV file too short"));
        }
        if &mmap[0..4] != &MAGIC {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PBIV magic"));
        }
        let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
        Ok(Self { mmap, n, 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) -> bool {
        (self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
    }

    // SAFETY: mmap is page-aligned, HEADER_SIZE=16 divisible by 8 → u64-aligned.
    fn data_words(&self) -> &[u64] {
        let nw  = n_words(self.n);
        let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
        unsafe { std::slice::from_raw_parts(ptr, nw) }
    }

    pub fn view(&self) -> BitSliceView<'_> {
        BitSliceView::new(self.data_words(), self.n)
    }

    pub fn words(&self) -> &[u64] { self.data_words() }

    pub fn count_ones(&self)  -> u64 { self.view().count_ones() }
    pub fn count_zeros(&self) -> u64 { self.view().count_zeros() }

    pub fn partial_jaccard_dist(&self, other: &PersistentBitVec) -> (u64, u64) {
        self.view().partial_jaccard_dist(other.view())
    }
    pub fn jaccard_dist(&self, other: &PersistentBitVec) -> f64 {
        self.view().jaccard_dist(other.view())
    }
    pub fn hamming_dist(&self, other: &PersistentBitVec) -> u64 {
        self.view().hamming_dist(other.view())
    }

    pub fn iter(&self) -> BitIter<'_> {
        BitIter { words: self.data_words(), slot: 0, n: self.n }
    }
}

impl<'a> IntoIterator for &'a PersistentBitVec {
    type Item = bool;
    type IntoIter = BitIter<'a>;
    fn into_iter(self) -> BitIter<'a> { self.iter() }
}

// ── BitIter ───────────────────────────────────────────────────────────────────

pub struct BitIter<'a> {
    words: &'a [u64],
    slot:  usize,
    n:     usize,
}

impl ExactSizeIterator for BitIter<'_> {}

impl Iterator for BitIter<'_> {
    type Item = bool;
    fn next(&mut self) -> Option<bool> {
        if self.slot >= self.n { return None; }
        let v = (self.words[self.slot >> 6] >> (self.slot & 63)) & 1 != 0;
        self.slot += 1;
        Some(v)
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let rem = self.n - self.slot;
        (rem, Some(rem))
    }
}

// ── PersistentBitVecBuilder ───────────────────────────────────────────────────

pub struct PersistentBitVecBuilder {
    mmap: MmapMut,
    n:    usize,
    path: PathBuf,
}

impl PersistentBitVecBuilder {
    pub fn new(n: usize, path: &Path) -> io::Result<Self> {
        let file_size = HEADER_SIZE + n_bytes_for_words(n);
        let mut file = OpenOptions::new()
            .read(true).write(true).create(true).truncate(true)
            .open(path)?;
        file.write_all(&MAGIC)?;
        file.write_all(&[0u8; 4])?;
        file.write_all(&(n as u64).to_le_bytes())?;
        file.seek(SeekFrom::Start(0))?;
        file.set_len(file_size as u64)?;
        let mmap = unsafe { MmapMut::map_mut(&file)? };
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn from_raw_bytes(bytes: &[u8], n: usize, path: &Path) -> io::Result<Self> {
        let file_size = HEADER_SIZE + n_bytes_for_words(n);
        let file = OpenOptions::new()
            .read(true).write(true).create(true).truncate(true)
            .open(path)?;
        file.set_len(file_size as u64)?;
        let mut mmap = unsafe { MmapMut::map_mut(&file)? };
        mmap[0..4].copy_from_slice(&MAGIC);
        mmap[8..16].copy_from_slice(&(n as u64).to_le_bytes());
        mmap[HEADER_SIZE..HEADER_SIZE + bytes.len()].copy_from_slice(bytes);
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    /// Create an all-ones bit vector of length `n` at `path`.
    ///
    /// More efficient than `new(n, path)` + `not()`: the data is written as
    /// 0xFF bytes in a single sequential pass, with no intermediate all-zeros state.
    pub fn new_ones(n: usize, path: &Path) -> io::Result<Self> {
        let nw        = n_words(n);
        let file_size = HEADER_SIZE + nw * 8;
        let mut file  = OpenOptions::new()
            .read(true).write(true).create(true).truncate(true)
            .open(path)?;
        file.write_all(&MAGIC)?;
        file.write_all(&[0u8; 4])?;
        file.write_all(&(n as u64).to_le_bytes())?;
        file.write_all(&vec![0xFFu8; nw * 8])?;
        file.seek(SeekFrom::Start(0))?;
        file.set_len(file_size as u64)?;
        let mut mmap = unsafe { MmapMut::map_mut(&file)? };
        // Clear padding bits in the last word so trailing bits are always 0.
        let rem = n % 64;
        if rem != 0 {
            let ptr   = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
            let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
            words[nw - 1] &= (1u64 << rem) - 1;
        }
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn build_from(source: &PersistentBitVec, path: &Path) -> io::Result<Self> {
        fs::copy(source.path(), path)?;
        let file = OpenOptions::new().read(true).write(true).open(path)?;
        let mmap = unsafe { MmapMut::map_mut(&file)? };
        let n = source.len();
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn build_from_counts(source: &PersistentCompactIntVec, threshold: u32, path: &Path) -> io::Result<Self> {
        let n = source.len();
        let file_size = HEADER_SIZE + n_bytes_for_words(n);
        let mut file = OpenOptions::new()
            .read(true).write(true).create(true).truncate(true)
            .open(path)?;
        file.write_all(&MAGIC)?;
        file.write_all(&[0u8; 4])?;
        file.write_all(&(n as u64).to_le_bytes())?;
        file.seek(SeekFrom::Start(0))?;
        file.set_len(file_size as u64)?;
        let mut mmap = unsafe { MmapMut::map_mut(&file)? };
        {
            let nw  = n_words(n);
            let ptr = mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
            let words = unsafe { std::slice::from_raw_parts_mut(ptr, nw) };
            for (slot, count) in source.iter().enumerate() {
                if count >= threshold { words[slot >> 6] |= 1u64 << (slot & 63); }
            }
        }
        Ok(Self { mmap, n, path: path.to_path_buf() })
    }

    pub fn build_from_presence(source: &PersistentCompactIntVec, path: &Path) -> io::Result<Self> {
        Self::build_from_counts(source, 1, path)
    }

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

    pub fn get(&self, slot: usize) -> bool {
        (self.mmap[HEADER_SIZE + (slot >> 3)] >> (slot & 7)) & 1 != 0
    }

    pub fn set(&mut self, slot: usize, value: bool) {
        let bit = 1u64 << (slot & 63);
        if value { self.data_words_mut()[slot >> 6] |=  bit; }
        else     { self.data_words_mut()[slot >> 6] &= !bit; }
    }

    fn data_words(&self) -> &[u64] {
        let nw  = n_words(self.n);
        let ptr = self.mmap[HEADER_SIZE..].as_ptr() as *const u64;
        unsafe { std::slice::from_raw_parts(ptr, nw) }
    }

    // SAFETY: same alignment argument as PersistentBitVec::data_words.
    fn data_words_mut(&mut self) -> &mut [u64] {
        let nw  = n_words(self.n);
        let ptr = self.mmap[HEADER_SIZE..].as_mut_ptr() as *mut u64;
        unsafe { std::slice::from_raw_parts_mut(ptr, nw) }
    }

    pub fn view(&self) -> BitSliceView<'_> {
        BitSliceView::new(self.data_words(), self.n)
    }

    pub fn words(&self) -> &[u64] { self.data_words() }

    pub fn copy_from(&mut self, src: BitSliceView<'_>) {
        assert_eq!(self.n, src.len(), "BitSliceView length mismatch");
        self.data_words_mut().copy_from_slice(src.words());
    }

    pub fn and(&mut self, other: BitSliceView<'_>) {
        assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
        for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w &= o; }
    }

    pub fn or(&mut self, other: BitSliceView<'_>) {
        assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
        for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w |= o; }
    }

    pub fn xor(&mut self, other: BitSliceView<'_>) {
        assert_eq!(self.n, other.len(), "BitSliceView length mismatch");
        for (w, &o) in self.data_words_mut().iter_mut().zip(other.words()) { *w ^= o; }
    }

    pub fn not(&mut self) {
        let rem   = self.n % 64;
        let words = self.data_words_mut();
        for w in words.iter_mut() { *w ^= u64::MAX; }
        if rem != 0 {
            if let Some(last) = words.last_mut() { *last &= (1u64 << rem) - 1; }
        }
    }

    /// OR in bits at slots where `pred(col[slot])` is true.
    pub fn or_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
        assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
        let n = self.n;
        let primary = col.primary_bytes();
        let words = self.data_words_mut();
        let nw = n_words(n);
        for wi in 0..nw {
            let base  = wi * 64;
            let limit = (base + 64).min(n);
            let mut mask = 0u64;
            for bit in 0..(limit - base) {
                let b = primary[base + bit];
                if b < 255 && pred(b as u32) { mask |= 1u64 << bit; }
            }
            words[wi] |= mask;
        }
        for (slot, val) in col.overflow_entries() {
            if pred(val) { words[slot >> 6] |= 1u64 << (slot & 63); }
        }
    }

    /// Clear bits at slots where `pred(col[slot])` is false.
    pub fn and_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
        assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
        let n = self.n;
        let primary = col.primary_bytes();
        let words = self.data_words_mut();
        let nw = n_words(n);
        for wi in 0..nw {
            let base  = wi * 64;
            let limit = (base + 64).min(n);
            let mut mask = 0u64;
            for bit in 0..(limit - base) {
                let b = primary[base + bit];
                if b < 255 && !pred(b as u32) { mask |= 1u64 << bit; }
            }
            words[wi] &= !mask;
        }
        for (slot, val) in col.overflow_entries() {
            if !pred(val) { words[slot >> 6] &= !(1u64 << (slot & 63)); }
        }
    }

    /// Toggle bits at slots where `pred(col[slot])` is true.
    pub fn xor_where(&mut self, col: IntSliceView<'_>, pred: impl Fn(u32) -> bool) {
        assert_eq!(self.n, col.len(), "IntSliceView length mismatch");
        let n = self.n;
        let primary = col.primary_bytes();
        let words = self.data_words_mut();
        let nw = n_words(n);
        for wi in 0..nw {
            let base  = wi * 64;
            let limit = (base + 64).min(n);
            let mut mask = 0u64;
            for bit in 0..(limit - base) {
                let b = primary[base + bit];
                if b < 255 && pred(b as u32) { mask |= 1u64 << bit; }
            }
            words[wi] ^= mask;
        }
        for (slot, val) in col.overflow_entries() {
            if pred(val) { words[slot >> 6] ^= 1u64 << (slot & 63); }
        }
    }

    pub fn iter(&self) -> BitSliceIter<'_> {
        self.view().iter()
    }

    pub fn close(self) -> io::Result<()> { self.mmap.flush() }

    pub fn finish(self) -> io::Result<PersistentBitVec> {
        let path = self.path.clone();
        self.close()?;
        PersistentBitVec::open(&path)
    }
}
