use std::collections::HashMap;
use std::hash::Hash;
use std::ops::Deref;
use std::sync::OnceLock;

/// Defines the item type, key type and key extraction logic for a [`KeyedVec`].
///
/// Implementations of this trait are generated by the Skir code generator.
/// You should not write your own implementations.
pub trait KeyedVecSpec {
    type Item: 'static;
    /// The owned key type stored in the internal index.
    type StorageKey: Eq + Hash;
    /// Controls how [`KeyedVec::find_by_key`] accepts its argument.
    /// Use [`internal::CopyLookup`] for `Copy` key types, [`internal::BorrowLookup`] for `String` keys.
    type Lookup: internal::Lookup<Self::StorageKey>;

    fn get_key(item: &Self::Item) -> Self::StorageKey;
    /// Returns the dot-separated key path used in type descriptors.
    fn key_extractor() -> &'static str;
    /// Returns a `'static` reference to the default item, used by
    /// [`KeyedVec::find_by_key_or_default`] when no matching element is found.
    fn default_item() -> &'static Self::Item;
}

/// An immutable vector that supports O(1) lookup by key.
///
/// The index is built lazily on the first call to [`KeyedVec::find_by_key`] and
/// cached for subsequent calls. Building and caching the index is thread-safe.
///
/// `S` must implement [`KeyedVecSpec`], which defines both the element type and the
/// lookup key type. When multiple elements share the same key,
/// [`KeyedVec::find_by_key`] returns the first one.
pub struct KeyedVec<S: KeyedVecSpec> {
    items: Vec<S::Item>,
    index: OnceLock<HashMap<S::StorageKey, usize>>,
    _marker: std::marker::PhantomData<S>,
}

impl<S: KeyedVecSpec> KeyedVec<S> {
    pub fn new(items: Vec<S::Item>) -> Self {
        Self {
            items,
            index: OnceLock::new(),
            _marker: std::marker::PhantomData,
        }
    }

    /// Returns the first element whose key equals `key`, or the
    /// [`KeyedVecSpec::default_item`] if no match is found.
    ///
    /// The index is built on the first call and reused on all subsequent calls.
    pub fn find_by_key_or_default<'a>(
        &self,
        key: <S::Lookup as internal::Lookup<S::StorageKey>>::Key<'a>,
    ) -> &S::Item {
        match self.find_by_key(key) {
            Some(item) => item,
            None => S::default_item(),
        }
    }

    /// Returns the first element whose key equals `key`, or `None`.
    ///
    /// The index is built on the first call and reused on all subsequent calls.
    pub fn find_by_key<'a>(
        &self,
        key: <S::Lookup as internal::Lookup<S::StorageKey>>::Key<'a>,
    ) -> Option<&S::Item> {
        let index = self.index.get_or_init(|| {
            let mut map = HashMap::with_capacity(self.items.len());
            for (i, item) in self.items.iter().enumerate() {
                map.entry(S::get_key(item)).or_insert(i);
            }
            map
        });
        <S::Lookup as internal::Lookup<S::StorageKey>>::find_index(index, key)
            .map(|i| &self.items[i])
    }
}

impl<S: KeyedVecSpec> Deref for KeyedVec<S> {
    type Target = [S::Item];

    fn deref(&self) -> &[S::Item] {
        &self.items
    }
}

impl<'a, S: KeyedVecSpec> IntoIterator for &'a KeyedVec<S> {
    type Item = &'a S::Item;
    type IntoIter = std::slice::Iter<'a, S::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.items.iter()
    }
}

impl<S: KeyedVecSpec> IntoIterator for KeyedVec<S> {
    type Item = S::Item;
    type IntoIter = std::vec::IntoIter<S::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.items.into_iter()
    }
}

impl<S: KeyedVecSpec> Clone for KeyedVec<S>
where
    S::Item: Clone,
{
    fn clone(&self) -> Self {
        // The index is not cloned; it will be lazily rebuilt on next access.
        Self::new(self.items.clone())
    }
}

impl<S: KeyedVecSpec> std::fmt::Debug for KeyedVec<S>
where
    S::Item: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.items.fmt(f)
    }
}

impl<S: KeyedVecSpec> Default for KeyedVec<S> {
    fn default() -> Self {
        Self::new(Vec::new())
    }
}

impl<S: KeyedVecSpec> PartialEq for KeyedVec<S>
where
    S::Item: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.items == other.items
    }
}

