use hashbrown::hash_map::{self, HashMap};

use std::{
    ffi::CStr,
    hash::{Hash, Hasher},
    mem,
};

pub enum HashTable {
    PointerKeyed(HashMap<*mut u8, *mut u8>),
    StringKeyed(HashMap<StringKey, *mut u8>),
}

impl HashTable {
    pub fn new_pointer_keyed() -> Self {
        HashTable::PointerKeyed(HashMap::new())
    }

    pub fn new_string_keyed() -> Self {
        HashTable::StringKeyed(HashMap::new())
    }

    pub fn clear(&mut self) {
        match self {
            HashTable::PointerKeyed(m) => m.clear(),
            HashTable::StringKeyed(m) => m.clear(),
        }
    }

    pub fn len(&self) -> usize {
        match self {
            HashTable::PointerKeyed(m) => m.len(),
            HashTable::StringKeyed(m) => m.len(),
        }
    }

    pub unsafe fn get(&self, key: *const u8) -> Option<*mut u8> {
        match self {
            HashTable::PointerKeyed(m) => {
                let key = key.cast_mut();
                m.get(&key).copied()
            }
            HashTable::StringKeyed(m) => {
                let key = StringKey(key.cast_mut());
                m.get(&key).copied()
            }
        }
    }

    pub unsafe fn insert(&mut self, key: *mut u8, data: *mut u8) -> Option<*mut u8> {
        match self {
            HashTable::PointerKeyed(m) => m.insert(key, data),
            HashTable::StringKeyed(m) => m.insert(StringKey(key), data),
        }
    }

    pub unsafe fn remove(&mut self, key: *const u8) {
        match self {
            HashTable::PointerKeyed(m) => {
                let key = key.cast_mut();
                m.remove(&key);
            }
            HashTable::StringKeyed(m) => {
                let key = StringKey(key.cast_mut());
                m.remove(&key);
                mem::forget(key);
            }
        }
    }
}

#[repr(transparent)]
pub struct StringKey(*const u8);

impl PartialEq for StringKey {
    fn eq(&self, other: &Self) -> bool {
        match (self.0.is_null(), other.0.is_null()) {
            (true, true) => true,
            (false, false) => unsafe { CStr::from_ptr(self.0.cast()) == CStr::from_ptr(other.0.cast()) },
            _ => false,
        }
    }
}

impl Eq for StringKey {}

impl std::fmt::Debug for StringKey {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "StringKey({:?})", unsafe { std::ffi::CStr::from_ptr(self.0) })
    }
}

impl Hash for StringKey {
    fn hash<H: Hasher>(&self, h: &mut H) {
        if self.0.is_null() {
            h.write_usize(0)
        } else {
            unsafe { CStr::from_ptr(self.0.cast()).hash(h) }
        }
    }
}

