use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};

use log::error;

use super::{Dictionary, DictionaryInfo, Entries, LookupStrategy, Phrase, UpdateDictionaryError};
use crate::{
    dictionary::{DictionaryUsage, TrieBuf},
    zhuyin::Syllable,
};

/// A collection of dictionaries that returns the union of the lookup results.
/// # Examples
///
/// ```
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
///
/// use chewing::{
///     dictionary::{Dictionary, Layered, LookupStrategy, Phrase, TrieBuf},
///     syl,
///     zhuyin::Bopomofo,
/// };
///
/// let sys_dict = TrieBuf::from([(
///     vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
///     vec![("測", 1), ("冊", 1), ("側", 1)]
/// )]);
/// let user_dict = TrieBuf::from([(
///     vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
///     vec![("策", 100), ("冊", 100)]
/// )]);
///
/// let dict = Layered::new(vec![Box::new(sys_dict), Box::new(user_dict)]);
/// assert_eq!(
///     [
///         ("側", 1, 0).into(),
///         ("冊", 100, 0).into(),
///         ("測", 1, 0).into(),
///         ("策", 100, 0).into(),
///     ]
///     .into_iter()
///     .collect::<Vec<Phrase>>(),
///     dict.lookup(&[
///         syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]
///     ], LookupStrategy::Standard),
/// );
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct Layered {
    dicts: Vec<Box<dyn Dictionary>>,
    user_dict_index: usize,
}

impl Layered {
    /// Creates a new `Layered` with the list of dictionaries.
    pub fn new(mut dicts: Vec<Box<dyn Dictionary>>) -> Layered {
        let user_dict_index = dicts.iter().enumerate().find_map(|d| {
            if d.1.about().usage == DictionaryUsage::User {
                Some(d.0)
            } else {
                None
            }
        });
        if user_dict_index.is_none() {
            let mut mem_dict = TrieBuf::new_in_memory();
            mem_dict.set_usage(DictionaryUsage::User);
            dicts.push(Box::new(mem_dict));
        }
        let user_dict_index = user_dict_index.unwrap_or(dicts.len() - 1);
        Layered {
            dicts,
            user_dict_index,
        }
    }
    pub fn user_dict(&self) -> &dyn Dictionary {
        self.dicts[self.user_dict_index].as_ref()
    }
    pub fn user_dict_mut(&mut self) -> &mut dyn Dictionary {
        self.dicts[self.user_dict_index].as_mut()
    }
    fn enabled_non_user_dicts(&self) -> impl Iterator<Item = &Box<dyn Dictionary>> {
        self.dicts.iter().filter(|d| {
            !matches!(
                d.about().usage,
                DictionaryUsage::User | DictionaryUsage::ExcludeList,
            )
        })
    }
    fn enabled_dicts(&self) -> impl Iterator<Item = &Box<dyn Dictionary>> {
        self.dicts
            .iter()
            .filter(|d| d.about().usage != DictionaryUsage::ExcludeList)
    }
    fn exclusion_dicts(&self) -> impl Iterator<Item = &Box<dyn Dictionary>> {
        self.dicts
            .iter()
            .filter(|d| d.about().usage == DictionaryUsage::ExcludeList)
    }
    fn exclusion_dicts_mut(&mut self) -> impl Iterator<Item = &mut Box<dyn Dictionary>> {
        self.dicts
            .iter_mut()
            .filter(|d| d.about().usage == DictionaryUsage::ExcludeList)
    }
    pub(crate) fn is_excluded(&self, syllables: &[Syllable], phrase: &str) -> bool {
        self.exclusion_dicts().any(|d| {
            d.lookup(syllables, LookupStrategy::Standard)
                .iter()
                .any(|p| p.text.as_ref() == phrase)
        })
    }
}