impl<S: KeyedVecSpec> AsRef<[S::Item]> for KeyedVec<S> {
    fn as_ref(&self) -> &[S::Item] {
        &self.items
    }
}

impl<S: KeyedVecSpec> std::borrow::Borrow<[S::Item]> for KeyedVec<S> {
    fn borrow(&self) -> &[S::Item] {
        &self.items
    }
}

impl<S: KeyedVecSpec> FromIterator<S::Item> for KeyedVec<S> {
    fn from_iter<I: IntoIterator<Item = S::Item>>(iter: I) -> Self {
        Self::new(iter.into_iter().collect())
    }
}

impl<S: KeyedVecSpec> From<Vec<S::Item>> for KeyedVec<S> {
    fn from(items: Vec<S::Item>) -> Self {
        Self::new(items)
    }
}

pub mod internal {
    use std::collections::HashMap;
    use std::hash::Hash;

    /// Lookup strategy for [`super::KeyedVecSpec`] implementors using a `Copy` key type.
    /// Pass this as `type Lookup = CopyLookup` when the key (e.g. `u32`) is `Copy`.
    pub struct CopyLookup;

    /// Lookup strategy for [`super::KeyedVecSpec`] implementors using `String` as the stored key.
    /// Pass this as `type Lookup = BorrowLookup`; callers pass a `&str` with no allocation.
    pub struct BorrowLookup;

