use bitvec::prelude::*;

#[derive(Clone)]
pub struct PaletteStorage<T: Copy + Eq> {
    len: usize,
    bits: usize,
    data: BitVec,
    entries: Vec<PaletteEntry<T>>,
}

impl<T: Copy + Eq> PaletteStorage<T> {
    pub fn new(len: usize, fill: T) -> Self {
        Self {
            len,
            bits: 0,
            data: bitvec![],
            entries: vec![PaletteEntry {
                value: Some(fill),
                used: len,
            }],
        }
    }

    pub fn new_with_default(len: usize) -> Self
    where
        T: Default,
    {
        Self::new(len, Default::default())
    }

    /// Gets the length of this `PaletteStorage`. That is, how many
    /// elements can be accessed via the `get` and `set` functions.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Gets the number of bits each element takes up when stored.
    ///
    /// More bits means more palette entries can be stored, so more unique
    /// values can be stored in this `PaletteStorage`, at the cost of more
    /// memory. This will automatically increase as necessary.
    pub fn bits(&self) -> usize {
        self.bits
    }

    /// Gets the number of times the specific value occurs in this `PaletteStorage`.
    pub fn get_used_count(&self, value: T) -> usize {
        self.find_existing_entry(value).map(|e| e.used).unwrap_or(0)
    }

    pub fn get(&self, index: usize) -> T {
        let palette_index = self.get_palette_index(index);
        let entry = &self.entries[palette_index];
        entry.value.expect("invalid palette entry")
    }

    pub fn set(&mut self, index: usize, value: T) -> T {
        let prev_palette_index = self.get_palette_index(index);
        let prev_entry = &mut self.entries[prev_palette_index];
        let prev_value = prev_entry.value.expect("invalid palette entry");

        // If entry found at `index` already has this value, return early.
        if prev_value == value {
            return value;
        }

        // Reduce the number of times the previously used palette for this
        // `index` is in use. This potentially allows this entry to be reused.
        prev_entry.used -= 1;

        // Find a palette entry for this value. Resizes palette if necessary.
        let (new_palette_index, new_entry) = self.find_or_create_entry(value);
        // increase the number of times it's in use
        new_entry.used += 1;

        // Update the palette index in the actual `data` BitVec.
        self.set_palette_index(index, new_palette_index);

        prev_value
    }

    /// Gets the bit range into `data` for the specified index.
    fn bit_range(&self, index: usize) -> std::ops::Range<usize> {
        (index * self.bits)..((index + 1) * self.bits)
    }

    /// Looks up the palette index (into `entries`) at the specified index.
    fn get_palette_index(&self, index: usize) -> usize {
        if self.bits > 0 {
            let bit_range = self.bit_range(index);
            self.data[bit_range].load()
        } else {
            0
        }
    }

    fn set_palette_index(&mut self, index: usize, value: usize) {
        let bit_range = self.bit_range(index);
        self.data[bit_range].store(value)
    }

    fn find_existing_entry(&self, value: T) -> Option<&PaletteEntry<T>> {
        self.entries.iter().find(|e| e.value == Some(value))
    }

    fn find_or_create_entry(&mut self, value: T) -> (usize, &mut PaletteEntry<T>) {
        match self.find_entry_index(value) {
            // The palette entry already exists, so just return it.
            FindResult::Existing(index) => {
                let entry = &mut self.entries[index];
                (index, entry)
            }
            // The entry didn't exist, but we found one we can reuse.
            FindResult::Uninitialized(index) | FindResult::Unused(index) => {
                let entry = &mut self.entries[index];
                entry.value = Some(value);
                (index, entry)
            }
            // Everything in the palette is already in use. RESIZE!
            FindResult::None => {
                let index = self.entries.len();
                self.resize_palette(self.bits + 1);
                let entry = &mut self.entries[index];
                entry.value = Some(value);
                (index, entry)
            }
        }
    }

    fn find_entry_index(&self, value: T) -> FindResult {
        let mut result = FindResult::None;
        for (index, entry) in self.entries.iter().enumerate() {
            match entry {
                // If the specified value is found in the palette, return it immediately.
                PaletteEntry { value: v, .. } if *v == Some(value) => {
                    return FindResult::Existing(index)
                }
                // Store the first uninitialized entry in case we don't find the specified value.
                PaletteEntry { value: None, .. } => match result {
                    FindResult::Existing(_) => unreachable!(),
                    FindResult::Uninitialized(_) => {}
                    _ => result = FindResult::Uninitialized(index),
                },
                // Otherwise, pick the first initialized entry that's currently used.
                PaletteEntry { used: 0, .. } => match result {
                    FindResult::Existing(_) => unreachable!(),
                    FindResult::Uninitialized(_) | FindResult::Unused(_) => {}
                    FindResult::None => result = FindResult::Unused(index),
                },
                _ => {}
            }
        }
        result
    }

    fn resize_palette(&mut self, new_bits: usize) {
        if new_bits == self.bits {
            return;
        }

        // TODO: Revisit this to see if we can optimize it.
        // Create a new data BitVec and copy the bits from the old one over.
        let mut new_data = BitVec::with_capacity(self.len * new_bits);
        if new_bits == 0 {
            // Nothing to do if we have nothing to copy to.
        } else if self.bits == 0 {
            // No data to copy from, so just fill with zeroes.
            new_data.resize(self.len * new_bits, false);
        } else if new_bits > self.bits {
            let additional_bits = new_bits - self.bits;
            for chunk in self.data.chunks_exact(self.bits) {
                new_data.extend(chunk);
                for _ in 0..additional_bits {
                    new_data.push(false);
                }
            }
        } else {
            for chunk in self.data.chunks_exact(self.bits) {
                new_data.extend(&chunk[0..new_bits]);
            }
        }
        self.data = new_data;
        self.bits = new_bits;

        // Resize the palette itself.
        let num_entries = 2usize.pow(new_bits as u32);
        self.entries.resize_with(num_entries, Default::default);
    }
}

enum FindResult {
    Existing(usize),
    Uninitialized(usize),
    Unused(usize),
    None,
}

// #[derive(Default)]
#[derive(Clone)]
struct PaletteEntry<T> {
    value: Option<T>,
    used: usize,
}

// NOTE: For some weird reason, deriving `Default` doesn't quite do what we want
//       when later calling `Vec::resize_with`, so we're manually implementing
//       the trait instead. [insert confused noises here]
impl<T> Default for PaletteEntry<T> {
    fn default() -> Self {
        Self {
            value: None,
            used: 0,
        }
    }
}