use std::{convert::TryFrom, marker::PhantomData, mem};

use crate::{
    generated::bpf_map_type::BPF_MAP_TYPE_HASH,
    syscalls::{
        bpf_map_delete_elem, bpf_map_get_next_key, bpf_map_lookup_and_delete_elem,
        bpf_map_lookup_elem, bpf_map_update_elem,
    },
};

use super::{Map, MapError};
use crate::Pod;

pub struct HashMap<T: AsRef<Map>, K, V> {
    inner: T,
    _k: PhantomData<K>,
    _v: PhantomData<V>,
}

impl<T: AsRef<Map>, K: Pod, V: Pod> HashMap<T, K, V> {
    pub fn new(map: T) -> Result<HashMap<T, K, V>, MapError> {
        let inner = map.as_ref();
        let map_type = inner.obj.def.map_type;
        if map_type != BPF_MAP_TYPE_HASH {
            return Err(MapError::InvalidMapType {
                map_type: map_type as u32,
            })?;
        }
        let size = mem::size_of::<K>();
        let expected = inner.obj.def.key_size as usize;
        if size != expected {
            return Err(MapError::InvalidKeySize { size, expected });
        }

        let size = mem::size_of::<V>();
        let expected = inner.obj.def.value_size as usize;
        if size != expected {
            return Err(MapError::InvalidValueSize { size, expected });
        }

        Ok(HashMap {
            inner: map,
            _k: PhantomData,
            _v: PhantomData,
        })
    }

    pub unsafe fn get(&self, key: &K, flags: u64) -> Result<Option<V>, MapError> {
        let fd = self.inner.as_ref().fd_or_err()?;
        bpf_map_lookup_elem(fd, key, flags)
            .map_err(|(code, io_error)| MapError::LookupElementFailed { code, io_error })
    }

    pub unsafe fn iter<'coll>(&'coll self) -> MapIter<'coll, T, K, V> {
        MapIter::new(self)
    }

    pub unsafe fn keys<'coll>(&'coll self) -> MapKeys<'coll, T, K, V> {
        MapKeys::new(self)
    }
}

impl<T: AsRef<Map> + AsMut<Map>, K: Pod, V: Pod> HashMap<T, K, V> {
    pub fn insert(&mut self, key: K, value: V, flags: u64) -> Result<(), MapError> {
        let fd = self.inner.as_ref().fd_or_err()?;
        bpf_map_update_elem(fd, &key, &value, flags)
            .map_err(|(code, io_error)| MapError::UpdateElementFailed { code, io_error })?;
        Ok(())
    }

    pub unsafe fn pop(&mut self, key: &K) -> Result<Option<V>, MapError> {
        let fd = self.inner.as_ref().fd_or_err()?;
        bpf_map_lookup_and_delete_elem(fd, key)
            .map_err(|(code, io_error)| MapError::LookupAndDeleteElementFailed { code, io_error })
    }

    pub fn remove(&mut self, key: &K) -> Result<(), MapError> {
        let fd = self.inner.as_ref().fd_or_err()?;
        bpf_map_delete_elem(fd, key)
            .map(|_| ())
            .map_err(|(code, io_error)| MapError::DeleteElementFailed { code, io_error })
    }
}

impl<'a, K: Pod, V: Pod> TryFrom<&'a Map> for HashMap<&'a Map, K, V> {
    type Error = MapError;

    fn try_from(inner: &'a Map) -> Result<HashMap<&'a Map, K, V>, MapError> {
        HashMap::new(inner)
    }
}

impl<'a, K: Pod, V: Pod> TryFrom<&'a mut Map> for HashMap<&'a mut Map, K, V> {
    type Error = MapError;

    fn try_from(inner: &'a mut Map) -> Result<HashMap<&'a mut Map, K, V>, MapError> {
        HashMap::new(inner)
    }
}

pub struct MapKeys<'coll, T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> {
    map: &'coll HashMap<T, K, V>,
    err: bool,
    key: Option<K>,
}

impl<'coll, T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> MapKeys<'coll, T, K, V> {
    fn new(map: &'coll HashMap<T, K, V>) -> MapKeys<'coll, T, K, V> {
        MapKeys {
            map,
            err: false,
            key: None,
        }
    }
}

impl<T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> Iterator for MapKeys<'_, T, K, V> {
    type Item = Result<K, MapError>;