    /// Defines how a query key is used to look up an entry in the pre-built index.
    ///
    /// Two ready-made implementations are provided: [`CopyLookup`] and [`BorrowLookup`].
    /// You should not need to implement this trait yourself.
    pub trait Lookup<StorageKey: Eq + Hash> {
        /// The key type accepted by [`super::KeyedVec::find_by_key`].
        type Key<'a>;
        #[doc(hidden)]
        fn find_index<'a>(
            index: &HashMap<StorageKey, usize>,
            query: Self::Key<'a>,
        ) -> Option<usize>;
    }

    impl<K: Eq + Hash + Copy> Lookup<K> for CopyLookup {
        type Key<'a> = K;
        fn find_index<'a>(index: &HashMap<K, usize>, query: K) -> Option<usize> {
            index.get(&query).copied()
        }
    }

    impl Lookup<String> for BorrowLookup {
        type Key<'a> = &'a str;
        fn find_index(index: &HashMap<String, usize>, query: &str) -> Option<usize> {
            index.get(query).copied() // String: Borrow<str>, no allocation
        }
    }
}

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

    #[derive(Default)]
    struct Item {
        id: u32,
        name: String,
    }

    struct ItemById;

    impl KeyedVecSpec for ItemById {
        type Item = Item;
        type StorageKey = u32;
        type Lookup = internal::CopyLookup;
        fn get_key(item: &Item) -> u32 {
            item.id
        }
        fn key_extractor() -> &'static str {
            ""
        }
        fn default_item() -> &'static Item {
            static DEFAULT: std::sync::LazyLock<Item> = std::sync::LazyLock::new(Item::default);
            &DEFAULT
        }
    }

    fn make_vec() -> KeyedVec<ItemById> {
        KeyedVec::new(vec![
            Item {
                id: 1,
                name: "one".to_owned(),
            },
            Item {
                id: 2,
                name: "two".to_owned(),
            },
            Item {
                id: 3,
                name: "three".to_owned(),
            },
        ])
    }

    // --- find_by_key() ---

    #[test]
    fn search_existing_key_returns_element() {
        let kv = make_vec();
        let found = kv.find_by_key(2).expect("key 2 should be present");
        assert_eq!(found.name, "two");
    }

    #[test]
    fn search_missing_key_returns_none() {
        let kv = make_vec();
        assert!(kv.find_by_key(99).is_none());
    }

    #[test]
    fn search_duplicate_key_returns_first_occurrence() {
        let kv: KeyedVec<ItemById> = KeyedVec::new(vec![
            Item {
                id: 42,
                name: "first".to_owned(),
            },
            Item {
                id: 42,
                name: "second".to_owned(),
            },
        ]);
        let found = kv.find_by_key(42).expect("key 42 should be present");
        assert_eq!(found.name, "first");
    }

    #[test]
    fn search_empty_vec_returns_none() {
        let kv: KeyedVec<ItemById> = KeyedVec::new(vec![]);
        assert!(kv.find_by_key(1).is_none());
    }

    // Index is built once; calling find_by_key() a second time still works.
    #[test]
    fn search_called_twice_returns_same_result() {
        let kv = make_vec();
        assert!(kv.find_by_key(1).is_some());
        assert!(kv.find_by_key(1).is_some());
    }

    // --- Vec-like behaviour (Deref) ---

    #[test]
    fn len_reflects_number_of_items() {
        let kv = make_vec();
        assert_eq!(kv.len(), 3);
    }

    #[test]
    fn is_empty_false_for_non_empty_vec() {
        let kv = make_vec();
        assert!(!kv.is_empty());
    }

    #[test]
    fn is_empty_true_for_empty_vec() {
        let kv: KeyedVec<ItemById> = KeyedVec::new(vec![]);
        assert!(kv.is_empty());
    }

    #[test]
    fn index_operator_accesses_elements() {
        let kv = make_vec();
        assert_eq!(kv[0].name, "one");
        assert_eq!(kv[2].name, "three");
    }

    #[test]
    fn iter_visits_all_elements_in_order() {
        let kv = make_vec();
        let names: Vec<&str> = kv.iter().map(|item| item.name.as_str()).collect();
        assert_eq!(names, ["one", "two", "three"]);
    }

    // --- IntoIterator for &KeyedVec ---

    #[test]
    fn for_loop_visits_all_elements_in_order() {
        let kv = make_vec();
        let mut names = Vec::new();
        for item in &kv {
            names.push(item.name.as_str());
        }
        assert_eq!(names, ["one", "two", "three"]);
    }

    // --- String keys ---

    struct ItemByName;

    impl KeyedVecSpec for ItemByName {
        type Item = Item;
        type StorageKey = String;
        type Lookup = internal::BorrowLookup;
        fn get_key(item: &Item) -> String {
            item.name.clone()
        }
        fn key_extractor() -> &'static str {
            ""
        }
        fn default_item() -> &'static Item {
            static DEFAULT: std::sync::LazyLock<Item> = std::sync::LazyLock::new(Item::default);
            &DEFAULT
        }
    }

    #[test]
    fn string_key_found_with_str_slice() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![
            Item {
                id: 1,
                name: "alice".to_owned(),
            },
            Item {
                id: 2,
                name: "bob".to_owned(),
            },
        ]);
        // &str works directly because String: Borrow<str> — no allocation at call site
        let found = kv.find_by_key("bob").expect("bob should be present");
        assert_eq!(found.id, 2);
    }

    #[test]
    fn string_key_missing_returns_none() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![Item {
            id: 1,
            name: "alice".to_owned(),
        }]);
        assert!(kv.find_by_key("carol").is_none());
    }

    #[test]
    fn string_key_found_with_owned_string() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![
            Item {
                id: 1,
                name: "alice".to_owned(),
            },
            Item {
                id: 2,
                name: "bob".to_owned(),
            },
        ]);
        let name = String::from("bob");
        // Pass &name (i.e. &String, which coerces to &str) to satisfy QueryKey<'_> = &str.
        let found = kv.find_by_key(&name).expect("bob should be present");
        assert_eq!(found.id, 2);
    }

    #[test]
    fn string_key_missing_with_owned_string() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![Item {
            id: 1,
            name: "alice".to_owned(),
        }]);
        let name = String::from("carol");
        assert!(kv.find_by_key(&name).is_none());
    }

    // --- find_by_key_or_default() ---

    #[test]
    fn find_or_default_returns_item_when_found() {
        let kv = make_vec();
        let item = kv.find_by_key_or_default(2);
        assert_eq!(item.name, "two");
    }

    #[test]
    fn find_or_default_returns_default_when_missing() {
        let kv = make_vec();
        let item = kv.find_by_key_or_default(99);
        assert_eq!(item.id, 0);
        assert_eq!(item.name, "");
    }

    #[test]
    fn find_or_default_string_key_returns_item_when_found() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![
            Item {
                id: 1,
                name: "alice".to_owned(),
            },
            Item {
                id: 2,
                name: "bob".to_owned(),
            },
        ]);
        let item = kv.find_by_key_or_default("bob");
        assert_eq!(item.id, 2);
    }

    #[test]
    fn find_or_default_string_key_returns_default_when_missing() {
        let kv: KeyedVec<ItemByName> = KeyedVec::new(vec![Item {
            id: 1,
            name: "alice".to_owned(),
        }]);
        let item = kv.find_by_key_or_default("carol");
        assert_eq!(item.id, 0);
        assert_eq!(item.name, "");
    }
}