impl Dictionary for Layered {
    /// Lookup phrases from all underlying dictionaries.
    ///
    /// Phrases are ordered by their first apperance in the underlying
    /// dictionaries.
    ///
    /// When a phrase appears in multiple dictionaries, the final
    /// frequency is the max of all frequency in all dictionaries.
    ///
    /// Pseudo code
    ///
    /// ```pseudo_code
    /// Set phrases = list()
    /// Foreach d in d_layers
    ///   Foreach phrase, freq in d.lookup_syllables()
    ///     If phrase in phrases
    ///       Set phrases[phrase].freq = max(phrases[phrase].freq, freq)
    ///     Else
    ///       Add phrases <- (phrase, freq)
    /// ```
    fn lookup(&self, syllables: &[Syllable], strategy: LookupStrategy) -> Vec<Phrase> {
        let mut sort_map: BTreeMap<String, usize> = BTreeMap::new();
        let mut phrases: Vec<Phrase> = Vec::new();

        self.enabled_non_user_dicts().for_each(|d| {
            for phrase in d.lookup(syllables, strategy) {
                debug_assert!(!phrase.as_str().is_empty());
                match sort_map.entry(phrase.to_string()) {
                    Entry::Occupied(entry) => {
                        let index = *entry.get();
                        phrases[index].freq = phrase.freq.max(phrases[index].freq);
                        phrases[index].last_used =
                            match (phrases[index].last_used, phrase.last_used) {
                                (Some(orig), Some(new)) => Some(u64::max(orig, new)),
                                (Some(orig), None) => Some(orig),
                                (None, Some(new)) => Some(new),
                                (None, None) => None,
                            };
                    }
                    Entry::Vacant(entry) => {
                        entry.insert(phrases.len());
                        phrases.push(phrase);
                    }
                }
            }
        });

        // NB: User dictionary should override other dictionaries.
        // However, old libchewing and/or PIME may store user phrase with lower
        // frequency than built-in dictionaries. Thus it is important to only
        // override when user frequency is greater than the built-in ones.
        //
        // TODO: design a safe upgrade path to migrate to normalized learning.
        for phrase in self.user_dict().lookup(syllables, strategy) {
            match sort_map.entry(phrase.to_string()) {
                Entry::Occupied(entry) => {
                    let index = *entry.get();
                    phrases[index].freq = u32::max(phrases[index].freq, phrase.freq);
                    phrases[index].last_used = phrase.last_used;
                }
                Entry::Vacant(entry) => {
                    entry.insert(phrases.len());
                    phrases.push(phrase);
                }
            }
        }

        // Remove excluded
        let excluded: BTreeSet<Box<str>> = self
            .exclusion_dicts()
            .flat_map(|d| d.lookup(syllables, strategy))
            .map(|p| p.text)
            .collect();
        phrases
            .into_iter()
            .filter(|p| !excluded.contains(&p.text))
            .collect()
    }