    fn next(&mut self) -> Option<Result<K, MapError>> {
        if self.err {
            return None;
        }

        let fd = match self.map.inner.as_ref().fd_or_err() {
            Ok(fd) => fd,
            Err(e) => {
                self.err = true;
                return Some(Err(e));
            }
        };

        match bpf_map_get_next_key(fd, self.key.as_ref()) {
            Ok(Some(key)) => {
                self.key = Some(key);
                return Some(Ok(key));
            }
            Ok(None) => {
                self.key = None;
                return None;
            }
            Err((code, io_error)) => {
                self.err = true;
                return Some(Err(MapError::GetNextKeyFailed { code, io_error }));
            }
        }
    }
}

pub struct MapIter<'coll, T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> {
    inner: MapKeys<'coll, T, K, V>,
}

impl<'coll, T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> MapIter<'coll, T, K, V> {
    fn new(map: &'coll HashMap<T, K, V>) -> MapIter<'coll, T, K, V> {
        MapIter {
            inner: MapKeys::new(map),
        }
    }
}

impl<T: AsRef<Map>, K: Clone + Pod, V: Clone + Pod> Iterator for MapIter<'_, T, K, V> {
    type Item = Result<(K, V), MapError>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.inner.next() {
                Some(Ok(key)) => {
                    let value = unsafe { self.inner.map.get(&key, 0) };
                    match value {
                        Ok(None) => continue,
                        Ok(Some(value)) => return Some(Ok((key, value))),
                        Err(e) => return Some(Err(e)),
                    }
                }
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
    }
}

impl AsRef<Map> for &Map {
    fn as_ref(&self) -> &Map {
        self
    }
}

impl AsRef<Map> for &mut Map {
    fn as_ref(&self) -> &Map {
        self
    }
}

impl AsMut<Map> for &mut Map {
    fn as_mut(&mut self) -> &mut Map {
        self
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use libc::{EFAULT, ENOENT};

    use crate::{
        bpf_map_def,
        generated::{
            bpf_attr, bpf_cmd,
            bpf_map_type::{BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_PERF_EVENT_ARRAY},
        },
        obj,
        syscalls::{override_syscall, SysResult, Syscall},
    };

    use super::*;

    fn new_obj_map(name: &str) -> obj::Map {
        obj::Map {
            name: name.to_string(),
            def: bpf_map_def {
                map_type: BPF_MAP_TYPE_HASH,
                key_size: 4,
                value_size: 4,
                max_entries: 1024,
                map_flags: 0,
            },
            section_index: 0,
            data: Vec::new(),
        }
    }

    fn sys_error(value: i32) -> SysResult {
        Err((-1, io::Error::from_raw_os_error(value)))
    }

    #[test]
    fn test_wrong_key_size() {
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        assert!(matches!(
            HashMap::<_, u8, u32>::new(&map),
            Err(MapError::InvalidKeySize {
                size: 1,
                expected: 4
            })
        ));
    }

    #[test]
    fn test_wrong_value_size() {
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        assert!(matches!(
            HashMap::<_, u32, u16>::new(&map),
            Err(MapError::InvalidValueSize {
                size: 2,
                expected: 4
            })
        ));
    }

    #[test]
    fn test_try_from_wrong_map() {
        let map = Map {
            obj: obj::Map {
                name: "TEST".to_string(),
                def: bpf_map_def {
                    map_type: BPF_MAP_TYPE_PERF_EVENT_ARRAY,
                    key_size: 4,
                    value_size: 4,
                    max_entries: 1024,
                    map_flags: 0,
                },
                section_index: 0,
                data: Vec::new(),
            },
            fd: None,
        };

        assert!(matches!(
            HashMap::<_, u32, u32>::try_from(&map),
            Err(MapError::InvalidMapType { .. })
        ));
    }

    #[test]
    fn test_try_from_ok() {
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        assert!(HashMap::<_, u32, u32>::try_from(&map).is_ok())
    }

    #[test]
    fn test_insert_not_created() {
        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(
            hm.insert(1, 42, 0),
            Err(MapError::NotCreated { .. })
        ));
    }