pub enum Enumerator<'a> {
    PointerKeyed(hash_map::IterMut<'a, *mut u8, *mut u8>),
    StringKeyed(hash_map::IterMut<'a, StringKey, *mut u8>),
}

pub mod ffi {
    use std::{
        ffi::{c_int, c_uint, c_void},
        mem, ptr,
    };

    use super::{Enumerator, HashTable, StringKey};

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMCreateIdentityHashTable() -> *mut HashTable {
        Box::leak(Box::new(HashTable::new_pointer_keyed()))
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMCreateStringHashTable() -> *mut HashTable {
        Box::leak(Box::new(HashTable::new_string_keyed()))
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMFreeHashTable(table: *mut HashTable) {
        if !table.is_null() {
            let _ = unsafe { Box::from_raw(table) };
        }
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMResetHashTable(table: *mut HashTable) {
        if !table.is_null() {
            unsafe {
                (*table).clear();
            }
        }
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMCountHashTable(table: *mut HashTable) -> c_uint {
        if table.is_null() {
            0
        } else {
            (unsafe { (*table).len() }) as c_uint
        }
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMHashGet(table: *mut HashTable, key: *const c_void) -> *mut c_void {
        if table.is_null() {
            return ptr::null_mut();
        }
        let key = key.cast::<u8>();
        (unsafe { (*table).get(key) })
            .map(|v| v.cast::<c_void>())
            .unwrap_or(ptr::null_mut())
    }

    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMHashGetItemAndKey(
        table: *mut HashTable,
        key: *const c_void,
        value_dest: *mut *mut c_void,
        key_dest: *mut *const c_void,
    ) -> c_int {
        if table.is_null() {
            return 0;
        }
        let table = unsafe { &mut *table };
        match table {
            HashTable::PointerKeyed(m) => {
                let result = match m.get_key_value_mut(&key.cast::<u8>().cast_mut()) {
                    Some((k, v)) => {
                        unsafe {
                            *key_dest = k.cast::<c_void>();
                            *value_dest = v.cast::<c_void>();
                        }
                        1
                    }
                    None => 0,
                };
                result
            }
            HashTable::StringKeyed(m) => {
                let key = StringKey(key.cast::<u8>().cast_mut());
                let result = match m.get_key_value_mut(&key) {
                    Some((k, v)) => {
                        unsafe {
                            *key_dest = k.0.cast::<c_void>();
                            *value_dest = v.cast::<c_void>();
                        }
                        1
                    }
                    None => 0,
                };
                mem::forget(key);
                return result;
            }
        }
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMHashInsert(
        table: *mut HashTable,
        key: *mut c_void,
        data: *mut c_void,
    ) -> *mut c_void {
        if table.is_null() {
            return ptr::null_mut();
        }
        match unsafe { (*table).insert(key.cast::<u8>(), data.cast::<u8>()) } {
            Some(v) => v.cast::<c_void>(),
            None => ptr::null_mut(),
        }
    }

    #[unsafe(no_mangle)]
    #[allow(non_snake_case)]
    pub unsafe extern "C" fn WMHashRemove(table: *mut HashTable, key: *mut c_void) {
        if table.is_null() {
            return;
        }
        unsafe {
            (*table).remove(key.cast::<u8>());
        }
    }

    /// Important note: this may leak memory if you don't pass the enumerator
    /// back to [`WMFreeHashEnumerator`]. This is a breaking change from the
    /// original C implementation, which did not require any resource cleanup.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMEnumerateHashTable(
        table: *mut HashTable,
    ) -> *mut Enumerator<'static> {
        if table.is_null() {
            return ptr::null_mut();
        }
        let table = unsafe { &mut *table };
        match table {
            HashTable::PointerKeyed(m) => {
                Box::leak(Box::new(Enumerator::PointerKeyed(m.iter_mut())))
            }
            HashTable::StringKeyed(m) => Box::leak(Box::new(Enumerator::StringKeyed(m.iter_mut()))),
        }
    }

    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMFreeHashEnumerator(e: *mut Enumerator<'static>) {
        if !e.is_null() {
            let _ = unsafe { Box::from_raw(e) };
        }
    }

    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMNextHashEnumeratorItem(e: *mut Enumerator<'static>) -> *mut c_void {
        if e.is_null() {
            return ptr::null_mut();
        }
        let e = unsafe { &mut *e };
        match e {
            Enumerator::PointerKeyed(i) => match i.next() {
                Some((_, v)) => v.cast::<c_void>(),
                None => ptr::null_mut(),
            },
            Enumerator::StringKeyed(i) => match i.next() {
                Some((_, v)) => v.cast::<c_void>(),
                None => ptr::null_mut(),
            },
        }
    }

    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMNextHashEnumeratorKey(e: *mut Enumerator<'static>) -> *const c_void {
        if e.is_null() {
            return ptr::null_mut();
        }
        let e = unsafe { &mut *e };
        match e {
            Enumerator::PointerKeyed(i) => match i.next() {
                Some((k, _)) => k.cast::<c_void>(),
                None => ptr::null_mut(),
            },
            Enumerator::StringKeyed(i) => match i.next() {
                Some((k, _)) => k.0.cast::<c_void>(),
                None => ptr::null_mut(),
            },
        }
    }

    #[cfg(test)]
    mod test {
        use super::*;

        use std::{ffi::CString, ptr};

        #[test]
        fn enumerate_nonempty() {
            unsafe {
                let table = WMCreateStringHashTable();
                let k1 = CString::new("hello").unwrap().into_raw();
                let v1 = CString::new("world").unwrap().into_raw();
                let k2 = CString::new("foo").unwrap().into_raw();
                let v2 = CString::new("bar").unwrap().into_raw();
                WMHashInsert(table, k1.cast(), v1.cast());
                WMHashInsert(table, k2.cast(), v2.cast());

                let i = WMEnumerateHashTable(table);

                let v = WMNextHashEnumeratorItem(i);
                assert_ne!(v, ptr::null_mut());
                assert!(v == v1.cast() || v == v2.cast());

                let v = WMNextHashEnumeratorItem(i);
                assert_ne!(v, ptr::null_mut());
                assert!(v == v1.cast() || v == v2.cast());

                let v = WMNextHashEnumeratorItem(i);
                assert_eq!(v, ptr::null_mut());
                WMFreeHashEnumerator(i);

                let _ = CString::from_raw(k1);
                let _ = CString::from_raw(v1);
                let _ = CString::from_raw(k2);
                let _ = CString::from_raw(v2);
            }
        }
    }
}
