use std::collections::HashMap;
use std::fs::{self, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};

use memmap2::MmapMut;

use crate::format::{HEADER_SIZE, OVERFLOW_ENTRY_SIZE, finalize_pciv};
use crate::reader::PersistentCompactIntVec;

pub struct PersistentCompactIntVecBuilder {
    path: PathBuf,
    mmap: MmapMut,
    n: usize,
    overflow: HashMap<usize, u32>,
}

impl PersistentCompactIntVecBuilder {
    /// Create a new, zero-filled PCIV at `path`. Primary is mmapped immediately.
    pub fn new(n: usize, path: &Path) -> io::Result<Self> {
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)?;
        file.set_len((HEADER_SIZE + n) as u64)?;
        let mmap = unsafe { MmapMut::map_mut(&file)? };
        Ok(Self {
            path: path.to_path_buf(),
            mmap,
            n,
            overflow: HashMap::new(),
        })
    }

    /// Copy `source`'s file to `path`, mmap the primary section, load overflow into RAM.
    /// Avoids iterating all n slots: the file copy is OS-level, overflow loading is O(n_overflow).
    pub fn build_from(source: &PersistentCompactIntVec, 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();
        let n_overflow = u64::from_le_bytes(mmap[16..24].try_into().unwrap()) as usize;
        let data_offset = HEADER_SIZE + n;

        let mut overflow = HashMap::with_capacity(n_overflow);
        for i in 0..n_overflow {
            let off = data_offset + i * OVERFLOW_ENTRY_SIZE;
            let slot  = u64::from_le_bytes(mmap[off..off + 8].try_into().unwrap()) as usize;
            let value = u32::from_le_bytes(mmap[off + 8..off + 12].try_into().unwrap());
            overflow.insert(slot, value);
        }

        Ok(Self {
            path: path.to_path_buf(),
            mmap,
            n,
            overflow,
        })
    }

    /// Get the value at the given slot, handling overflow if necessary.
    pub fn get(&self, slot: usize) -> u32 {
        match self.mmap[HEADER_SIZE + slot] {
            255 => *self
                .overflow
                .get(&slot)
                .expect("sentinel without overflow entry"),
            v => v as u32,
        }
    }

    pub fn set(&mut self, slot: usize, value: u32) {
        if value < 255 {
            self.mmap[HEADER_SIZE + slot] = value as u8;
            self.overflow.remove(&slot);
        } else {
            self.mmap[HEADER_SIZE + slot] = 255;
            self.overflow.insert(slot, value);
        }
    }

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

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

    pub fn min(&mut self, other: &PersistentCompactIntVec) {
        assert_eq!(self.n, other.len(), "length mismatch");
        for (slot, other_val) in other.iter().enumerate() {
            if other_val < self.get(slot) {
                self.set(slot, other_val);
            }
        }
    }

    pub fn max(&mut self, other: &PersistentCompactIntVec) {
        assert_eq!(self.n, other.len(), "length mismatch");
        for (slot, other_val) in other.iter().enumerate() {
            if other_val > self.get(slot) {
                self.set(slot, other_val);
            }
        }
    }

    pub fn add(&mut self, other: &PersistentCompactIntVec) {
        assert_eq!(self.n, other.len(), "length mismatch");
        for (slot, other_val) in other.iter().enumerate() {
            let cur = self.get(slot);
            self.set(slot, cur.checked_add(other_val).expect("u32 overflow in add"));
        }
    }

    pub fn diff(&mut self, other: &PersistentCompactIntVec) {
        assert_eq!(self.n, other.len(), "length mismatch");
        for (slot, other_val) in other.iter().enumerate() {
            self.set(slot, self.get(slot).saturating_sub(other_val));
        }
    }

    /// Flush the primary mmap, then write sorted overflow data + index and fix the header.
    pub fn close(self) -> io::Result<()> {
        self.mmap.flush()?;
        let Self {
            path,
            mmap,
            n,
            overflow,
        } = self;
        drop(mmap);

        let mut entries: Vec<(usize, u32)> = overflow.into_iter().collect();
        entries.sort_unstable_by_key(|&(slot, _)| slot);

        finalize_pciv(&path, n, &entries)
    }
}