    #[test]
    fn test_insert_syscall_error() {
        override_syscall(|_| sys_error(EFAULT));

        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(
            hm.insert(1, 42, 0),
            Err(MapError::UpdateElementFailed { code: -1, io_error }) if io_error.raw_os_error() == Some(EFAULT)
        ));
    }

    #[test]
    fn test_insert_ok() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_UPDATE_ELEM,
                ..
            } => Ok(1),
            _ => sys_error(EFAULT),
        });

        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(hm.insert(1, 42, 0).is_ok());
    }

    #[test]
    fn test_remove_not_created() {
        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(hm.remove(&1), Err(MapError::NotCreated { .. })));
    }

    #[test]
    fn test_remove_syscall_error() {
        override_syscall(|_| sys_error(EFAULT));

        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(
            hm.remove(&1),
            Err(MapError::DeleteElementFailed { code: -1, io_error }) if io_error.raw_os_error() == Some(EFAULT)
        ));
    }

    #[test]
    fn test_remove_ok() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_DELETE_ELEM,
                ..
            } => Ok(1),
            _ => sys_error(EFAULT),
        });

        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(hm.remove(&1).is_ok());
    }

    #[test]
    fn test_get_not_created() {
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        assert!(matches!(
            unsafe { hm.get(&1, 0) },
            Err(MapError::NotCreated { .. })
        ));
    }

    #[test]
    fn test_get_syscall_error() {
        override_syscall(|_| sys_error(EFAULT));
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        assert!(matches!(
            unsafe { hm.get(&1, 0) },
            Err(MapError::LookupElementFailed { code: -1, io_error }) if io_error.raw_os_error() == Some(EFAULT)
        ));
    }

    #[test]
    fn test_get_not_found() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_ELEM,
                ..
            } => sys_error(ENOENT),
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        assert!(matches!(unsafe { hm.get(&1, 0) }, Ok(None)));
    }

    #[test]
    fn test_pop_not_created() {
        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: None,
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(
            unsafe { hm.pop(&1) },
            Err(MapError::NotCreated { .. })
        ));
    }

    #[test]
    fn test_pop_syscall_error() {
        override_syscall(|_| sys_error(EFAULT));
        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(
            unsafe { hm.pop(&1) },
            Err(MapError::LookupAndDeleteElementFailed { code: -1, io_error }) if io_error.raw_os_error() == Some(EFAULT)
        ));
    }

    #[test]
    fn test_pop_not_found() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_AND_DELETE_ELEM,
                ..
            } => sys_error(ENOENT),
            _ => sys_error(EFAULT),
        });
        let mut map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let mut hm = HashMap::<_, u32, u32>::new(&mut map).unwrap();

        assert!(matches!(unsafe { hm.pop(&1) }, Ok(None)));
    }

    fn bpf_key<T: Copy>(attr: &bpf_attr) -> Option<T> {
        match unsafe { attr.__bindgen_anon_2.key } as *const T {
            p if p.is_null() => None,
            p => Some(unsafe { *p }),
        }
    }

    fn set_next_key<T: Copy>(attr: &bpf_attr, next: T) {
        let key = unsafe { attr.__bindgen_anon_2.__bindgen_anon_1.next_key } as *const T as *mut T;
        unsafe { *key = next };
    }

    fn set_ret<T: Copy>(attr: &bpf_attr, ret: T) {
        let value = unsafe { attr.__bindgen_anon_2.__bindgen_anon_1.value } as *const T as *mut T;
        unsafe { *value = ret };
    }

    #[test]
    fn test_keys_empty() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                ..
            } => sys_error(ENOENT),
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();
        let keys = unsafe { hm.keys() }.collect::<Result<Vec<_>, _>>();
        assert!(matches!(keys, Ok(ks) if ks.is_empty()))
    }

    fn get_next_key(attr: &bpf_attr) -> SysResult {
        match bpf_key(&attr) {
            None => set_next_key(&attr, 10),
            Some(10) => set_next_key(&attr, 20),
            Some(20) => set_next_key(&attr, 30),
            Some(30) => return sys_error(ENOENT),
            Some(_) => return sys_error(EFAULT),
        };

        Ok(1)
    }

    fn lookup_elem(attr: &bpf_attr) -> SysResult {
        match bpf_key(&attr) {
            Some(10) => set_ret(&attr, 100),
            Some(20) => set_ret(&attr, 200),
            Some(30) => set_ret(&attr, 300),
            Some(_) => return sys_error(ENOENT),
            None => return sys_error(EFAULT),
        };

        Ok(1)
    }

    #[test]
    fn test_keys() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => get_next_key(&attr),
            _ => sys_error(EFAULT),
        });

        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        let keys = unsafe { hm.keys() }.collect::<Result<Vec<_>, _>>().unwrap();
        assert_eq!(&keys, &[10, 20, 30])
    }

    #[test]
    fn test_keys_error() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => {
                match bpf_key(&attr) {
                    None => set_next_key(&attr, 10),
                    Some(10) => set_next_key(&attr, 20),
                    Some(_) => return sys_error(EFAULT),
                };

                Ok(1)
            }
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        let mut keys = unsafe { hm.keys() };
        assert!(matches!(keys.next(), Some(Ok(10))));
        assert!(matches!(keys.next(), Some(Ok(20))));
        assert!(matches!(keys.next(), Some(Err(MapError::GetNextKeyFailed { .. }))));
        assert!(matches!(keys.next(), None));
    }

    #[test]
    fn test_iter() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => get_next_key(&attr),
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_ELEM,
                attr,
            } => lookup_elem(&attr),
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();
        let items = unsafe { hm.iter() }.collect::<Result<Vec<_>, _>>().unwrap();
        assert_eq!(&items, &[(10, 100), (20, 200), (30, 300)])
    }

    #[test]
    fn test_iter_key_deleted() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => get_next_key(&attr),
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_ELEM,
                attr,
            } => {
                match bpf_key(&attr) {
                    Some(10) => set_ret(&attr, 100),
                    Some(20) => return sys_error(ENOENT),
                    Some(30) => set_ret(&attr, 300),
                    Some(_) => return sys_error(ENOENT),
                    None => return sys_error(EFAULT),
                };

                Ok(1)
            }
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        let items = unsafe { hm.iter() }.collect::<Result<Vec<_>, _>>().unwrap();
        assert_eq!(&items, &[(10, 100), (30, 300)])
    }

    #[test]
    fn test_iter_key_error() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => {
                match bpf_key(&attr) {
                    None => set_next_key(&attr, 10),
                    Some(10) => set_next_key(&attr, 20),
                    Some(20) => return sys_error(EFAULT),
                    Some(30) => return sys_error(ENOENT),
                    Some(_) => panic!(),
                };

                Ok(1)
            }
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_ELEM,
                attr,
            } => lookup_elem(&attr),
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        let mut iter = unsafe { hm.iter() };
        assert!(matches!(iter.next(), Some(Ok((10, 100)))));
        assert!(matches!(iter.next(), Some(Ok((20, 200)))));
        assert!(matches!(iter.next(), Some(Err(MapError::GetNextKeyFailed { .. }))));
        assert!(matches!(iter.next(), None));
    }

    #[test]
    fn test_iter_value_error() {
        override_syscall(|call| match call {
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_GET_NEXT_KEY,
                attr,
            } => get_next_key(&attr),
            Syscall::Bpf {
                cmd: bpf_cmd::BPF_MAP_LOOKUP_ELEM,
                attr,
            } => {
                match bpf_key(&attr) {
                    Some(10) => set_ret(&attr, 100),
                    Some(20) => return sys_error(EFAULT),
                    Some(30) => set_ret(&attr, 300),
                    Some(_) => return sys_error(ENOENT),
                    None => return sys_error(EFAULT),
                };

                Ok(1)
            }
            _ => sys_error(EFAULT),
        });
        let map = Map {
            obj: new_obj_map("TEST"),
            fd: Some(42),
        };
        let hm = HashMap::<_, u32, u32>::new(&map).unwrap();

        let mut iter = unsafe { hm.iter() };
        assert!(matches!(iter.next(), Some(Ok((10, 100)))));
        assert!(matches!(iter.next(), Some(Err(MapError::LookupElementFailed { .. }))));
        assert!(matches!(iter.next(), Some(Ok((30, 300)))));
        assert!(matches!(iter.next(), None));
    }
}