    /// Returns all entries from all dictionaries.
    ///
    /// **NOTE**: Duplicate entries are not removed.
    fn entries(&self) -> Entries<'_> {
        Box::new(self.enabled_dicts().flat_map(|dict| dict.entries()))
    }

    fn about(&self) -> DictionaryInfo {
        DictionaryInfo {
            name: "Built-in Layered".to_string(),
            ..Default::default()
        }
    }

    fn path(&self) -> Option<&std::path::Path> {
        None
    }

    fn set_usage(&mut self, _usage: DictionaryUsage) {}

    fn reopen(&mut self) -> Result<(), UpdateDictionaryError> {
        self.exclusion_dicts_mut().for_each(|d| {
            if let Err(error) = d.reopen() {
                error!("Failed to reopen exclusion dictionary: {error}");
            }
        });
        self.user_dict_mut().reopen()
    }

    fn flush(&mut self) -> Result<(), UpdateDictionaryError> {
        self.exclusion_dicts_mut().for_each(|d| {
            if let Err(error) = d.flush() {
                error!("Failed to flush exclusion dictionary: {error}");
            }
        });
        self.user_dict_mut().flush()
    }

    fn add_phrase(
        &mut self,
        syllables: &[Syllable],
        phrase: Phrase,
    ) -> Result<(), UpdateDictionaryError> {
        if phrase.as_str().is_empty() {
            error!("BUG! added phrase is empty");
            return Ok(());
        }
        self.exclusion_dicts_mut().for_each(|d| {
            if let Err(error) = d.remove_phrase(syllables, &phrase.text) {
                error!(
                    "Failed to remove {phrase} {syllables:?} from exclusion dictionary: {error}"
                );
            }
        });
        self.user_dict_mut().add_phrase(syllables, phrase)
    }

    fn update_phrase(
        &mut self,
        syllables: &[Syllable],
        phrase: Phrase,
        user_freq: u32,
        time: u64,
    ) -> Result<(), UpdateDictionaryError> {
        if phrase.as_str().is_empty() {
            error!("BUG! added phrase is empty");
            return Ok(());
        }
        self.user_dict_mut()
            .update_phrase(syllables, phrase, user_freq, time)
    }

    fn remove_phrase(
        &mut self,
        syllables: &[Syllable],
        phrase_str: &str,
    ) -> Result<(), UpdateDictionaryError> {
        // Only exclude phrases if it's not a single character
        if syllables.len() > 1 {
            self.exclusion_dicts_mut().for_each(|d| {
                if let Err(error) = d.add_phrase(syllables, (phrase_str, 0).into()) {
                    error!(
                        "Failed to add {phrase_str} {syllables:?} to exclusion dictionary: {error}"
                    );
                }
            });
        }
        self.user_dict_mut().remove_phrase(syllables, phrase_str)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        error::Error,
        io::{Cursor, Seek},
    };

    use super::Layered;
    use crate::{
        dictionary::{
            Dictionary, DictionaryBuilder, DictionaryUsage, LookupStrategy, Phrase, Trie, TrieBuf,
            TrieBuilder,
        },
        syl,
        zhuyin::Bopomofo,
    };

    #[test]
    fn test_entries() -> Result<(), Box<dyn Error>> {
        let sys_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("測", 1), ("冊", 1), ("側", 1)],
        )]);
        let user_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("策", 100), ("冊", 100)],
        )]);

        let dict = Layered::new(vec![Box::new(sys_dict), Box::new(user_dict)]);
        assert_eq!(
            [
                (
                    vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                    ("側", 1, 0).into()
                ),
                (
                    vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                    ("冊", 1, 0).into()
                ),
                (
                    vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                    ("測", 1, 0).into()
                ),
                (
                    vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                    ("冊", 100, 0).into()
                ),
                (
                    vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                    ("策", 100, 0).into()
                ),
            ]
            .into_iter()
            .collect::<Vec<_>>(),
            dict.entries().collect::<Vec<_>>(),
        );
        Ok(())
    }

    #[test]
    fn test_lookup() -> Result<(), Box<dyn Error>> {
        let sys_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("測", 1), ("冊", 1), ("側", 1)],
        )]);
        let user_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("策", 100), ("冊", 100)],
        )]);

        let dict = Layered::new(vec![Box::new(sys_dict), Box::new(user_dict)]);
        assert_eq!(
            Some(("側", 1, 0).into()),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        assert_eq!(
            [
                ("側", 1, 0).into(),
                ("冊", 100, 0).into(),
                ("測", 1, 0).into(),
                ("策", 100, 0).into(),
            ]
            .into_iter()
            .collect::<Vec<Phrase>>(),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            ),
        );
        Ok(())
    }

    #[test]
    fn test_user_dict_override() -> Result<(), Box<dyn Error>> {
        let sys_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("測", 100), ("冊", 100), ("側", 100)],
        )]);
        let mut user_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("策", 1), ("冊", 1)],
        )]);
        user_dict.set_usage(DictionaryUsage::User);

        let dict = Layered::new(vec![Box::new(sys_dict), Box::new(user_dict)]);
        assert_eq!(
            Some(("側", 100, 0).into()),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        assert_eq!(
            [
                ("側", 100, 0).into(),
                ("冊", 100, 0).into(),
                ("測", 100, 0).into(),
                ("策", 1, 0).into(),
            ]
            .into_iter()
            .collect::<Vec<Phrase>>(),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            ),
        );
        Ok(())
    }

    #[test]
    fn test_exclusion_dict() -> Result<(), Box<dyn Error>> {
        let mut sys_dict = TrieBuf::from([
            (
                vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                vec![("測", 100)],
            ),
            (
                vec![
                    syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],
                    syl![Bopomofo::SH, Bopomofo::TONE4],
                ],
                vec![("測試", 100)],
            ),
        ]);
        sys_dict.set_usage(DictionaryUsage::BuiltIn);

        let mut exclude_dict = TrieBuf::new_in_memory();
        exclude_dict.set_usage(DictionaryUsage::ExcludeList);

        let mut user_dict = TrieBuf::from([
            (
                vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                vec![("測", 1000)],
            ),
            (
                vec![
                    syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],
                    syl![Bopomofo::SH, Bopomofo::TONE4],
                ],
                vec![("測試", 1000)],
            ),
        ]);
        user_dict.set_usage(DictionaryUsage::User);

        let mut dict = Layered::new(vec![
            Box::new(sys_dict),
            Box::new(user_dict),
            Box::new(exclude_dict),
        ]);
        // Verify baseline
        assert_eq!(
            Some(("測", 1000, 0).into()),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        assert_eq!(
            Some(("測試", 1000, 0).into()),
            dict.lookup(
                &vec![
                    syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],
                    syl![Bopomofo::SH, Bopomofo::TONE4],
                ],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        // Remove a phrase should also exclude it
        dict.remove_phrase(
            &vec![
                syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],
                syl![Bopomofo::SH, Bopomofo::TONE4],
            ],
            "測試",
        )?;
        assert_eq!(
            None,
            dict.lookup(
                &vec![
                    syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],
                    syl![Bopomofo::SH, Bopomofo::TONE4],
                ],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        // Remove a character should only remove it from the user dict
        dict.remove_phrase(&vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]], "測")?;
        assert_eq!(
            Some(("測", 100, 0).into()),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4],],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        Ok(())
    }

    #[test]
    fn test_readonly_user_dict() -> Result<(), Box<dyn Error>> {
        let sys_dict = TrieBuf::from([(
            vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            vec![("測", 1), ("冊", 1), ("側", 1)],
        )]);
        let mut builder = TrieBuilder::new();
        builder.set_usage(DictionaryUsage::User);
        builder.insert(
            &[syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            ("策", 100, 0).into(),
        )?;
        builder.insert(
            &[syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
            ("冊", 100, 0).into(),
        )?;
        let mut cursor = Cursor::new(vec![]);
        builder.write(&mut cursor)?;
        cursor.rewind()?;
        let mut user_dict = Trie::new(&mut cursor)?;
        user_dict.set_usage(DictionaryUsage::User);

        let mut dict = Layered::new(vec![Box::new(sys_dict), Box::new(user_dict)]);
        assert_eq!(
            Some(("側", 1, 0).into()),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            )
            .first()
            .cloned(),
        );
        assert_eq!(
            [
                ("側", 1, 0).into(),
                ("冊", 100, 0).into(),
                ("測", 1, 0).into(),
                ("策", 100, 0).into(),
            ]
            .into_iter()
            .collect::<Vec<Phrase>>(),
            dict.lookup(
                &vec![syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                LookupStrategy::Standard
            ),
        );
        let _ = dict.about();
        assert!(dict.reopen().is_err());
        assert!(dict.flush().is_err());
        assert!(
            dict.add_phrase(
                &[syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                ("冊", 100).into()
            )
            .is_err()
        );
        assert!(
            dict.update_phrase(
                &[syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]],
                ("冊", 100).into(),
                0,
                0,
            )
            .is_err()
        );
        assert!(
            dict.remove_phrase(&[syl![Bopomofo::C, Bopomofo::E, Bopomofo::TONE4]], "冊")
                .is_err()
        );
        Ok(())
    }
}
