use std::{
    cell::{Ref, RefCell, RefMut},
    fmt,
    rc::Rc,
    slice,
};

use crate::{
    error::MachineError,
    vm::{Value, env::Environment, machine::Machine, value::convert::AnyFunction},
};

#[derive(Debug, Clone, PartialEq)]
pub struct HashTable(RefCell<HashTableData>);

#[derive(Clone)]
pub struct HashTableData {
    buckets: Box<[Vec<(Value, Value)>]>,
    length: usize,
}

pub struct HashTableIter<'a> {
    buckets: slice::Iter<'a, Vec<(Value, Value)>>,
    inner: Option<slice::Iter<'a, (Value, Value)>>,
}
pub struct HashTableIterMut<'a> {
    buckets: slice::IterMut<'a, Vec<(Value, Value)>>,
    inner: Option<slice::IterMut<'a, (Value, Value)>>,
}

impl HashTable {
    pub fn borrow(&self) -> Ref<'_, HashTableData> {
        self.0.borrow()
    }

    pub fn borrow_mut(&self) -> RefMut<'_, HashTableData> {
        self.0.borrow_mut()
    }
}

impl fmt::Display for HashTable {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&*self.0.borrow(), f)
    }
}

impl HashTableData {
    pub fn new(bucket_count: usize) -> Self {
        let buckets = vec![vec![]; bucket_count].into_boxed_slice();
        Self { buckets, length: 0 }
    }

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

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

    pub fn clear(&mut self) {
        for bucket in self.buckets.iter_mut() {
            bucket.clear();
        }
        self.length = 0;
    }

    // Returns the value inserted
    pub fn insert(&mut self, key: Value, value: Value) -> Result<Value, MachineError> {
        let hash = key
            .hash()
            .ok_or_else(|| MachineError::InvalidHashTableKey(key.clone()))?;

        let bucket_index = (hash % self.buckets.len() as u64) as usize;
        let slot_index = self.buckets[bucket_index]
            .iter()
            .position(|(k, _)| k == &key);
        if let Some(slot_index) = slot_index {
            self.buckets[bucket_index][slot_index].1 = value.clone();
        } else {
            self.buckets[bucket_index].push((key, value.clone()));
            self.length += 1;
        }

        Ok(value)
    }

    pub fn retain_invoke(
        &mut self,
        vm: &mut Machine,
        env: &Rc<Environment>,
        predicate: &AnyFunction,
    ) -> Result<(), MachineError> {
        for bucket in self.buckets.iter_mut() {
            let mut index = 0;
            loop {
                if index >= bucket.len() {
                    break;
                }

                let (key, value) = &bucket[index];
                let output = predicate.invoke(vm, env, &[key.clone(), value.clone()])?;
                if !output.is_trueish() {
                    bucket.remove(index);
                    self.length -= 1;
                } else {
                    index += 1;
                }
            }
        }
        Ok(())
    }

    pub fn contains_key(&self, key: &Value) -> bool {
        let Some(hash) = key.hash() else { return false };
        let bucket_index = (hash % self.buckets.len() as u64) as usize;

        self.buckets[bucket_index]
            .iter()
            .find(|(k, _)| k == key)
            .is_some()
    }

    pub fn get(&self, key: &Value) -> Option<&Value> {
        let hash = key.hash()?;
        let bucket_index = (hash % self.buckets.len() as u64) as usize;

        self.buckets[bucket_index]
            .iter()
            .find(|(k, _)| k == key)
            .map(|(_, v)| v)
    }

    pub fn get_mut(&mut self, key: &Value) -> Option<&mut Value> {
        let hash = key.hash()?;
        let bucket_index = (hash % self.buckets.len() as u64) as usize;

        self.buckets[bucket_index]
            .iter_mut()
            .find(|(k, _)| k == key)
            .map(|(_, v)| v)
    }

    pub fn remove(&mut self, key: &Value) -> Option<Value> {
        let hash = key.hash()?;
        let bucket_index = (hash % self.buckets.len() as u64) as usize;
        let slot_index = self.buckets[bucket_index]
            .iter()
            .position(|(k, _)| k == key)?;
        let value = self.buckets[bucket_index].remove(slot_index);
        self.length -= 1;
        Some(value.1)
    }

    pub fn iter(&self) -> HashTableIter<'_> {
        HashTableIter {
            buckets: self.buckets.iter(),
            inner: None,
        }
    }

    pub fn iter_mut(&mut self) -> HashTableIterMut<'_> {
        HashTableIterMut {
            buckets: self.buckets.iter_mut(),
            inner: None,
        }
    }
}

impl PartialEq for HashTableData {
    fn eq(&self, other: &Self) -> bool {
        if self.length != other.length {
            return false;
        }
        for (key, value1) in self.iter() {
            let Some(value2) = other.get(key) else {
                return false;
            };

            if value1 != value2 {
                return false;
            }
        }

        true
    }
}

impl fmt::Debug for HashTableData {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut list = f.debug_list();
        for bucket in self.buckets.iter() {
            for cell in bucket.iter() {
                list.entry(cell);
            }
        }
        list.finish()
    }
}

impl fmt::Display for HashTableData {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut index = 0;
        for bucket in self.buckets.iter() {
            for (key, value) in bucket.iter() {
                if index != 0 {
                    write!(f, " ")?;
                }
                write!(f, "({key} . {value})")?;
                index += 1;
            }
        }
        Ok(())
    }
}

impl From<HashTableData> for HashTable {
    fn from(value: HashTableData) -> Self {
        Self(RefCell::new(value))
    }
}

impl<'a> Iterator for HashTableIter<'a> {
    type Item = (&'a Value, &'a Value);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(inner) = &mut self.inner
                && let Some((key, value)) = inner.next()
            {
                return Some((key, value));
            }

            let next_bucket = self.buckets.next()?;
            self.inner = Some(next_bucket.iter());
        }
    }
}
impl<'a> Iterator for HashTableIterMut<'a> {
    type Item = (&'a Value, &'a mut Value);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(inner) = &mut self.inner
                && let Some((key, value)) = inner.next()
            {
                return Some((key, value));
            }

            let next_bucket = self.buckets.next()?;
            self.inner = Some(next_bucket.iter_mut());
        }
    }
}
